Optimize detokenization without HF decode kwargs (#25309)
This commit is contained in:
@@ -44,6 +44,7 @@ from sglang.srt.utils import (
|
|||||||
)
|
)
|
||||||
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
||||||
from sglang.srt.utils.network import get_zmq_socket
|
from sglang.srt.utils.network import get_zmq_socket
|
||||||
|
from sglang.srt.utils.patch_tokenizer import decode_without_hf_kwargs
|
||||||
from sglang.srt.utils.watchdog import Watchdog
|
from sglang.srt.utils.watchdog import Watchdog
|
||||||
from sglang.utils import (
|
from sglang.utils import (
|
||||||
TypeBasedDispatcher,
|
TypeBasedDispatcher,
|
||||||
@@ -190,6 +191,12 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
|
|||||||
) -> List[str]:
|
) -> List[str]:
|
||||||
"""Batch decode with grouping by (skip_special_tokens, spaces_between_special_tokens)."""
|
"""Batch decode with grouping by (skip_special_tokens, spaces_between_special_tokens)."""
|
||||||
|
|
||||||
|
if not getattr(self.tokenizer, "is_fast", False):
|
||||||
|
return [
|
||||||
|
decode_without_hf_kwargs(self.tokenizer, ids, skip)
|
||||||
|
for ids, skip in zip(ids_list, skip_list)
|
||||||
|
]
|
||||||
|
|
||||||
# fast path
|
# fast path
|
||||||
first_skip, first_space = skip_list[0], space_list[0]
|
first_skip, first_space = skip_list[0], space_list[0]
|
||||||
if all(
|
if all(
|
||||||
|
|||||||
@@ -29,6 +29,15 @@ def _is_kimi_tiktoken_tokenizer(tokenizer):
|
|||||||
return class_name == "TikTokenTokenizer" and "tokenization_kimi" in module_name
|
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 = set(tokenizer.all_special_ids)
|
||||||
|
token_ids = [tid for tid in token_ids if tid not in special_ids]
|
||||||
|
return tokenizer.decode(token_ids)
|
||||||
|
|
||||||
|
|
||||||
class _SpecialTokensCachePatcher:
|
class _SpecialTokensCachePatcher:
|
||||||
_PATCHED_FLAG = "_sglang_special_tokens_patched"
|
_PATCHED_FLAG = "_sglang_special_tokens_patched"
|
||||||
_CACHED_TOKENS_ATTR = "_sglang_cached_special_tokens"
|
_CACHED_TOKENS_ATTR = "_sglang_cached_special_tokens"
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ from transformers import AutoTokenizer
|
|||||||
|
|
||||||
from sglang.srt.utils.patch_tokenizer import (
|
from sglang.srt.utils.patch_tokenizer import (
|
||||||
_SpecialTokensCachePatcher,
|
_SpecialTokensCachePatcher,
|
||||||
|
decode_without_hf_kwargs,
|
||||||
unpatch_tokenizer,
|
unpatch_tokenizer,
|
||||||
)
|
)
|
||||||
from sglang.test.ci.ci_register import register_cpu_ci
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
@@ -150,6 +151,19 @@ class TestPatchTokenizerUnitTest(unittest.TestCase):
|
|||||||
|
|
||||||
unpatch_tokenizer(tokenizer)
|
unpatch_tokenizer(tokenizer)
|
||||||
|
|
||||||
|
def test_decode_without_hf_kwargs_uses_native_decode(self):
|
||||||
|
tokenizer = _FakeDecodeTokenizer()
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
decode_without_hf_kwargs(tokenizer, [1, 99, 2], True),
|
||||||
|
"ab",
|
||||||
|
)
|
||||||
|
self.assertEqual(
|
||||||
|
decode_without_hf_kwargs(tokenizer, [1, 99, 2], False),
|
||||||
|
"a<special>b",
|
||||||
|
)
|
||||||
|
self.assertEqual(tokenizer.decode_calls, [[1, 2], [1, 99, 2]])
|
||||||
|
|
||||||
|
|
||||||
def _get_class_attr_ids(cls):
|
def _get_class_attr_ids(cls):
|
||||||
return {
|
return {
|
||||||
@@ -174,5 +188,18 @@ def _patched_tokenizer():
|
|||||||
unpatch_tokenizer(tokenizer)
|
unpatch_tokenizer(tokenizer)
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeDecodeTokenizer:
|
||||||
|
all_special_ids_set = {99}
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.decode_calls = []
|
||||||
|
|
||||||
|
def decode(self, token_ids):
|
||||||
|
token_ids = list(token_ids)
|
||||||
|
self.decode_calls.append(token_ids)
|
||||||
|
token_text = {1: "a", 2: "b", 99: "<special>"}
|
||||||
|
return "".join(token_text[token_id] for token_id in token_ids)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
Reference in New Issue
Block a user