[Test] Add unit tests for srt/tokenizer/tiktoken_tokenizer (#21107)

This commit is contained in:
ROSINE HE
2026-04-06 10:23:19 +08:00
committed by GitHub
parent 2f00e42555
commit 8d8aca8f26
@@ -0,0 +1,149 @@
"""Unit tests for tiktoken_tokenizer — no server, no model loading."""
import unittest
from unittest.mock import MagicMock, patch
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
from sglang.srt.tokenizer.tiktoken_tokenizer import (
CONTROL_TOKEN_TEXTS,
DEFAULT_CONTROL_TOKENS,
DEFAULT_SPECIAL_TOKENS,
EOS,
PAD,
RESERVED_TOKEN_TEXTS,
SEP,
TiktokenProcessor,
TiktokenTokenizer,
)
class TestConstants(CustomTestCase):
def test_reserved_token_count(self):
self.assertEqual(len(RESERVED_TOKEN_TEXTS), 125)
def test_reserved_token_format(self):
self.assertEqual(RESERVED_TOKEN_TEXTS[0], "<|reserved_3|>")
self.assertEqual(RESERVED_TOKEN_TEXTS[-1], "<|reserved_127|>")
def test_control_token_count(self):
self.assertEqual(len(CONTROL_TOKEN_TEXTS), 704)
def test_control_token_format(self):
self.assertEqual(CONTROL_TOKEN_TEXTS[0], "<|control1|>")
self.assertEqual(CONTROL_TOKEN_TEXTS[-1], "<|control704|>")
def test_special_token_values(self):
self.assertEqual(PAD, "<|pad|>")
self.assertEqual(EOS, "<|eos|>")
self.assertEqual(SEP, "<|separator|>")
def test_default_special_tokens_contains_all(self):
self.assertIn(PAD, DEFAULT_SPECIAL_TOKENS)
self.assertIn(EOS, DEFAULT_SPECIAL_TOKENS)
self.assertIn(SEP, DEFAULT_SPECIAL_TOKENS)
def test_default_control_tokens_values(self):
# Note: "sep" maps to EOS and "eos" maps to SEP in the source code
self.assertEqual(DEFAULT_CONTROL_TOKENS["pad"], PAD)
self.assertEqual(DEFAULT_CONTROL_TOKENS["sep"], EOS)
self.assertEqual(DEFAULT_CONTROL_TOKENS["eos"], SEP)
class TestTiktokenProcessor(CustomTestCase):
def setUp(self):
tokenizer_patcher = patch(
"sglang.srt.tokenizer.tiktoken_tokenizer.TiktokenTokenizer"
)
tokenizer_patcher.start()
self.addCleanup(tokenizer_patcher.stop)
self.processor = TiktokenProcessor(name="dummy")
def test_image_processor_returns_dict(self):
result = self.processor.image_processor("fake_image")
self.assertIsInstance(result, dict)
def test_image_processor_has_pixel_values_key(self):
result = self.processor.image_processor("fake_image")
self.assertIn("pixel_values", result)
def test_image_processor_wraps_image_in_list(self):
image = "fake_image_data"
result = self.processor.image_processor(image)
self.assertEqual(result["pixel_values"], [image])
def test_image_processor_with_none(self):
result = self.processor.image_processor(None)
self.assertEqual(result["pixel_values"], [None])
class TestTiktokenTokenizer(CustomTestCase):
def setUp(self):
from jinja2 import Template
self.tok = TiktokenTokenizer.__new__(TiktokenTokenizer)
self.mock_tokenizer = MagicMock()
self.tok.tokenizer = self.mock_tokenizer
self.tok.chat_template = "dummy"
self.tok.chat_template_jinja = Template(
"{% for message in messages %}"
"{{ message['role'] }}: {{ message['content'] }}"
"{% endfor %}"
"{% if add_generation_prompt %}assistant:{% endif %}"
)
def test_encode_delegates_to_tokenizer(self):
self.mock_tokenizer.encode.return_value = [1, 2, 3]
result = self.tok.encode("hello")
self.mock_tokenizer.encode.assert_called_once_with("hello")
self.assertEqual(result, [1, 2, 3])
def test_decode_delegates_to_tokenizer(self):
self.mock_tokenizer.decode.return_value = "hello"
result = self.tok.decode([1, 2, 3])
self.mock_tokenizer.decode.assert_called_once_with([1, 2, 3])
self.assertEqual(result, "hello")
def test_batch_decode_list_of_lists(self):
self.mock_tokenizer.decode_batch.return_value = ["hello", "world"]
result = self.tok.batch_decode([[1, 2], [3, 4]])
self.mock_tokenizer.decode_batch.assert_called_once_with([[1, 2], [3, 4]])
self.assertEqual(result, ["hello", "world"])
def test_batch_decode_flat_list_wraps_each(self):
self.mock_tokenizer.decode_batch.return_value = ["a", "b"]
self.tok.batch_decode([1, 2])
self.mock_tokenizer.decode_batch.assert_called_once_with([[1], [2]])
def test_call_returns_input_ids(self):
self.mock_tokenizer.encode.return_value = [1, 2, 3]
result = self.tok(["hello", "world"])
self.assertIn("input_ids", result)
self.assertEqual(len(result["input_ids"]), 2)
def test_apply_chat_template_no_tokenize(self):
messages = [{"role": "user", "content": "hello"}]
result = self.tok.apply_chat_template(
messages=messages,
tokenize=False,
add_generation_prompt=False,
)
self.assertIsInstance(result, str)
self.assertIn("hello", result)
def test_apply_chat_template_with_tokenize(self):
self.mock_tokenizer.encode.return_value = [1, 2, 3]
messages = [{"role": "user", "content": "hello"}]
result = self.tok.apply_chat_template(
messages=messages,
tokenize=True,
add_generation_prompt=False,
)
self.assertIsInstance(result, list)
if __name__ == "__main__":
unittest.main()