diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index d1842c567..31c2637bb 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -21,6 +21,18 @@ class ThinkingMode(str, Enum): import jinja2 import orjson from fastapi import Request + +try: + from mistral_common.exceptions import MistralCommonException + + _MISTRAL_COMMON_ERRORS: tuple[type[BaseException], ...] = (MistralCommonException,) +except ImportError: + _MISTRAL_COMMON_ERRORS = () + +_CHAT_TEMPLATE_CLIENT_ERRORS: tuple[type[BaseException], ...] = ( + jinja2.TemplateError, + TypeError, +) + _MISTRAL_COMMON_ERRORS from fastapi.responses import ORJSONResponse, StreamingResponse from jsonschema import Draft202012Validator, SchemaError @@ -1367,7 +1379,7 @@ class OpenAIServingChat(OpenAIServingBase): prompt_ids = self.tokenizer_manager.tokenizer.encode( rendered_prompt, **encode_kwargs ) - except (jinja2.TemplateError, TypeError) as template_error: + except _CHAT_TEMPLATE_CLIENT_ERRORS as template_error: # Template errors (e.g., from raise_exception in Jinja templates) # and TypeError (e.g., tojson filter on Jinja2 Undefined variables) # should be treated as client errors (400 BadRequest) diff --git a/python/sglang/srt/utils/hf_transformers/mistral_utils.py b/python/sglang/srt/utils/hf_transformers/mistral_utils.py index f75d65dca..8cc17278b 100644 --- a/python/sglang/srt/utils/hf_transformers/mistral_utils.py +++ b/python/sglang/srt/utils/hf_transformers/mistral_utils.py @@ -634,4 +634,47 @@ def patch_mistral_common_tokenizer(tokenizer): return tokenizer._orig_apply_chat_template(messages, **kwargs) tokenizer.apply_chat_template = _safe_apply_chat_template + + def init_xgrammar(): + from xgrammar import TokenizerInfo + + tekken = getattr( + getattr(tokenizer.tokenizer, "instruct_tokenizer", None), "tokenizer", None + ) + if tekken is None or not hasattr(tekken, "id_to_byte_piece"): + logger.warning( + "Cannot build XGrammar TokenizerInfo: no Tekkenizer found under %s", + type(tokenizer).__name__, + ) + return None, None + + try: + placeholder = "<|xg_special_token_{}|>" + encoded_vocab = [] + for token_id in range(tekken.n_words): + piece = ( + tekken.id_to_piece(token_id) + if token_id < tekken.num_special_tokens + else tekken.id_to_byte_piece(token_id) + ) + # XGrammar reserves b"\x00"-prefixed tokens as special markers. + if isinstance(piece, bytes) and piece.startswith(b"\x00"): + piece = placeholder.format(f"nul{token_id}") + encoded_vocab.append(piece) + + eos_token_id = getattr(tokenizer, "eos_token_id", None) + override_stop_tokens = [eos_token_id] if eos_token_id is not None else None + tokenizer_info = TokenizerInfo( + encoded_vocab, stop_token_ids=override_stop_tokens + ) + except Exception as e: + logger.warning( + "Failed to build XGrammar TokenizerInfo for %s: %s", + type(tokenizer).__name__, + e, + ) + return None, None + return tokenizer_info, override_stop_tokens + + tokenizer.init_xgrammar = init_xgrammar return tokenizer diff --git a/test/registered/unit/constrained/test_mistral_common_xgrammar.py b/test/registered/unit/constrained/test_mistral_common_xgrammar.py new file mode 100644 index 000000000..9fe118585 --- /dev/null +++ b/test/registered/unit/constrained/test_mistral_common_xgrammar.py @@ -0,0 +1,112 @@ +import sys + +import pytest + +from sglang.srt.utils.hf_transformers.mistral_utils import ( + patch_mistral_common_tokenizer, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + +VOCAB_SIZE = 300 +NUM_SPECIAL = 8 + + +class _StubTekkenizer: + def __init__( + self, + vocab_size=VOCAB_SIZE, + num_special=NUM_SPECIAL, + fail_on_byte_piece=False, + ): + self.n_words = vocab_size + self.num_special_tokens = num_special + self.fail_on_byte_piece = fail_on_byte_piece + + def id_to_piece(self, token_id): + return f"" + + def id_to_byte_piece(self, token_id): + if self.fail_on_byte_piece: + raise RuntimeError("byte-piece conversion failed") + if token_id == self.num_special_tokens: + return b"\x00" + return bytes([token_id % 256]) + + +class _StubMistralTokenizer: + def __init__(self, tekken=None): + inner = type("InstructTokenizer", (), {"tokenizer": tekken})() + self.tokenizer = type("MistralTokenizer", (), {"instruct_tokenizer": inner})() + self.eos_token_id = 2 + self.chat_template = "x" + + def add_special_tokens(self, *args, **kwargs): + return 0 + + def convert_tokens_to_ids(self, val): + return 0 + + def decode(self, *args, **kwargs): + return "" + + def batch_decode(self, *args, **kwargs): + return [] + + def apply_chat_template(self, *args, **kwargs): + return [] + + +class _MistralCommonStub(_StubMistralTokenizer): + pass + + +def _patched(tekken): + return patch_mistral_common_tokenizer(_MistralCommonStub(tekken)) + + +def _is_allowed(mask, token_id): + return bool((int(mask[0][token_id // 32]) >> (token_id % 32)) & 1) + + +def test_builds_tokenizer_info_over_full_vocab(): + info, stop_tokens = _patched(_StubTekkenizer()).init_xgrammar() + + assert info is not None + assert info.vocab_size == VOCAB_SIZE + assert stop_tokens == [2] + + +def test_json_schema_compiles_and_constrains(): + from xgrammar import GrammarCompiler, GrammarMatcher, allocate_token_bitmask + + info, _ = _patched(_StubTekkenizer()).init_xgrammar() + grammar = GrammarCompiler(tokenizer_info=info).compile_json_schema( + '{"type":"object","properties":{"a":{"type":"integer"}},"required":["a"]}' + ) + mask = allocate_token_bitmask(1, info.vocab_size) + GrammarMatcher(grammar).fill_next_token_bitmask(mask) + + assert _is_allowed(mask, ord("{")) + assert not _is_allowed(mask, ord("z")) + + +def test_returns_none_without_a_tekkenizer(): + info, stop_tokens = _patched(object()).init_xgrammar() + + assert info is None + assert stop_tokens is None + + +def test_returns_none_when_vocab_extraction_fails(): + info, stop_tokens = _patched( + _StubTekkenizer(fail_on_byte_piece=True) + ).init_xgrammar() + + assert info is None + assert stop_tokens is None + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__]))