Add Mistral guidance

Signed-off-by: Julien Denize <[email protected]>

Author: juliendenize

tests/tool_parsers/test_mistral_tool_parser.py

@@ -81,7 +81,7 @@
     validate_request_params,
 )
 from vllm.tool_parsers import ToolParser
-from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
+from vllm.tool_parsers.mistral_tool_parser import MistralToolCall, MistralToolParser
 from vllm.tool_parsers.utils import partial_json_loads
 from vllm.utils.collection_utils import as_list
 from vllm.v1.sample.logits_processor import validate_logits_processors_parameters
@@ -142,6 +142,9 @@ def __init__(
             enable_auto_tools=enable_auto_tools,
             model_name=self.model_config.model,
         )
+        if self.tool_parser == MistralToolParser and self.reasoning_parser is not None:
+            self.tool_parser.reasoning = True
+
         self.exclude_tools_when_tool_choice_none = exclude_tools_when_tool_choice_none
 
         self.enable_prompt_tokens_details = enable_prompt_tokens_details

vllm/entrypoints/openai/chat_completion/serving.py

@@ -110,7 +110,9 @@
 from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs
 from vllm.sampling_params import BeamSearchParams, SamplingParams
 from vllm.tokenizers import TokenizerLike
+from vllm.tokenizers.mistral import MistralTokenizer
 from vllm.tool_parsers import ToolParser
+from vllm.tool_parsers.mistral_tool_parser import MistralToolParser
 from vllm.tracing import (
     contains_trace_headers,
     extract_trace_headers,
@@ -1228,24 +1230,33 @@ def _parse_tool_calls_from_content(
         tool_parser_cls: Callable[[TokenizerLike], ToolParser] | None,
         content: str | None = None,
     ) -> tuple[list[FunctionCall] | None, str | None]:
+        is_mistral_flow = tool_parser_cls == MistralToolParser and isinstance(
+            tokenizer, MistralTokenizer
+        )
         function_calls = list[FunctionCall]()
-        if request.tool_choice and isinstance(request.tool_choice, ToolChoiceFunction):
+        if (
+            request.tool_choice
+            and isinstance(request.tool_choice, ToolChoiceFunction)
+            and not is_mistral_flow
+        ):
             assert content is not None
             # Forced Function Call
             function_calls.append(
                 FunctionCall(name=request.tool_choice.name, arguments=content)
             )
             content = None  # Clear content since tool is called.
-        elif request.tool_choice and isinstance(
-            request.tool_choice, ChatCompletionNamedToolChoiceParam
+        elif (
+            request.tool_choice
+            and isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
+            and not is_mistral_flow
         ):
             assert content is not None
             # Forced Function Call
             function_calls.append(
                 FunctionCall(name=request.tool_choice.function.name, arguments=content)
             )
             content = None  # Clear content since tool is called.
-        elif request.tool_choice == "required":
+        elif request.tool_choice == "required" and not is_mistral_flow:
             assert content is not None
             tool_calls = TypeAdapter(list[FunctionDefinition]).validate_json(content)
             function_calls.extend(
@@ -1259,7 +1270,8 @@ def _parse_tool_calls_from_content(
             )
             content = None  # Clear content since tool is called.
         elif (
-            tool_parser_cls
+            is_mistral_flow
+            or tool_parser_cls
             and enable_auto_tools
             and (request.tool_choice == "auto" or request.tool_choice is None)
         ):
@@ -1270,6 +1282,7 @@ def _parse_tool_calls_from_content(
 
             # Automatic Tool Call Parsing
             try:
+                assert tool_parser_cls is not None
                 tool_parser = tool_parser_cls(tokenizer)
             except RuntimeError as e:
                 logger.exception("Error in tool parser creation.")

vllm/entrypoints/openai/engine/serving.py

@@ -37,6 +37,7 @@ class StructuredOutputsParams:
     regex: str | None = None
     choice: list[str] | None = None
     grammar: str | None = None
+    lark: str | None = None
     json_object: bool | None = None
     # These are other options that can be set.
     disable_fallback: bool = False
@@ -58,6 +59,7 @@ def __post_init__(self):
                 self.regex is not None,
                 self.choice is not None,
                 self.grammar is not None,
+                self.lark is not None,
                 self.json_object is not None,
                 self.structural_tag is not None,
             ]
@@ -84,6 +86,7 @@ def all_constraints_none(self) -> bool:
                 "regex",
                 "choice",
                 "grammar",
+                "lark",
                 "json_object",
                 "structural_tag",
             )
@@ -100,6 +103,7 @@ def all_non_structural_tag_constraints_none(self) -> bool:
                 "regex",
                 "choice",
                 "grammar",
+                "lark",
                 "json_object",
             )
         )

vllm/sampling_params.py

@@ -3,6 +3,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast, overload
 
+import llguidance as llg
+import regex as re
 from mistral_common.protocol.instruct.request import (
     ChatCompletionRequest as MistralChatCompletionRequest,
 )
@@ -11,8 +13,15 @@
 from mistral_common.tokens.tokenizers.base import (
     SpecialTokenPolicy,
     SpecialTokens,
+    Tokenizer,
+)
+from mistral_common.tokens.tokenizers.instruct import (
+    InstructTokenizerBase,
+    InstructTokenizerV13,
+)
+from mistral_common.tokens.tokenizers.mistral import (
+    MistralTokenizer as MistralCommonTokenizer,
 )
-from mistral_common.tokens.tokenizers.instruct import InstructTokenizerV13
 from mistral_common.tokens.tokenizers.sentencepiece import (
     SentencePieceTokenizer,
 )
