Fix detokenizer for byte-level tokenizers with an SPM-style decoder (fixes #1041)

mlx_lm/tokenizer_utils.py

@@ -543,6 +543,26 @@ def _is_bpe_decoder(decoder):
     return isinstance(decoder, dict) and decoder.get("type", None) == "ByteLevel"
 
 
+def _has_byte_level_pretokenizer(tokenizer_content):
+    """Whether the tokenizer encodes its vocabulary with byte-level markers.
+
+    A ``ByteLevel`` pre-tokenizer (possibly nested in a ``Sequence``) means the
+    vocabulary uses GPT-2-style byte markers (e.g. ``Ġ`` for spaces), which only
+    ``BPEStreamingDetokenizer`` knows how to decode -- regardless of what the
+    ``decoder`` field looks like.
+    """
+
+    def _is_byte_level(node):
+        return isinstance(node, dict) and node.get("type", None) == "ByteLevel"
+
+    pre_tokenizer = tokenizer_content.get("pre_tokenizer")
+    if _is_byte_level(pre_tokenizer):
+        return True
+    if isinstance(pre_tokenizer, dict) and pre_tokenizer.get("type") == "Sequence":
+        return any(_is_byte_level(p) for p in pre_tokenizer.get("pretokenizers", []))
+    return False
+
+
 def _infer_tool_parser(chat_template):
     """Attempt to auto-infer a tool parser from the chat template."""
     if not isinstance(chat_template, str):
@@ -596,12 +616,20 @@ def load(
                 raise JSONDecodeError("Failed to parse tokenizer.json", e.doc, e.pos)
 
         if "decoder" in tokenizer_content:
-            if _is_spm_decoder(tokenizer_content["decoder"]):
+            decoder = tokenizer_content["decoder"]
+            if _is_bpe_decoder(decoder) or _has_byte_level_pretokenizer(
+                tokenizer_content
+            ):
+                # A byte-level vocabulary must use the BPE detokenizer even when
+                # the tokenizer ships an SPM-style decoder (e.g. Mistral tekken
+                # v13): the SPM detokenizer only strips the ``▁`` marker and so
+                # leaves the byte-level space marker ``Ġ`` (U+0120) in the
+                # output. See issue #1041.
+                detokenizer_class = BPEStreamingDetokenizer
+            elif _is_spm_decoder(decoder):
                 detokenizer_class = SPMStreamingDetokenizer
-            elif _is_spm_decoder_no_space(tokenizer_content["decoder"]):
+            elif _is_spm_decoder_no_space(decoder):
                 detokenizer_class = partial(SPMStreamingDetokenizer, trim_space=False)
-            elif _is_bpe_decoder(tokenizer_content["decoder"]):
-                detokenizer_class = BPEStreamingDetokenizer
 
     if isinstance(eos_token_ids, int):
         eos_token_ids = [eos_token_ids]

tests/test_tokenizers.py

@@ -10,6 +10,7 @@
     NaiveStreamingDetokenizer,
     SPMStreamingDetokenizer,
     TokenizerWrapper,
+    _has_byte_level_pretokenizer,
 )
 from mlx_lm.utils import load_tokenizer
 
@@ -123,6 +124,53 @@ def test_find_token(self):
         self.assertEqual(find(prompt, [THINK_START], start=0), 1)
         self.assertEqual(find(prompt, [THINK_START], start=0, reverse=True), 3)
 
+    def test_has_byte_level_pretokenizer(self):
+        byte_level = {"type": "ByteLevel"}
+        self.assertTrue(_has_byte_level_pretokenizer({"pre_tokenizer": byte_level}))
+        self.assertTrue(
+            _has_byte_level_pretokenizer(
+                {
+                    "pre_tokenizer": {
+                        "type": "Sequence",
+                        "pretokenizers": [{"type": "Split"}, byte_level],
+                    }
+                }
+            )
+        )
+        self.assertFalse(
+            _has_byte_level_pretokenizer(
+                {"pre_tokenizer": {"type": "Metaspace", "replacement": "▁"}}
+            )
+        )
+        self.assertFalse(_has_byte_level_pretokenizer({}))
+
+    def test_byte_level_vocab_with_spm_decoder(self):
+        # Mistral tekken v13 tokenizers carry an SPM-style decoder but a
+        # byte-level vocabulary (ByteLevel pre-tokenizer). They must use the BPE
+        # detokenizer; the SPM detokenizer only strips the "▁" marker and so
+        # leaves the byte-level space marker "Ġ" (U+0120) in the output.
+        # Regression test for #1041. Download tokenizer files only -- the model
+        # weights are irrelevant to detokenization and very large.
+        repo = "mlx-community/Devstral-Small-2-24B-Instruct-2512-bf16"
+        path = Path(
+            snapshot_download(
+                repo,
+                allow_patterns=["*.json", "*.model", "*tokenizer*", "*.jinja"],
+            )
+        )
+        tokenizer = load_tokenizer(path)
+        self.assertIsInstance(tokenizer.detokenizer, BPEStreamingDetokenizer)
+
+        text = "Hello! How can I assist you today?"
+        tokens = tokenizer.encode(text, add_special_tokens=False)
+        detokenizer = tokenizer.detokenizer
+        detokenizer.reset()
+        for t in tokens:
+            detokenizer.add_token(t)
+        detokenizer.finalize()
+        self.assertNotIn("Ġ", detokenizer.text)
+        self.assertEqual(detokenizer.text, text)
+
 
 if __name__ == "__main__":
     unittest.main()