diff --git a/llama_cpp/llama_tokenizer.py b/llama_cpp/llama_tokenizer.py index 1375e1392..fc81b30be 100644 --- a/llama_cpp/llama_tokenizer.py +++ b/llama_cpp/llama_tokenizer.py @@ -81,9 +81,22 @@ def __init__(self, hf_tokenizer: Any): def tokenize( self, text: bytes, add_bos: bool = True, special: bool = True ) -> List[int]: - return self.hf_tokenizer.encode( - text.decode("utf-8", errors="ignore"), add_special_tokens=special - ) + text = text.decode("utf-8", errors="ignore") + tokens = self.hf_tokenizer.encode(text, add_special_tokens=special) + bos_token_id = getattr(self.hf_tokenizer, "bos_token_id", None) + if bos_token_id is None: + return tokens + + if add_bos: + if not tokens or tokens[0] != bos_token_id: + tokens.insert(0, bos_token_id) + elif special and tokens and tokens[0] == bos_token_id: + raw_tokens = self.hf_tokenizer.encode(text, add_special_tokens=False) + # Preserve a BOS token from the input; remove only a prefixed one. + if tokens[1 : len(raw_tokens) + 1] == raw_tokens: + tokens = tokens[1:] + + return tokens def detokenize( self, diff --git a/tests/test_llama_tokenizer.py b/tests/test_llama_tokenizer.py new file mode 100644 index 000000000..dfd423f89 --- /dev/null +++ b/tests/test_llama_tokenizer.py @@ -0,0 +1,36 @@ +from llama_cpp.llama_tokenizer import LlamaHFTokenizer + + +class SyntheticHFTokenizer: + bos_token_id = 1 + eos_token_id = 2 + + def encode(self, text, add_special_tokens=True): + tokens = [self.bos_token_id, 3] if text.startswith("") else [3] + if add_special_tokens: + return [self.bos_token_id] + tokens + [self.eos_token_id] + return tokens + + +def test_hf_tokenizer_respects_add_bos_false(): + tokenizer = LlamaHFTokenizer(SyntheticHFTokenizer()) + + assert tokenizer.tokenize(b"prompt", add_bos=False, special=True) == [3, 2] + + +def test_hf_tokenizer_adds_bos_when_special_tokens_are_disabled(): + tokenizer = LlamaHFTokenizer(SyntheticHFTokenizer()) + + assert tokenizer.tokenize(b"prompt", add_bos=True, special=False) == [1, 3] + + +def test_hf_tokenizer_preserves_explicit_bos_when_add_bos_is_false(): + tokenizer = LlamaHFTokenizer(SyntheticHFTokenizer()) + + assert tokenizer.tokenize(b"prompt", add_bos=False, special=True) == [1, 3, 2] + + +def test_hf_tokenizer_does_not_duplicate_bos_when_special_tokens_add_it(): + tokenizer = LlamaHFTokenizer(SyntheticHFTokenizer()) + + assert tokenizer.tokenize(b"prompt", add_bos=True, special=True) == [1, 3, 2]