[Model] Add Mistral Tokenization to improve robustness and chat encoding

requirements-common.txt

@@ -26,3 +26,4 @@ librosa # Required for audio processing
 soundfile # Required for audio processing
 gguf == 0.9.1
 importlib_metadata
+mistral_common >= 1.3.4

vllm/config.py

@@ -238,10 +238,10 @@ def _init_multimodal_config(
 
     def _verify_tokenizer_mode(self) -> None:
         tokenizer_mode = self.tokenizer_mode.lower()
-        if tokenizer_mode not in ["auto", "slow"]:
+        if tokenizer_mode not in ["auto", "slow", "mistral"]:
             raise ValueError(
                 f"Unknown tokenizer mode: {self.tokenizer_mode}. Must be "
-                "either 'auto' or 'slow'.")
+                "either 'auto', 'slow' or 'mistral'.")
         self.tokenizer_mode = tokenizer_mode
 
     def _verify_embedding_mode(self) -> None:

vllm/engine/arg_utils.py

@@ -197,7 +197,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             '--tokenizer-mode',
             type=str,
             default=EngineArgs.tokenizer_mode,
-            choices=['auto', 'slow'],
+            choices=['auto', 'slow', 'mistral'],
             help='The tokenizer mode.\n\n* "auto" will use the '
             'fast tokenizer if available.\n* "slow" will '
             'always use the slow tokenizer.')

vllm/entrypoints/chat_utils.py

@@ -2,8 +2,16 @@
 from dataclasses import dataclass
 from functools import lru_cache
 from pathlib import Path
-from typing import (Any, Awaitable, Iterable, List, Literal, Optional, Tuple,
-                    Union)
+from typing import (
+    Any,
+    Awaitable,
+    Iterable,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+)
 
 # yapf conflicts with isort for this block
 # yapf: disable
@@ -21,9 +29,11 @@
 from vllm.config import ModelConfig
 from vllm.logger import init_logger
 from vllm.multimodal import MultiModalDataDict
-from vllm.multimodal.utils import (async_get_and_parse_audio,
-                                   async_get_and_parse_image)
-from vllm.transformers_utils.tokenizer import AnyTokenizer
+from vllm.multimodal.utils import (
+    async_get_and_parse_audio,
+    async_get_and_parse_image,
+)
+from vllm.transformers_utils.tokenizer import AnyTokenizer, AnyHFTokenizer
 
 logger = init_logger(__name__)
 
@@ -56,6 +66,7 @@ class CustomChatCompletionContentPartParam(TypedDict, total=False):
 
 class CustomChatCompletionMessageParam(TypedDict, total=False):
     """Enables custom roles in the Chat Completion API."""
+
     role: Required[str]
     """The role of the message's author."""
 
@@ -87,7 +98,7 @@ class ChatMessageParseResult:
 
 
 def load_chat_template(
-        chat_template: Optional[Union[Path, str]]) -> Optional[str]:
+    chat_template: Optional[Union[Path, str]]) -> Optional[str]:
     if chat_template is None:
         return None
     try:
