From 6d5cfe697c232307ee07ce03d3d69c013cde3390 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebasti=C3=A1n=20Real?= Date: Mon, 31 Aug 2026 10:22:05 -0500 Subject: [PATCH] Fix GGUF user-defined token atomicity --- python/freetoken/models/gguf/tokenizer.py | 38 ++++++++ tests/models/test_gguf_tokenizer.py | 110 ++++++++++++++++++++++ 2 files changed, 148 insertions(+) create mode 100644 tests/models/test_gguf_tokenizer.py diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 6d5481c17..4ac42b92d 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -10,6 +10,8 @@ from typing import Any +from tokenizers import AddedToken + from .reader import gguf_architecture, load_gguf_metadata # GGUF architecture -> transformers GGUF tokenizer-converter key. @@ -46,6 +48,42 @@ def tok_for(id_key: str, default: str) -> str: unk_token=tok_for("unknown_token_id", ""), pad_token=tok_for("padding_token_id", ""), ) + + # GGUF user-defined tokens already exist in the model vocabulary with fixed + # IDs, but transformers' GGUF conversion may not register them as AddedToken + # entries. In that state convert_tokens_to_ids("") returns the correct + # vocabulary ID while encode("") incorrectly splits it into ordinary + # subword tokens. Restore their atomic-token behavior without marking them + # special and without changing vocabulary size or IDs. + token_types = tok_dict.get("token_type") + if token_types is not None: + added_tokens = [] + for token_id, (token, token_type) in enumerate(zip(tokens, token_types)): + if int(token_type) != 4: + continue + if tokenizer.convert_tokens_to_ids(token) != token_id: + continue + if tokenizer.encode(token, add_special_tokens=False) == [token_id]: + continue + added_tokens.append( + AddedToken( + token, + single_word=False, + lstrip=False, + rstrip=False, + normalized=False, + special=False, + ) + ) + + if added_tokens: + vocab_size_before = len(tokenizer) + tokenizer.add_tokens(added_tokens, special_tokens=False) + if len(tokenizer) != vocab_size_before: + raise RuntimeError( + "restoring GGUF user-defined tokens unexpectedly changed vocabulary size" + ) + chat_template = meta.get("tokenizer.chat_template") if chat_template: tokenizer.chat_template = chat_template diff --git a/tests/models/test_gguf_tokenizer.py b/tests/models/test_gguf_tokenizer.py new file mode 100644 index 000000000..ef69eaba5 --- /dev/null +++ b/tests/models/test_gguf_tokenizer.py @@ -0,0 +1,110 @@ +from __future__ import annotations + + +def test_gguf_user_defined_tokens_preserve_atomic_vocab_ids(monkeypatch): + """GGUF USER_DEFINED tokens must remain atomic after GGUF -> HF conversion. + + Regression: Qwen3.6 GGUF contains and as USER_DEFINED + vocabulary entries. transformers' GGUF conversion can preserve their vocab + IDs while still allowing the pre-tokenizer to split their text into ordinary + subword tokens. That changes the actual prompt token IDs seen by the model. + """ + from tokenizers import Tokenizer + from tokenizers.models import WordLevel + from tokenizers.pre_tokenizers import Whitespace + + import freetoken.models.gguf.tokenizer as gguf_tokenizer + + tokens = [ + "", # 0 + "", # 1 + "", # 2 + "", # 3 + "<", # 4 + "think", # 5 + ">", # 6 + "", # 8 USER_DEFINED + "", # 9 USER_DEFINED + "hello", # 10 + ] + + # GGUF TokenType: + # NORMAL=1, UNKNOWN=2, CONTROL=3, USER_DEFINED=4. + token_types = [ + 2, + 3, + 3, + 3, + 1, + 1, + 1, + 1, + 4, + 4, + 1, + ] + + metadata = { + "tokenizer.ggml.tokens": tokens, + "tokenizer.ggml.token_type": token_types, + "tokenizer.ggml.unknown_token_id": 0, + "tokenizer.ggml.bos_token_id": 1, + "tokenizer.ggml.eos_token_id": 2, + "tokenizer.ggml.padding_token_id": 3, + } + + backend = Tokenizer( + WordLevel( + vocab={token: idx for idx, token in enumerate(tokens)}, + unk_token="", + ) + ) + + # This intentionally reproduces the bug: punctuation is pre-tokenized, + # so the vocabulary entry exists at ID 8 but plain encoding would + # otherwise produce "<" + "think" + ">". + backend.pre_tokenizer = Whitespace() + + monkeypatch.setattr( + gguf_tokenizer, + "load_gguf_metadata", + lambda _path: metadata, + ) + monkeypatch.setattr( + gguf_tokenizer, + "gguf_architecture", + lambda _path: "synthetic", + ) + + def fake_convert_gguf_tokenizer(_arch, _tok_dict): + return backend, {} + + monkeypatch.setattr( + "transformers.integrations.ggml.convert_gguf_tokenizer", + fake_convert_gguf_tokenizer, + ) + + tokenizer = gguf_tokenizer.load_gguf_tokenizer("synthetic.gguf") + + # Independent source of truth: these IDs come from the synthetic GGUF + # vocabulary above, not from the implementation under test. + assert tokenizer.convert_tokens_to_ids("") == 8 + assert tokenizer.convert_tokens_to_ids("") == 9 + + assert tokenizer.encode( + "", + add_special_tokens=False, + ) == [8] + + assert tokenizer.encode( + "", + add_special_tokens=False, + ) == [9] + + # Restoring atomicity must not create new embedding/vocabulary IDs. + assert len(tokenizer) == len(tokens) + + # USER_DEFINED does not mean HF "special token". + assert "" not in tokenizer.all_special_tokens + assert "" not in tokenizer.all_special_tokens