@@ -21,20 +30,21 @@
 from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.logger import init_logger
+from vllm.tokenizers.protocol import TokenizerLike
+
+try:
+    # Transformers v5
+    from transformers.tokenization_mistral_common import MistralCommonBackend
+except ImportError:
+    # Transformers v4
+    from transformers.tokenization_mistral_common import (
+        MistralCommonTokenizer as MistralCommonBackend,
+    )
 
-from .protocol import TokenizerLike
 
 if TYPE_CHECKING:
     from transformers import BatchEncoding
 
-    try:
-        # Transformers v5
-        from transformers.tokenization_mistral_common import MistralCommonBackend
-    except ImportError:
-        # Transformers v4
-        from transformers.tokenization_mistral_common import (
-            MistralCommonTokenizer as MistralCommonBackend,
-        )
 
 logger = init_logger(__name__)
 
@@ -217,15 +227,6 @@ def from_pretrained(
         download_dir: str | None = None,
         **kwargs,
     ) -> "MistralTokenizer":
-        try:
-            # Transformers v5
-            from transformers.tokenization_mistral_common import MistralCommonBackend
-        except ImportError:
-            # Transformers v4
-            from transformers.tokenization_mistral_common import (
-                MistralCommonTokenizer as MistralCommonBackend,
-            )
-
         tokenizer = MistralCommonBackend.from_pretrained(
             path_or_repo_id,
             *args,
@@ -240,10 +241,10 @@ def from_pretrained(
     def __init__(self, tokenizer: "MistralCommonBackend") -> None:
         super().__init__()
 
-        self.transformers_tokenizer = tokenizer
-        self.mistral = tokenizer.tokenizer
-        self.instruct = self.mistral.instruct_tokenizer
-        self.tokenizer = self.instruct.tokenizer
+        self.transformers_tokenizer: MistralCommonBackend = tokenizer
+        self.mistral: MistralCommonTokenizer = tokenizer.tokenizer
+        self.instruct: InstructTokenizerBase = self.mistral.instruct_tokenizer
+        self.tokenizer: Tokenizer = self.instruct.tokenizer
 
         mode = self.mistral._chat_completion_request_validator._mode
         if mode != ValidationMode.test:
@@ -509,7 +510,7 @@ def convert_ids_to_tokens(
             return [self.tokenizer.id_to_piece(token_id) for token_id in ids]
 
         non_skip_special_tokens_ids = {
-            self.tokenizer.get_control_token(SpecialTokens.tool_calls),
+            self.tokenizer.get_special_token(SpecialTokens.tool_calls),
         }
         if isinstance(self.instruct, InstructTokenizerV13):
             if self.instruct.BEGIN_THINK:
@@ -541,3 +542,66 @@ def convert_ids_to_tokens(
             ]
 
         return tokens
+
+
+class MistralLLGTokenizer:
+    """Wraps a mistral tokenizer for use with llguidance."""
+
+    eos_token_id: int
+    bos_token_id: int
+    tokens: list[bytes]
+    special_token_ids: list[int]
+
+    def __init__(self, tokenizer: MistralTokenizer) -> None:
+        self._tokenizer = tokenizer.tokenizer
+        self.eos_token_id = self._tokenizer.eos_id
+        self.bos_token_id = self._tokenizer.bos_id
+
+        self.tokens: list[bytes] = []
+        self.special_token_ids: list[int] = []
+
+        seen_special_tokens: set[str] = set()
+        for i in range(self._tokenizer.n_words):
+            # Convert square brackets to angle brackets for special tokens,
+            # since llg only recognizes the latter.
+            if self._tokenizer.is_special(i):
+                token_rep = self._tokenizer.id_to_piece(i)
+                if match := re.fullmatch(r"\[(.*)\]", token_rep):
+                    token_rep_llg = f"<{match.group(1)}>"
+                else:
+                    token_rep_llg = token_rep
+
+                if not re.fullmatch(r"<.*>", token_rep_llg):
+                    raise ValueError(
+                        f"Invalid special token: {token_rep_llg} ({token_rep})"
+                    )
+                assert token_rep_llg not in seen_special_tokens, (
+                    token_rep_llg,
+                    seen_special_tokens,
+                )
+                seen_special_tokens.add(token_rep_llg)
+                self.special_token_ids.append(i)
+                self.tokens.append(token_rep_llg.encode("utf-8"))
+            else:
+                token = self._tokenizer.id_to_byte_piece(i, SpecialTokenPolicy.RAISE)
+                self.tokens.append(token)
+
+        assert len(self.special_token_ids) == self._tokenizer.num_special_tokens, (
+            len(self.special_token_ids),
+            self._tokenizer.num_special_tokens,
+        )
+
+    def __call__(self, s: str, *args, **kwds) -> list[int]:
+        # HACK: add a null byte to the start of the string to avoid the tokenizer
+        # absorbing the first character of tokens that start with "▁".
+        # we then ignore the first two tokens the "▁" and the null byte.
+        # This gives us the pure tokenized text without SP shit.
+        if isinstance(self._tokenizer, SentencePieceTokenizer):
+            return self._tokenizer.encode("\00" + s, bos=False, eos=False)[2:]
+        else:
+            return self._tokenizer.encode(s, bos=False, eos=False)
+
+
+def guidance_tokenizer_from_mistral_tokenizer(tokenizer: Tokenizer) -> llg.LLTokenizer:
+    tokenizer_data = MistralLLGTokenizer(tokenizer)
+    return llg.LLTokenizer(llg.TokenizerWrapper(tokenizer_data))