main
py 572 lines 21.3 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Sequence
4 from functools import cached_property
5 from pathlib import Path
6 from typing import TYPE_CHECKING, Any, cast, overload
7
8 from mistral_common.guidance.grammar_factory import GrammarFactory
9 from mistral_common.guidance.tokenizer import from_mistral_tokenizer
10 from mistral_common.protocol.instruct.request import (
11 ChatCompletionRequest as MistralChatCompletionRequest,
12 )
13 from mistral_common.protocol.instruct.request import (
14 ReasoningEffort,
15 )
16 from mistral_common.protocol.instruct.validator import ValidationMode
17 from mistral_common.tokens.tokenizers.base import (
18 SpecialTokenPolicy,
19 SpecialTokens,
20 Tokenizer,
21 )
22 from mistral_common.tokens.tokenizers.instruct import (
23 InstructTokenizerBase,
24 )
25 from mistral_common.tokens.tokenizers.mistral import (
26 MistralTokenizer as MistralCommonTokenizer,
27 )
28 from mistral_common.tokens.tokenizers.sentencepiece import (
29 SentencePieceTokenizer,
30 )
31 from mistral_common.tokens.tokenizers.tekken import Tekkenizer
32 from pydantic import ValidationError
33 from transformers.tokenization_mistral_common import MistralCommonBackend
34
35 from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
36 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
37 from vllm.exceptions import VLLMValidationError
38 from vllm.logger import init_logger
39 from vllm.tokenizers.protocol import TokenizerLike
40
41 if TYPE_CHECKING:
42 import llguidance
43 from transformers import BatchEncoding
44
45 logger = init_logger(__name__)
46
47
48 def maybe_serialize_tool_calls(request: "MistralChatCompletionRequest"):
49 # SEE: https://github.com/vllm-project/vllm/pull/9951
50 # Credits go to: @gcalmettes
51 # NOTE: There is currently a bug in pydantic where attributes
52 # declared as iterables are replaced in the instances by
53 # pydantic-core ValidatorIterator instance. In particular, this
54 # affects tool_calls defined in ChatCompletionAssistantMessageParam
55 # model:
56 # see:
57 # - https://github.com/pydantic/pydantic/issues/9467
58 # As a result, tool_calls from assistant messages are never
59 # deserialized in the request object if the tool_calls iterator is
60 # not consumed. This affect messages passed to the MistralTokenizer
61 # since no chat template is applied and therefore the tools_calls
62 # iterator is not directly consumed.
63 # Issue is tracked on Pydantic side, with resolution planned for
64 # v2.11 release. In the meantime, the official workaround is to
65 # consume the iterator so the tool_calls are correctly deserialized
66 # in the OpenAI ChatCompletionAssistantMessageParam object
67 # https://github.com/pydantic/pydantic/issues/9467#issuecomment-2442097291 # noqa: E501
68 # Official Pydantic Issues:
69 # - https://github.com/pydantic/pydantic/issues/9541
70 # TODO: remove when pydantic v2.11 is released
71 for i, message in enumerate(request.messages):
72 if message.get("role") == "assistant":
73 if (tool_calls_validator := message.get("tool_calls", None)) is not None:
74 try:
75 validated_tool_calls = list(tool_calls_validator)
76 except ValidationError as e:
77 raise ValueError(
78 "Validating messages' `tool_calls` raised an error. "
79 "Please ensure `tool_calls` are iterable of tool calls."
80 ) from e
81 else:
82 validated_tool_calls = []
83
84 request.messages[i]["tool_calls"] = validated_tool_calls
85
86
87 def truncate_tool_call_ids(request: "MistralChatCompletionRequest"):
88 """Truncates tool call IDs for Mistral's ID requirements."""
89 for i, message in enumerate(request.messages):
90 if message.get("role") == "assistant":
91 tool_calls = message.get("tool_calls", [])
92 for tool_call in tool_calls:
93 if len(tool_call["id"]) > 9:
94 logger.warning(
95 "Truncating tool call ID: %s to %s",
96 tool_call["id"],
97 tool_call["id"][-9:],
98 )
99 tool_call["id"] = tool_call["id"][-9:]
100
101 request.messages[i]["tool_calls"] = tool_calls
102
103 elif message.get("role") in {"tool_results", "tool"}:
104 if "tool_call_id" in message:
105 tool_call_id = message["tool_call_id"]
106
107 if len(tool_call_id) > 9:
108 logger.warning(
109 "Truncating tool_call_id: %s to %s",
110 tool_call_id,
111 tool_call_id[-9:],
112 )
113 tool_call_id = tool_call_id[-9:]
114 request.messages[i]["tool_call_id"] = tool_call_id
115
116
117 def _validate_apply_chat_template_args(
118 messages: list["ChatCompletionMessageParam"],
119 continue_final_message: bool = False,
120 add_generation_prompt: bool = False,
121 ) -> None:
122 if add_generation_prompt and continue_final_message:
123 raise ValueError(
124 "Cannot set both `add_generation_prompt` and "
125 "`continue_final_message` to True."
126 )
127
128 last_message = cast(dict[str, Any], messages[-1])
129 # add_generation_prompt is directly handled by the tokenizer but we
130 # check if the user is trying to use it with a final assistant message
131 # which is probably not what they want.
132 # If add_generation_prompt is False, we don't need to check anything.
133 if add_generation_prompt and last_message["role"] == "assistant":
134 raise ValueError(
135 "Cannot set `add_generation_prompt` to True when "
136 "the last message is from the assistant. Consider "
137 "using `continue_final_message` instead."
138 )
139 if continue_final_message and last_message["role"] != "assistant":
140 raise ValueError(
141 "Cannot set `continue_final_message` to True when "
142 "the last message is not from the assistant."
143 )
144
145
146 def validate_request_params(request: "ChatCompletionRequest"):
147 if request.chat_template is not None or request.chat_template_kwargs is not None:
148 raise VLLMValidationError(
149 "chat_template is not supported for Mistral tokenizers.",
150 parameter="chat_template",
151 )
152
153 if request.reasoning_effort and request.reasoning_effort not in list(
154 ReasoningEffort
155 ):
156 raise VLLMValidationError(
157 f"reasoning_effort={request.reasoning_effort} is not supported by "
158 "Mistral models. Supported values are: "
159 f"{[e.value for e in ReasoningEffort]}.",
160 parameter="reasoning_effort",
161 )
162
163
164 def _tekken_token_to_id(tokenizer: "Tekkenizer", t: str | bytes) -> int:
165 assert isinstance(tokenizer, Tekkenizer), type(tokenizer)
166
167 t_bytes = t.encode("utf-8") if not isinstance(t, bytes) else t
168 shift = tokenizer.num_special_tokens
169 try:
170 return shift + tokenizer._tekken_token2id_nospecial[t_bytes]
171 except KeyError:
172 t_str = t_bytes.decode("utf-8")
173 if t_str in tokenizer._special_tokens_reverse_vocab:
174 return tokenizer._special_tokens_reverse_vocab[t_str]
175 logger.warning(
176 "Failed to convert token %s to id, replacing with <unk>", t_bytes
177 )
178 return tokenizer.unk_id
179
180
181 def mistral_common_tekkenizer(tokenizer: object) -> "Tekkenizer | None":
182 """Return the underlying `Tekkenizer` for a `MistralCommonBackend`."""
183 mistral = getattr(tokenizer, "tokenizer", None)
184 instruct = getattr(mistral, "instruct_tokenizer", None)
185 tekken = getattr(instruct, "tokenizer", None)
186 return tekken if isinstance(tekken, Tekkenizer) else None
187
188
189 def tekken_convert_ids_to_tokens(
190 tokenizer: "Tekkenizer", ids: Sequence[int]
191 ) -> list[str | bytes]:
192 """Convert ids to pieces, using raw `bytes` for byte-fallback tokens."""
193 tokens: list[str | bytes] = [tokenizer.id_to_piece(i) for i in ids]
194 if any("�" in t for t in tokens):
195 tokens = [
196 tokenizer.id_to_byte_piece(i, SpecialTokenPolicy.KEEP)
197 if i >= tokenizer.num_special_tokens
198 else tokenizer.decode([i], SpecialTokenPolicy.KEEP)
199 for i in ids
200 ]
201 return tokens
202
203
204 def tekken_convert_tokens_to_string(
205 tokenizer: "Tekkenizer", tokens: Sequence[str | bytes]
206 ) -> str:
207 """Reassemble pieces from `tekken_convert_ids_to_tokens` into text."""
208 if any(isinstance(t, bytes) for t in tokens):
209 ids = [_tekken_token_to_id(tokenizer, t) for t in tokens]
210 return tokenizer.decode(ids, SpecialTokenPolicy.KEEP)
211 return "".join(cast(Sequence[str], tokens))
212
213
214 class MistralTokenizer(TokenizerLike):
215 IS_MISTRAL_TOKENIZER = True # used by vllm.utils.mistral
216
217 @classmethod
218 def from_pretrained(
219 cls,
220 path_or_repo_id: str | Path,
221 *args,
222 trust_remote_code: bool = False,
223 revision: str | None = None,
224 download_dir: str | None = None,
225 **kwargs,
226 ) -> "MistralTokenizer":
227 tokenizer = MistralCommonBackend.from_pretrained(
228 path_or_repo_id,
229 *args,
230 mode=ValidationMode.test,
231 cache_dir=download_dir,
232 revision="main" if revision is None else revision,
233 **kwargs,
234 )
235
236 return cls(tokenizer)
237
238 def __init__(self, tokenizer: MistralCommonBackend) -> None:
239 super().__init__()
240
241 self.transformers_tokenizer: MistralCommonBackend = tokenizer
242 self.mistral: MistralCommonTokenizer = tokenizer.tokenizer
243 self.instruct: InstructTokenizerBase = self.mistral.instruct_tokenizer
244 self.tokenizer: Tokenizer = self.instruct.tokenizer
245
246 mode = self.mistral._chat_completion_request_validator._mode
247 if mode != ValidationMode.test:
248 raise ValueError(
249 "Mistral tokenizer must be in test mode. Make sure to "
250 "set `mode='ValidationMode.test'` when creating the "
251 "Mistral tokenizer."
252 )
253
254 _mistral_version_str = str(self.tokenizer.version.value)
255 self.version: int = int(_mistral_version_str.split("v")[-1])
256
257 self.is_tekken = isinstance(self.tokenizer, Tekkenizer)
258 self.is_spm = isinstance(self.tokenizer, SentencePieceTokenizer)
259 if not (self.is_tekken or self.is_spm):
260 raise TypeError(f"Unsupported tokenizer: {type(self.tokenizer)}")
261
262 # Reverse order to ensure that the lowest token id is kept.
263 self._vocab_dict = {
264 self.convert_ids_to_tokens([i], skip_special_tokens=False)[0]: i
265 for i in range(self.vocab_size - 1, -1, -1)
266 }
267 # Sort the dict for convenience
268 self._vocab_dict = dict(sorted(self._vocab_dict.items(), key=lambda x: x[1]))
269
270 # Vocab sorted by token id.
271 self._vocab = self.tokenizer.vocab()
272 self._max_token_id = self.vocab_size - 1
273 self._max_chars_per_token = max(len(tok) for tok in self._vocab)
274
275 # Cache special tokens for faster access.
276 self._special_token_ids = self._get_special_token_ids()
277 self._special_token_ids_set = set(self._special_token_ids)
278 self._special_tokens = self._get_special_tokens(self._special_token_ids)
279 self._special_tokens_set = set(self._special_tokens)
280
281 def _get_special_token_ids(self) -> list[int]:
282 return [i for i in range(len(self._vocab)) if self.tokenizer.is_special(i)]
283
284 def _get_special_tokens(self, all_special_ids: list[int]) -> list[str]:
285 return [
286 self.tokenizer.decode([i], special_token_policy=SpecialTokenPolicy.KEEP)
287 for i in all_special_ids
288 ]
289
290 def num_special_tokens_to_add(self) -> int:
291 return len(self.encode(""))
292
293 # the following attributes are set to fit vLLM's design and are used
294 # by the structured output backends.
295 @property
296 def all_special_tokens(self) -> list[str]:
297 return self._special_tokens
298
299 @property
300 def all_special_ids(self) -> list[int]:
301 return self._special_token_ids
302
303 @property
304 def bos_token_id(self) -> int:
305 return self.tokenizer.bos_id
306
307 @property
308 def eos_token_id(self) -> int:
309 return self.tokenizer.eos_id
310
311 @property
312 def pad_token_id(self) -> int:
313 return self.tokenizer.pad_id
314
315 @property
316 def is_fast(self) -> bool:
317 return True
318
319 @property
320 def vocab_size(self) -> int:
321 return self.transformers_tokenizer.vocab_size
322
323 @property
324 def max_token_id(self) -> int:
325 return self._max_token_id
326
327 @property
328 def max_chars_per_token(self) -> int:
329 return self._max_chars_per_token
330
331 @property
332 def truncation_side(self) -> str:
333 return self.transformers_tokenizer.truncation_side
334
335 def _is_special_token_id(self, token_id: int) -> bool:
336 return token_id in self._special_token_ids_set
337
338 def __hash__(self) -> int:
339 return hash(id(self))
340
341 def __len__(self) -> int:
342 return self.vocab_size
343
344 def __call__(
345 self,
346 text: str | list[str],
347 text_pair: str | None = None,
348 add_special_tokens: bool = True,
349 truncation: bool = False,
350 max_length: int | None = None,
351 ) -> "BatchEncoding":
352 if text_pair is not None:
353 raise ValueError(
354 "`text_pair` is not supported by `MistralTokenizer.__call__`."
355 )
356
357 return self.transformers_tokenizer(
358 text=text,
359 text_pair=text_pair,
360 add_special_tokens=add_special_tokens,
361 truncation=truncation,
362 max_length=max_length,
363 )
364
365 @property
366 def vocab(self) -> list[str]:
367 return self._vocab
368
369 def get_vocab(self) -> dict[str, int]:
370 return self._vocab_dict
371
372 def get_added_vocab(self) -> dict[str, int]:
373 # Mistral tokenizers have no added vocabulary
374 return {}
375
376 def encode(
377 self,
378 text: str,
379 truncation: bool | None = None,
380 max_length: int | None = None,
381 add_special_tokens: bool = True,
382 ) -> list[int]:
383 # NOTE: transformers guards truncation with `if max_length and ...`, so
384 # `max_length=0` returns the full sequence instead of an empty one. vLLM
385 # treats `truncate_prompt_tokens=0` as an empty prompt, so keep slicing
386 # here until that is fixed upstream.
387 encoded = self.tokenizer.encode(text, bos=add_special_tokens, eos=False)
388
389 if truncation is not False and max_length is not None:
390 return encoded[:max_length]
391 else:
392 return encoded
393
394 def apply_chat_template(
395 self,
396 messages: list["ChatCompletionMessageParam"],
397 tools: list[dict[str, Any]] | None = None,
398 **kwargs,
399 ) -> list[int]:
400 add_generation_prompt = kwargs.pop("add_generation_prompt", False)
401 continue_final_message = kwargs.get("continue_final_message", False)
402 tokenize = kwargs.get("tokenize", True)
403 padding = kwargs.get("padding", False)
404 truncation = kwargs.get("truncation", False)
405 max_length = kwargs.get("max_length")
406
407 version_kwargs = {}
408 # NOTE: This is for backward compatibility.
409 # Transformers should be passed arguments it knows.
410 if self.version >= 15:
411 version_kwargs["reasoning_effort"] = kwargs.get("reasoning_effort")
412
413 _validate_apply_chat_template_args(
414 messages, continue_final_message, add_generation_prompt
415 )
416
417 return self.transformers_tokenizer.apply_chat_template(
418 conversation=messages,
419 tools=tools,
420 continue_final_message=continue_final_message,
421 tokenize=tokenize,
422 padding=padding,
423 truncation=truncation,
424 max_length=max_length,
425 return_tensors=None,
426 return_dict=False,
427 **version_kwargs,
428 )
429
430 def decode(
431 self, ids: Sequence[int] | int, skip_special_tokens: bool = False
432 ) -> str:
433 return self.transformers_tokenizer.decode(
434 ids, skip_special_tokens=skip_special_tokens
435 )
436
437 def batch_decode(
438 self, ids: list[list[int]] | list[int], skip_special_tokens: bool = False
439 ) -> str:
440 return self.transformers_tokenizer.batch_decode(
441 ids, skip_special_tokens=skip_special_tokens
442 )
443
444 @overload
445 def convert_tokens_to_ids(self, tokens: str) -> int: ...
446
447 @overload
448 def convert_tokens_to_ids(self, tokens: list[str]) -> list[int]: ...
449
450 def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]:
451 return self.transformers_tokenizer.convert_tokens_to_ids(tokens)
452
453 def convert_tokens_to_string(self, tokens: list[str]) -> str:
454 to_decode_special_tokens = {
455 SpecialTokens.tool_calls,
456 SpecialTokens.args,
457 SpecialTokens.begin_think,
458 SpecialTokens.end_think,
459 }
460 if self.is_tekken:
461 assert isinstance(self.tokenizer, Tekkenizer), type(self.tokenizer)
462 tokens = [
463 t
464 for t in tokens
465 if (t in to_decode_special_tokens or t not in self._special_tokens_set)
466 ]
467
468 # We filtered unwanted special tokens before so we can decode the rest.
469 decoded = tekken_convert_tokens_to_string(self.tokenizer, tokens)
470 else:
471 # make sure certain special tokens like Tool calls are
472 # not decoded
473 assert isinstance(self.tokenizer, SentencePieceTokenizer), type(
474 self.tokenizer
475 )
476
477 regular_tokens: list[str] = []
478 decoded_list: list[str] = []
479 decoded = ""
480
481 for token in tokens:
482 if token in to_decode_special_tokens:
483 if regular_tokens:
484 decoded_list.append(
485 self.tokenizer.decode(
486 regular_tokens, SpecialTokenPolicy.IGNORE
487 )
488 )
489 regular_tokens = []
490 decoded_list.append(token)
491 else:
492 regular_tokens.append(token)
493
494 if regular_tokens:
495 decoded_list.append(
496 self.tokenizer.decode(regular_tokens, SpecialTokenPolicy.IGNORE)
497 )
498 decoded = "".join(decoded_list)
499
500 return decoded
501
502 def convert_ids_to_tokens(
503 self,
504 ids: Sequence[int],
505 skip_special_tokens: bool = False,
506 ) -> list[str]:
507 if not skip_special_tokens:
508 return [self.tokenizer.id_to_piece(token_id) for token_id in ids]
509
510 non_skip_special_tokens_ids = {
511 self.tokenizer.get_special_token(SpecialTokens.tool_calls),
512 }
513 # [ARGS] only exists in v11+ tool-call tokenizers; older tokenizers
514 # raise (Tekken) or return unk (SPM) for it.
515 if self.tokenizer.is_special(SpecialTokens.args):
516 non_skip_special_tokens_ids.add(
517 self.tokenizer.get_special_token(SpecialTokens.args)
518 )
519 # [THINK]/[/THINK] only exist in v13+ reasoning tokenizers; use the
520 # same is_special gate as [ARGS] above so newer versions are covered
521 # without an isinstance check on the instruct tokenizer.
522 for think_token in (SpecialTokens.begin_think, SpecialTokens.end_think):
523 if self.tokenizer.is_special(think_token):
524 non_skip_special_tokens_ids.add(
525 self.tokenizer.get_special_token(think_token)
526 )
527
528 ids_kept = [
529 i
530 for i in ids
531 if i in non_skip_special_tokens_ids or not self._is_special_token_id(i)
532 ]
533
534 # We filtered unwanted special tokens so we can decode the rest.
535 tokens = [self.tokenizer.id_to_piece(token_id) for token_id in ids_kept]
536
537 if any("�" in t for t in tokens) and self.is_tekken:
538 # if a decoded token contains the replacement character, then the
539 # token has an incomplete UTF-8 character so we must use bytes
540 # See: https://github.com/vllm-project/vllm/pull/8640
541 # https://github.com/vllm-project/vllm/pull/9625
542 # if underlying tokenizer is sentencepiece, we just add "�".
543 # We filtered unwanted special tokens so we can decode the rest.
544 tokens = [
545 self.tokenizer.id_to_byte_piece(token_id, SpecialTokenPolicy.KEEP)
546 if token_id not in self._special_token_ids_set
547 else self.tokenizer.decode([token_id], SpecialTokenPolicy.KEEP)
548 for token_id in ids_kept
549 ]
550
551 return tokens
552
553 @property
554 def supports_grammar(self) -> bool:
555 return GrammarFactory.is_supported(self.mistral)
556
557 @cached_property
558 def grammar_factory(self) -> GrammarFactory:
559 if not self.supports_grammar:
560 raise AttributeError(
561 "This tokenizer does not support `grammar_factory`. "
562 "This is only supported for tekken tokenizers with "
563 "version >= 11."
564 )
565 # Cache grammar factory to avoid creating a llguidance tokenizer at every usage.
566 return GrammarFactory(self.mistral)
567
568 @cached_property
569 def llg_tokenizer(self) -> "llguidance.LLTokenizer":
570 if not self.is_tekken:
571 raise ValueError("`llg_tokenizer` is only supported for Tekkenizers.")
572 return from_mistral_tokenizer(self.mistral)