huggingface/transformers · #48446

[Fix] Sparse TikToken tokenizers silently fail

remi-or · merged Sep 1, 20263 files · 177 + / 71
src/transformers/convert_slow_tokenizer.py5 + / 2
@@ -1931,15 +1931,18 @@ def __init__(             extra_special_tokens.keys() if isinstance(extra_special_tokens, dict) else extra_special_tokens         ) -    def extract_vocab_merges_from_model(self, tiktoken_url: str):+    @staticmethod+    def load_tiktoken_bpe(tiktoken_url: str) -> dict[bytes, int]:         try:             from tiktoken.load import load_tiktoken_bpe         except Exception:             raise ValueError(                 "`tiktoken` is required to read a `tiktoken` file. Install it with `pip install tiktoken`."             )+        return load_tiktoken_bpe(tiktoken_url) -        bpe_ranks = load_tiktoken_bpe(tiktoken_url)+    def extract_vocab_merges_from_model(self, tiktoken_url: str):+        bpe_ranks = self.load_tiktoken_bpe(tiktoken_url)         byte_encoder = bytes_to_unicode()          def token_bytes_to_string(b):
src/transformers/tokenization_utils_tokenizers.py108 + / 68
@@ -62,6 +62,9 @@ # Slow tokenizers have an additional added tokens files ADDED_TOKENS_FILE = "added_tokens.json" +# Some repos use a tiktoken-style tokenizer with an obvious name+TIKTOKEN_LEGACY_NAME = "tiktoken.model"+ INIT_TOKENIZER_DOCSTRING += """         tokenizer_object ([`tokenizers.Tokenizer`]):             A [`tokenizers.Tokenizer`] object from 🤗 tokenizers to instantiate from. See [Using tokenizers from 🤗@@ -210,77 +213,18 @@ def convert_to_native_format(cls, trust_remote_code=False, **kwargs):          # SentencePiece model (with TikToken fallback)         if isinstance(vocab_file, str) and os.path.isfile(vocab_file) and vocab_file.endswith(".model"):-            try:-                from .convert_slow_tokenizer import SentencePieceExtractor--                # 1. Extract vocab, merges, and spm_precompiled from the .model proto-                extractor = SentencePieceExtractor(vocab_file)-                local_kwargs = extractor.extract(cls.model, **local_kwargs)--                # 2. If a model-specific converter exists, use it.+            # Unless the name is a known Tiktoken pattern, try to use SentencePiece+            if os.path.basename(vocab_file) != TIKTOKEN_LEGACY_NAME:                 try:-                    from .convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS--                    converter_class = SLOW_TO_FAST_CONVERTERS.get(cls.__name__)-                    if converter_class is not None and hasattr(converter_class, "convert_from_spm"):-                        local_kwargs = converter_class.convert_from_spm(**local_kwargs)-                except Exception as e:+                    local_kwargs = cls._convert_from_sentencepiece(vocab_file, local_kwargs)+                    return local_kwargs+                except Exception as e:  # TODO only catch deserialization error here!                     logger.warning(-                        f"Could not reorder vocab using converter for {cls.__name__} due to {e}. Falling back to raw SentencePiece extraction."-                    )-                if hasattr(cls, "convert_from_spm_model"):-                    local_kwargs = cls.convert_from_spm_model(**local_kwargs)--                # 3. For non-model specific tokenizers (e.g. TokenizersBackend used-                #    for MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS), build a _tokenizer-                #    from the proto so normalizer/decoder are configured correctly.-                if "tokenizer_object" not in local_kwargs and (-                    cls is TokenizersBackend or "__init__" not in cls.__dict__-                ):-                    vocab = local_kwargs.pop("vocab", None)-                    merges = local_kwargs.pop("merges", None)--                    # Replace placeholder tokens as specified in added_tokens_decoder-                    added_tokens_decoder = local_kwargs.get("added_tokens_decoder") or {}-                    if vocab is not None and added_tokens_decoder:-                        id_to_token = {token_id: token for token, token_id in vocab.items()}-                        for token_id, new_token in added_tokens_decoder.items():-                            token_id = int(token_id)-                            new_token = str(new_token)-                            current_token = id_to_token.get(token_id)-                            if current_token and current_token != new_token and new_token not in vocab:-                                vocab[new_token] = vocab.pop(current_token)-                                id_to_token[token_id] = new_token--                    tokenizer_object = SpmConverter.build_tokenizer_from_spm_proto(-                        proto=extractor.proto,-                        vocab=vocab,-                        merges=merges,+                        f"Could not extract SentencePiece model from {vocab_file} using sentencepiece library due to"+                        f" {e}. Falling back to TikToken extractor."                     )-                    if tokenizer_object is not None:-                        local_kwargs["tokenizer_object"] = tokenizer_object-                        # Set bos/eos tokens from proto spec if available. This is needed when-                        # building a tokenizer_object directly from a .model file because the-                        # tokenizer_object does not have bos/eos set.-                        proto_spec = extractor.proto.trainer_spec-                        if proto_spec.bos_id >= 0:-                            local_kwargs.setdefault("bos_token", proto_spec.bos_piece or "<s>")-                        if proto_spec.eos_id >= 0:-                            local_kwargs.setdefault("eos_token", proto_spec.eos_piece or "</s>")-                        if proto_spec.unk_id >= 0:-                            local_kwargs.setdefault("unk_token", proto_spec.unk_piece or "<unk>")--            except Exception as e:  # TODO only catch deserialization error here!-                logger.warning(-                    f"Could not extract SentencePiece model from {vocab_file} using sentencepiece library due to {e}. "-                    "Falling back to TikToken extractor."-                )-                from .convert_slow_tokenizer import TikTokenConverter--                converter = TikTokenConverter(-                    vocab_file=vocab_file, extra_special_tokens=local_kwargs.get("extra_special_tokens")-                )-                local_kwargs["tokenizer_object"] = converter.converted()+            # If the name is a known Tiktoken pattern or the SentencePiece extraction failed, use TikToken+            local_kwargs = cls._convert_from_tiktoken(vocab_file, local_kwargs)             return local_kwargs          # Fallback to standard vocab/merges files if they existed!@@ -325,6 +269,102 @@ def _iter_special_tokens(values: Iterable[Any]) -> list[str]:             local_kwargs["merges"] = merges         return local_kwargs +    @classmethod+    def _convert_from_tiktoken(cls, vocab_file: str, local_kwargs: dict[str, Any]) -> dict[str, Any]:+        """Path to a vocab file and a list of extra special tokens."""+        from .convert_slow_tokenizer import TikTokenConverter++        # `added_tokens_decoder` is a (token_id -> token) dict of tokens that may not be contiguous+        added_tokens_decoder: dict[int, str] = {+            int(token_id): token["content"] if isinstance(token, dict) else str(token)+            for token_id, token in (local_kwargs.get("added_tokens_decoder") or {}).items()+        }+        # `extra_special_tokens` is a list of tokens to append to the base vocabulary+        extra_special_tokens: list[str] = local_kwargs.get("extra_special_tokens") or []+        if isinstance(extra_special_tokens, dict):+            extra_special_tokens = list(extra_special_tokens.keys())++        # Retrieve the vocab size from the vocab file+        base_vocab_size = len(TikTokenConverter.load_tiktoken_bpe(vocab_file))++        # If any of the tokens in `added_tokens_decoder` are not in the base vocabulary, they need to be added. To do so+        # we bake them in the `extra_special_tokens` list, with placeholder if `added_tokens_decoder` is not contiguous+        max_added_token_id = max(added_tokens_decoder.keys()) if added_tokens_decoder else 0+        if max_added_token_id >= base_vocab_size:+            new_extras = [+                added_tokens_decoder.get(index, f"<|reserved_token_{index}|>")+                for index in range(base_vocab_size, max_added_token_id + 1)+            ]+            set_new_extras = set(new_extras)+            for special_token in extra_special_tokens:+                if special_token not in set_new_extras:+                    new_extras.append(special_token)+            extra_special_tokens = new_extras++        converter = TikTokenConverter(vocab_file=vocab_file, extra_special_tokens=extra_special_tokens)+        local_kwargs["tokenizer_object"] = converter.converted()+        return local_kwargs++    @classmethod+    def _convert_from_sentencepiece(cls, vocab_file: str, local_kwargs: dict[str, Any]) -> dict[str, Any]:+        from .convert_slow_tokenizer import SentencePieceExtractor++        # 1. Extract vocab, merges, and spm_precompiled from the .model proto+        extractor = SentencePieceExtractor(vocab_file)+        local_kwargs = extractor.extract(cls.model, **local_kwargs)++        # 2. If a model-specific converter exists, use it.+        try:+            from .convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS++            converter_class = SLOW_TO_FAST_CONVERTERS.get(cls.__name__)+            if converter_class is not None and hasattr(converter_class, "convert_from_spm"):+                local_kwargs = converter_class.convert_from_spm(**local_kwargs)+        except Exception as e:+            logger.warning(+                f"Could not reorder vocab using converter for {cls.__name__} due to {e}. Falling back to raw SentencePiece extraction."+            )+        if hasattr(cls, "convert_from_spm_model"):+            local_kwargs = cls.convert_from_spm_model(**local_kwargs)++        # 3. For non-model specific tokenizers (e.g. TokenizersBackend used+        #    for MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS), build a _tokenizer+        #    from the proto so normalizer/decoder are configured correctly.+        if "tokenizer_object" not in local_kwargs and (cls is TokenizersBackend or "__init__" not in cls.__dict__):+            vocab = local_kwargs.pop("vocab", None)+            merges = local_kwargs.pop("merges", None)++            # Replace placeholder tokens as specified in added_tokens_decoder+            added_tokens_decoder = local_kwargs.get("added_tokens_decoder") or {}+            if vocab is not None and added_tokens_decoder:+                id_to_token = {token_id: token for token, token_id in vocab.items()}+                for token_id, new_token in added_tokens_decoder.items():+                    token_id = int(token_id)+                    new_token = str(new_token)+                    current_token = id_to_token.get(token_id)+                    if current_token and current_token != new_token and new_token not in vocab:+                        vocab[new_token] = vocab.pop(current_token)+                        id_to_token[token_id] = new_token++            tokenizer_object = SpmConverter.build_tokenizer_from_spm_proto(+                proto=extractor.proto,+                vocab=vocab,+                merges=merges,+            )+            if tokenizer_object is not None:+                local_kwargs["tokenizer_object"] = tokenizer_object+                # Set bos/eos tokens from proto spec if available. This is needed when+                # building a tokenizer_object directly from a .model file because the+                # tokenizer_object does not have bos/eos set.+                proto_spec = extractor.proto.trainer_spec+                if proto_spec.bos_id >= 0:+                    local_kwargs.setdefault("bos_token", proto_spec.bos_piece or "<s>")+                if proto_spec.eos_id >= 0:+                    local_kwargs.setdefault("eos_token", proto_spec.eos_piece or "</s>")+                if proto_spec.unk_id >= 0:+                    local_kwargs.setdefault("unk_token", proto_spec.unk_piece or "<unk>")+        return local_kwargs+     def __init__(self, *args, **kwargs):         # Truncation/padding dicts extracted from tokenizer.json by convert_to_native_format         # when a class with a custom __init__ rebuilds the backend tokenizer from scratch.
tests/tokenization/test_tokenization_fast.py64 + / 1
@@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import base64 import concurrent.futures import json import os@@ -23,7 +24,7 @@ from tokenizers.models import BPE, WordLevel  from transformers import AutoTokenizer, PreTrainedTokenizerFast-from transformers.testing_utils import require_tokenizers+from transformers.testing_utils import require_tiktoken, require_tokenizers   @require_tokenizers@@ -277,6 +278,68 @@ def test_non_bpe_tokenizer_still_cleans_up(self):         self.assertNotIn(" .", decoded)  +@require_tokenizers+@require_tiktoken+class TikTokenAddedTokensTest(unittest.TestCase):+    """Loading a tiktoken vocab must place `added_tokens_decoder` tokens on their declared ids.++    Repos such as `moonshotai/Kimi-K2.7-Code` ship a tiktoken vocab and declare their special tokens at sparse ids in a+    reserved range above it. If the ids are appended in a contiguous block, the tokenizer will silently produce a false+    mapping. The tests below mirror that shape at a small scale so the test needs no network.+    """++    BASE_TOKENS = list("abcdefghijklmnop")+    # Sparse and non-contiguous, like the reserved range of a real tiktoken checkpoint: a hole at 19+    # and a far tail pair, so declaration order and id order disagree.+    ADDED_TOKENS = {16: "[BOS]", 17: "[EOS]", 18: "<|im_end|>", 20: "<|start_header_id|>", 30: "[UNK]", 31: "[PAD]"}++    def _write_tiktoken_repo(self, directory, base_tokens, added_tokens):+        with open(os.path.join(directory, "tiktoken.model"), "w") as vocab_file:+            vocab_file.writelines(+                f"{base64.b64encode(token.encode()).decode()} {rank}\n" for rank, token in enumerate(base_tokens)+            )+        config = {+            "tokenizer_class": "PreTrainedTokenizerFast",+            "added_tokens_decoder": {+                str(token_id): {+                    "content": content,+                    "lstrip": False,+                    "normalized": False,+                    "rstrip": False,+                    "single_word": False,+                    "special": True,+                }+                for token_id, content in added_tokens.items()+            },+        }+        with open(os.path.join(directory, "tokenizer_config.json"), "w") as config_file:+            json.dump(config, config_file)++    def test_added_tokens_keep_their_declared_ids(self):+        """Every declared added token resolves to the id its config declares."""+        with tempfile.TemporaryDirectory() as directory:+            self._write_tiktoken_repo(directory, self.BASE_TOKENS, self.ADDED_TOKENS)+            tokenizer = AutoTokenizer.from_pretrained(directory)++            misplaced = {+                content: (token_id, tokenizer.convert_tokens_to_ids(content))+                for token_id, content in self.ADDED_TOKENS.items()+                if tokenizer.convert_tokens_to_ids(content) != token_id+            }+            self.assertEqual(misplaced, {}, f"tokens landed on the wrong ids (declared, actual): {misplaced}")++    def test_gaps_between_added_tokens_are_padded(self):+        """Ids with no declared token are filled with reserved placeholders, not left to collapse."""+        with tempfile.TemporaryDirectory() as directory:+            self._write_tiktoken_repo(directory, self.BASE_TOKENS, self.ADDED_TOKENS)+            tokenizer = AutoTokenizer.from_pretrained(directory)++            # 19 is the hole between "<|im_end|>" (18) and "<|start_header_id|>" (20).+            self.assertEqual(tokenizer.convert_ids_to_tokens(19), "<|reserved_token_19|>")+            # The vocabulary runs to the highest declared id and no further.+            self.assertEqual(len(tokenizer), max(self.ADDED_TOKENS) + 1)++ @require_tokenizers class TokenizerVersioningTest(unittest.TestCase):     def test_local_versioning(self):