@@ -113,8 +124,11 @@ def load_chat_template(
 
 
 @lru_cache(maxsize=None)
-def _mm_token_str(model_config: ModelConfig, tokenizer: AnyTokenizer,
-                  modality: Literal["image", "audio"]) -> Optional[str]:
+def _mm_token_str(
+    model_config: ModelConfig,
+    tokenizer: AnyHFTokenizer,
+    modality: Literal["image", "audio"],
+) -> Optional[str]:
     # TODO: Let user specify how to insert image tokens into prompt
     # (similar to chat template)
     if modality == "image":
@@ -259,7 +273,7 @@ def parse_chat_messages(
 
 
 def apply_chat_template(
-    tokenizer: AnyTokenizer,
+    tokenizer: AnyHFTokenizer,
     conversation: List[ConversationMessage],
     chat_template: Optional[str],
     *,

vllm/entrypoints/llm.py

@@ -5,22 +5,29 @@
 
 from vllm.engine.arg_utils import EngineArgs
 from vllm.engine.llm_engine import LLMEngine
-from vllm.entrypoints.chat_utils import (ChatCompletionMessageParam,
-                                         apply_chat_template,
-                                         parse_chat_messages)
+from vllm.entrypoints.chat_utils import (
+    ChatCompletionMessageParam,
+    apply_chat_template,
+    parse_chat_messages,
+)
 from vllm.inputs import PromptInputs, TextPrompt, TokensPrompt
 from vllm.inputs.parse import parse_and_batch_prompt
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.model_executor.guided_decoding import (
-    GuidedDecodingRequest, get_local_guided_decoding_logits_processor)
+    GuidedDecodingRequest,
+    get_local_guided_decoding_logits_processor,
+)
 from vllm.model_executor.guided_decoding.guided_fields import LLMGuidedOptions
 from vllm.outputs import EmbeddingRequestOutput, RequestOutput
 from vllm.pooling_params import PoolingParams
 from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import SamplingParams
-from vllm.transformers_utils.tokenizer import (AnyTokenizer,
-                                               get_cached_tokenizer)
+from vllm.transformers_utils.tokenizer import (
+    AnyTokenizer,
+    get_cached_tokenizer,
+    MistralTokenizer,
+)
 from vllm.transformers_utils.tokenizer_group import TokenizerGroup
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils import Counter, deprecate_kwargs
@@ -131,14 +138,14 @@ def __init__(
         disable_custom_all_reduce: bool = False,
         **kwargs,
     ) -> None:
-        '''
+        """
         LLM constructor.
 
         Note: if enforce_eager is unset (enforce_eager is None)
         it defaults to False for decoder-only models and True
         for encoder/decoder models, since encoder/decoder models
         do not currently support CUDAGraph.
-        '''
+        """
 
         if "disable_log_stats" not in kwargs:
             kwargs["disable_log_stats"] = True
@@ -273,15 +280,15 @@ def generate(
     def generate(
         self,
         prompts: Union[Union[PromptInputs, Sequence[PromptInputs]],
-                       Optional[Union[str, List[str]]]] = None,
+                       Optional[Union[str, List[str]]], ] = None,
         sampling_params: Optional[Union[SamplingParams,
                                         Sequence[SamplingParams]]] = None,
         prompt_token_ids: Optional[Union[List[int], List[List[int]]]] = None,
         use_tqdm: bool = True,
         lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
         prompt_adapter_request: Optional[PromptAdapterRequest] = None,
         guided_options_request: Optional[Union[LLMGuidedOptions,
-                                               GuidedDecodingRequest]] = None
+                                               GuidedDecodingRequest]] = None,
     ) -> List[RequestOutput]:
         """Generates the completions for the input prompts.
 
@@ -340,7 +347,8 @@ def generate(
             params=sampling_params,
             lora_request=lora_request,
             prompt_adapter_request=prompt_adapter_request,
-            guided_options=guided_options_request)
+            guided_options=guided_options_request,
+        )
 
         outputs = self._run_engine(use_tqdm=use_tqdm)
         return LLMEngine.validate_outputs(outputs, RequestOutput)
@@ -374,7 +382,7 @@ def chat(
             lora_request: LoRA request to use for generation, if any.
             chat_template: The template to use for structuring the chat.
               If not provided, the model's default chat template will be used.
-            add_generation_template: If True, adds a generation template 
+            add_generation_template: If True, adds a generation template
                 to each message.
 
         Returns:
@@ -383,20 +391,29 @@ def chat(
         """
 
         tokenizer = self.get_tokenizer()
-        model_config = self.llm_engine.get_model_config()
 
-        conversations, _ = parse_chat_messages(messages, model_config,
-                                               tokenizer)
+        if isinstance(tokenizer, MistralTokenizer):
+            prompt_token_ids = tokenizer.encode_messages(messages)
+            prompts = None
+        else:
+            model_config = self.llm_engine.get_model_config()
+
+            conversations, _ = parse_chat_messages(messages, model_config,
+                                                   tokenizer)
 
-        prompts = apply_chat_template(
-            tokenizer,
-            conversations,
-            chat_template=chat_template,
-            add_generation_template=add_generation_template)
+            prompts = apply_chat_template(
+                tokenizer,
+                conversations,
+                chat_template=chat_template,
+                add_generation_template=add_generation_template,
+            )
+
+            prompt_token_ids = None
 
         return self.generate(
-            prompts,
+            prompts,  # type: ignore[arg-type]
             sampling_params,
+            prompt_token_ids=prompt_token_ids,
             use_tqdm=use_tqdm,
             lora_request=lora_request,
         )
@@ -484,7 +501,7 @@ def encode(
     def encode(
         self,
         prompts: Union[Union[PromptInputs, Sequence[PromptInputs]],
-                       Optional[Union[str, List[str]]]] = None,
+                       Optional[Union[str, List[str]]], ] = None,
         pooling_params: Optional[Union[PoolingParams,
                                        Sequence[PoolingParams]]] = None,
         prompt_token_ids: Optional[Union[List[int], List[List[int]]]] = None,
@@ -564,8 +581,8 @@ def _convert_v1_inputs(
         if prompts is not None:
             num_requests = len(prompts)
         if prompt_token_ids is not None:
-            if (num_requests is not None
-                    and num_requests != len(prompt_token_ids)):
+            if num_requests is not None and num_requests != len(
+                    prompt_token_ids):
                 raise ValueError("The lengths of prompts and prompt_token_ids "
                                  "must be the same.")
 
@@ -593,7 +610,7 @@ def _validate_and_add_requests(
         self,
         inputs: Union[PromptInputs, Sequence[PromptInputs]],
         params: Union[SamplingParams, Sequence[SamplingParams], PoolingParams,
-                      Sequence[PoolingParams]],
+                      Sequence[PoolingParams], ],
         lora_request: Optional[Union[Sequence[LoRARequest], LoRARequest]],
         prompt_adapter_request: Optional[PromptAdapterRequest],
         guided_options: Optional[GuidedDecodingRequest] = None,
@@ -648,17 +665,21 @@ def _add_request(
         )
 
     def _add_guided_processor(
-            self,
-            params: SamplingParams,
-            guided_options: Optional[GuidedDecodingRequest] = None):
+        self,
+        params: SamplingParams,
+        guided_options: Optional[GuidedDecodingRequest] = None,
+    ):
         if guided_options:
             if guided_options.guided_decoding_backend is None:
                 decoding_config = self.llm_engine.get_decoding_config()
                 guided_options.guided_decoding_backend = (
                     decoding_config.guided_decoding_backend)
-            guided_logits_processor = get_local_guided_decoding_logits_processor(  #noqa
-                guided_options.guided_decoding_backend, guided_options,
-                self.get_tokenizer())
+            guided_logits_processor = (
+                get_local_guided_decoding_logits_processor(  # noqa
+                    guided_options.guided_decoding_backend,
+                    guided_options,
+                    self.get_tokenizer(),
+                ))
             if guided_logits_processor:
                 if params.logits_processors is None:
                     params.logits_processors = []