chore: import upstream snapshot with attribution

python/sglang/srt/utils/patch_tokenizer.py

@@ -0,0 +1,127 @@
+import logging
+
+from sglang.srt.environ import envs
+
+logger = logging.getLogger(__name__)
+
+
+def patch_tokenizer(tokenizer):
+    if not envs.SGLANG_PATCH_TOKENIZER.get():
+        return tokenizer
+
+    if _is_kimi_tiktoken_tokenizer(tokenizer):
+        logger.info(
+            f"Applying special tokens cache patch for Kimi tokenizer: {type(tokenizer)}"
+        )
+        return _SpecialTokensCachePatcher.patch(tokenizer)
+
+    return tokenizer
+
+
+def unpatch_tokenizer(tokenizer):
+    return _SpecialTokensCachePatcher.unpatch(tokenizer)
+
+
+def _is_kimi_tiktoken_tokenizer(tokenizer):
+    cls = type(tokenizer)
+    class_name = cls.__name__
+    module_name = cls.__module__ or ""
+    return class_name == "TikTokenTokenizer" and "tokenization_kimi" in module_name
+
+
+def decode_without_hf_kwargs(tokenizer, token_ids, skip_special_tokens):
+    if skip_special_tokens:
+        special_ids = getattr(tokenizer, "all_special_ids_set", None)
+        if special_ids is None:
+            special_ids = getattr(tokenizer, "all_special_ids", None)
+        if special_ids is not None:
+            special_ids_set = set(special_ids)
+            token_ids = [tid for tid in token_ids if tid not in special_ids_set]
+    return tokenizer.decode(token_ids)
+
+
+class _SpecialTokensCachePatcher:
+    _PATCHED_FLAG = "_sglang_special_tokens_patched"
+    _CACHED_TOKENS_ATTR = "_sglang_cached_special_tokens"
+    _CACHED_IDS_ATTR = "_sglang_cached_special_ids"
+
+    @classmethod
+    def patch(cls, tokenizer):
+        tokenizer_cls = type(tokenizer)
+
+        if getattr(tokenizer_cls, cls._PATCHED_FLAG, False):
+            return tokenizer
+
+        tokenizer_cls._original_all_special_tokens = (
+            tokenizer_cls.all_special_tokens.fget
+        )
+        tokenizer_cls._original_all_special_ids = tokenizer_cls.all_special_ids.fget
+        tokenizer_cls._original_add_special_tokens = tokenizer_cls.add_special_tokens
+        tokenizer_cls._original_add_tokens = tokenizer_cls.add_tokens
+
+        patched_all_special_tokens = _make_cached_property(
+            cls._CACHED_TOKENS_ATTR, tokenizer_cls._original_all_special_tokens
+        )
+        patched_all_special_ids = _make_cached_property(
+            cls._CACHED_IDS_ATTR, tokenizer_cls._original_all_special_ids
+        )
+
+        def patched_add_special_tokens(self, *args, **kwargs):
+            assert (
+                False
+            ), "Cannot modify special tokens after patch. Call unpatch_tokenizer first."
+
+        def patched_add_tokens(self, new_tokens, special_tokens=False):
+            assert (
+                not special_tokens
+            ), "Cannot add special tokens after patch. Call unpatch_tokenizer first."
+            return tokenizer_cls._original_add_tokens(
+                self, new_tokens, special_tokens=False
+            )
+
+        tokenizer_cls.all_special_tokens = patched_all_special_tokens
+        tokenizer_cls.all_special_ids = patched_all_special_ids
+        tokenizer_cls.add_special_tokens = patched_add_special_tokens
+        tokenizer_cls.add_tokens = patched_add_tokens
+        setattr(tokenizer_cls, cls._PATCHED_FLAG, True)
+
+        return tokenizer
+
+    @classmethod
+    def unpatch(cls, tokenizer):
+        tokenizer_cls = type(tokenizer)
+
+        if not getattr(tokenizer_cls, cls._PATCHED_FLAG, False):
+            return tokenizer
+
+        tokenizer_cls.all_special_tokens = property(
+            tokenizer_cls._original_all_special_tokens
+        )
+        tokenizer_cls.all_special_ids = property(
+            tokenizer_cls._original_all_special_ids
+        )
+        tokenizer_cls.add_special_tokens = tokenizer_cls._original_add_special_tokens
+        tokenizer_cls.add_tokens = tokenizer_cls._original_add_tokens
+
+        del tokenizer_cls._original_all_special_tokens
+        del tokenizer_cls._original_all_special_ids
+        del tokenizer_cls._original_add_special_tokens
+        del tokenizer_cls._original_add_tokens
+        delattr(tokenizer_cls, cls._PATCHED_FLAG)
+
+        for attr in [cls._CACHED_TOKENS_ATTR, cls._CACHED_IDS_ATTR]:
+            if hasattr(tokenizer, attr):
+                delattr(tokenizer, attr)
+
+        logger.info(f"Unpatched special tokens cache for {tokenizer_cls.__name__}")
+        return tokenizer
+
+
+def _make_cached_property(cache_attr, original_fn):
+    @property
+    def cached_prop(self):
+        if getattr(self, cache_attr, None) is None:
+            setattr(self, cache_attr, original_fn(self))
+        return getattr(self, cache_attr)
+
+    return cached_prop