From 33119c9bf73a6655f800168160ae6a9effa32472 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Fri, 10 Jul 2026 10:55:11 -0700 Subject: [PATCH] fix: guard remove_special_tokens against tokenizers without a BOS token (#7048) --- .../test_remove_special_tokens_no_bos.py | 46 +++++++++++++++++++ unsloth/chat_templates.py | 5 +- 2 files changed, 49 insertions(+), 2 deletions(-) create mode 100644 tests/python/test_remove_special_tokens_no_bos.py diff --git a/tests/python/test_remove_special_tokens_no_bos.py b/tests/python/test_remove_special_tokens_no_bos.py new file mode 100644 index 0000000000..94c5ea3027 --- /dev/null +++ b/tests/python/test_remove_special_tokens_no_bos.py @@ -0,0 +1,46 @@ +import ast +from pathlib import Path + + +def _load_remove_special_tokens(): + # Extract remove_special_tokens without importing unsloth (importing unsloth + # needs unsloth_zoo / a GPU). The function is pure Python and uses no imports, + # so it execs cleanly in an empty namespace. + source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py" + tree = ast.parse(source.read_text(encoding = "utf-8")) + funcs = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "remove_special_tokens" + ] + namespace = {} + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + exec(compile(module, str(source), "exec"), namespace) + return namespace["remove_special_tokens"] + + +class _StubTokenizer: + def __init__(self, bos_token): + self.bos_token = bos_token + + +def test_no_bos_tokenizer_does_not_crash(): + # Tokenizers such as Qwen2 / Qwen2.5, GPT-2, Falcon and GPT-NeoX have no BOS + # token, so tokenizer.bos_token is None. remove_special_tokens must leave the + # prompt untouched instead of raising + # "TypeError: startswith first arg must be str or a tuple of str, not NoneType". + remove_special_tokens = _load_remove_special_tokens() + assert remove_special_tokens(_StubTokenizer(None), "Hello world") == "Hello world" + + +def test_double_bos_is_stripped(): + # A tokenizer with a BOS token still has a single leading BOS removed. + remove_special_tokens = _load_remove_special_tokens() + assert remove_special_tokens(_StubTokenizer(""), "Hello world") == "Hello world" + + +def test_prompt_without_leading_bos_unchanged(): + # A BOS-bearing tokenizer leaves a prompt that does not start with BOS alone. + remove_special_tokens = _load_remove_special_tokens() + assert remove_special_tokens(_StubTokenizer(""), "Hello world") == "Hello world" diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index dd1e433471..2d3674fb04 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -2103,8 +2103,9 @@ def get_chat_template( def remove_special_tokens(tokenizer, prompt): # Removes double BOS token - if prompt.startswith(tokenizer.bos_token): - prompt = prompt[len(tokenizer.bos_token):] + bos_token = getattr(tokenizer, "bos_token", None) + if bos_token is not None and prompt.startswith(bos_token): + prompt = prompt[len(bos_token):] return prompt