From ebfc90571f051715ac6e1b182ddf43ad22c3c7d4 Mon Sep 17 00:00:00 2001 From: xiaoweis <39303645+Shixiaowei02@users.noreply.github.com> Date: Mon, 13 Jul 2026 08:08:56 -0700 Subject: [PATCH] [None][fix] Refresh special-token sets lazily for slow detokenization TransformersTokenizer snapshotted the tokenizer's special tokens in __init__ (all_special_ids for slow convert_ids_to_tokens, all_special_tokens for convert_tokens_to_string) and reused them forever. That snapshot is taken too early: a model's input processor can register special tokens after the wrapper is built but before serving (e.g. gemma4 adds <|video|>), so the affected tokens stopped being skipped and leaked into detokenized output when skip_special_tokens=True -- an accuracy regression. Compute the (ids, strings) sets once, lazily, on the first detokenization call. By request time all build-time special-token registration is done and nothing mutates it further, so a first-use cache is correct; it also keeps the slow-tokenizer detok path O(1) per streamed token (avoiding both the original per-id all_special_ids rebuild and any per-call refresh). Add regression tests on a genuine slow tokenizer: special tokens registered after wrapping must be reflected by both convert_ids_to_tokens and convert_tokens_to_string. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> --- tensorrt_llm/tokenizer/tokenizer.py | 53 +++++++------- .../unittest/llmapi/test_tokenizer_decode.py | 73 +++++++++++++++++++ 2 files changed, 99 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/tokenizer/tokenizer.py b/tensorrt_llm/tokenizer/tokenizer.py index 1422cc3804b5..1c2694a048c6 100644 --- a/tensorrt_llm/tokenizer/tokenizer.py +++ b/tensorrt_llm/tokenizer/tokenizer.py @@ -197,19 +197,8 @@ class TransformersTokenizer(TokenizerBase): def __init__(self, tokenizer): self.tokenizer = tokenizer - if hasattr(self.tokenizer, "all_special_tokens"): - self._all_special_tokens_set = set( - self.tokenizer.all_special_tokens) - else: - self._all_special_tokens_set = set() - # Cache of special-token ids used by convert_ids_to_tokens to work - # around an O(N*K) bug in transformers 5.x's slow tokenizer path - # (see convert_ids_to_tokens below). Computed once here so we don't - # rebuild it on every per-token streaming call. - try: - self._all_special_ids_set = set(self.tokenizer.all_special_ids) - except (AttributeError, NotImplementedError): - self._all_special_ids_set = set() + # (ids, strings) special-token sets, computed lazily on first detok. + self._special_tokens_sets = None def __reduce__(self): # In multi-node scenarios, AutoTokenizer.from_pretrained with @@ -334,24 +323,31 @@ def convert_ids_to_tokens( inner = self.tokenizer if (isinstance(ids, int) or not skip_special_tokens or getattr(inner, "is_fast", False)): - # Single id: no loop. Fast tokenizer: tokenization_utils_tokenizers - # already caches all_special_ids before the loop. Without skip: - # the buggy branch is short-circuited. + # Fast / no-skip / single-id paths skip the slow O(N*K) branch below. return inner.convert_ids_to_tokens( ids, skip_special_tokens=skip_special_tokens) - # Slow PreTrainedTokenizer.convert_ids_to_tokens in transformers 5.x - # tests `index in self.all_special_ids` inside the per-id loop, and - # `all_special_ids` is an @property that rebuilds the list on every - # access via convert_tokens_to_ids(self.all_special_tokens). The fast - # subclass already hoists this out of the loop (see - # tokenization_utils_tokenizers.py), but the slow base class wasn't - # updated. Mirror that fix using the set we cached in __init__, - # which keeps streaming detokenization at O(1) all-special-ids cost - # per call regardless of batch size or stream_interval. + # Slow path rebuilds all_special_ids per id; cache it: O(N*K) -> O(N). tokens = inner.convert_ids_to_tokens(ids, skip_special_tokens=False) - special_ids = self._all_special_ids_set + special_ids = self._special_tokens()[0] return [t for idx, t in zip(ids, tokens) if int(idx) not in special_ids] + def _special_tokens(self) -> Tuple[set, set]: + """(ids, strings) special-token sets, computed once on first detok. + + Deferred from __init__, which is too early: an input processor may + register specials after wrapping but before serving (e.g. gemma4's + <|video|>). Detok starts later and nothing mutates specials mid-stream, + so a first-use cache is correct and O(1) per streamed token. + """ + if self._special_tokens_sets is None: + inner = self.tokenizer + try: + self._special_tokens_sets = (set(inner.all_special_ids), + set(inner.all_special_tokens)) + except (AttributeError, NotImplementedError): + self._special_tokens_sets = (set(), set()) + return self._special_tokens_sets + def convert_tokens_to_string( self, tokens: List[str], @@ -362,10 +358,13 @@ def convert_tokens_to_string( if self.is_fast or not self.get_added_vocab(): return self.tokenizer.convert_tokens_to_string(tokens) + special_tokens = () + if skip_special_tokens: + special_tokens = self._special_tokens()[1] sub_texts: List[str] = [] current_sub_text: List[str] = [] for token in tokens: - if skip_special_tokens and token in self._all_special_tokens_set: + if skip_special_tokens and token in special_tokens: continue if token in self.get_added_vocab(): if current_sub_text: diff --git a/tests/unittest/llmapi/test_tokenizer_decode.py b/tests/unittest/llmapi/test_tokenizer_decode.py index b1df835b29b6..d4c601a824a6 100644 --- a/tests/unittest/llmapi/test_tokenizer_decode.py +++ b/tests/unittest/llmapi/test_tokenizer_decode.py @@ -17,10 +17,83 @@ from tokenizers.decoders import ByteFallback, Fuse, Sequence from tokenizers.models import WordLevel from transformers import PreTrainedTokenizerFast +from transformers.tokenization_utils import PreTrainedTokenizer from tensorrt_llm.tokenizer import TransformersTokenizer +class _ToySlowTokenizer(PreTrainedTokenizer): + """A minimal genuine *slow* tokenizer (``is_fast`` is False). + + The slow ``convert_ids_to_tokens`` path is the only one that filters + special tokens in TransformersTokenizer's wrapper; fast tokenizers delegate + straight to the backend. This toy tokenizer lets us exercise that path + deterministically without downloading weights. + """ + + def __init__(self, **kwargs): + self._vocab = {f"tok{i}": i for i in range(64)} + self._ids = {v: k for k, v in self._vocab.items()} + super().__init__(**kwargs) + + @property + def vocab_size(self) -> int: + return len(self._vocab) + + def get_vocab(self): + return dict(self._vocab) + + def _convert_token_to_id(self, token): + return self._vocab.get(token, 0) + + def _convert_id_to_token(self, index): + return self._ids.get(int(index), "tok0") + + def _tokenize(self, text): + return text.split() + + +def test_convert_ids_to_tokens_reflects_special_tokens_registered_after_wrap() -> None: + # Specials registered after wrapping (before first detok) must be reflected. + inner = _ToySlowTokenizer() + inner.add_special_tokens( + { + "bos_token": "tok1", + "eos_token": "tok2", + "additional_special_tokens": ["tok5"], + } + ) + tokenizer = TransformersTokenizer(inner) + assert inner.is_fast is False + + inner.add_special_tokens({"additional_special_tokens": ["tok5", "tok42"]}) + + ids = [10, 42, 11, 5, 12] + got = tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=True) + assert got == inner.convert_ids_to_tokens(ids, skip_special_tokens=True) + assert got == ["tok10", "tok11", "tok12"] + + +def test_convert_tokens_to_string_reflects_special_tokens_registered_after_wrap() -> None: + # The same must hold for convert_tokens_to_string's skip path. + inner = _ToySlowTokenizer() + inner.add_special_tokens( + { + "bos_token": "tok1", + "eos_token": "tok2", + "additional_special_tokens": ["tok5"], + } + ) + tokenizer = TransformersTokenizer(inner) + assert inner.is_fast is False + + inner.add_special_tokens({"additional_special_tokens": ["tok5", "tok42"]}) + + tokens = ["tok10", "tok42", "tok11"] + got = tokenizer.convert_tokens_to_string(tokens, skip_special_tokens=True) + assert got == "tok10 tok11" + + def test_hf_decode_incrementally_recovers_from_invalid_prefix() -> None: backend = Tokenizer( WordLevel(