diff --git a/mlx_lm/tokenizer_utils.py b/mlx_lm/tokenizer_utils.py index 4bb44eab1..9d2f14f5e 100644 --- a/mlx_lm/tokenizer_utils.py +++ b/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] diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py index d5c81e34a..ad9e5d4a6 100644 --- a/tests/test_tokenizers.py +++ b/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()