117 lines
3.9 KiB
Python
117 lines
3.9 KiB
Python
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
|
|
|
|
|
|
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
|