mistral.py 18.6 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
from typing import TYPE_CHECKING, Any, cast
5

6
from vllm.logger import init_logger
7
from vllm.transformers_utils.tokenizer_base import TokenizerBase
8

9
if TYPE_CHECKING:
10
11
12
13
14
15
    from mistral_common.protocol.instruct.request import (
        ChatCompletionRequest as MistralChatCompletionRequest,
    )
    from mistral_common.tokens.tokenizers.tekken import Tekkenizer
    from transformers.tokenization_mistral_common import (
        MistralCommonTokenizer as TransformersMistralTokenizer,
16
    )
17

18
    from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
19
    from vllm.entrypoints.openai.protocol import ChatCompletionRequest
20

21
22
logger = init_logger(__name__)

23

24
def maybe_serialize_tool_calls(request: "MistralChatCompletionRequest"):
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
    # SEE: https://github.com/vllm-project/vllm/pull/9951
    # Credits go to: @gcalmettes
    # NOTE: There is currently a bug in pydantic where attributes
    # declared as iterables are replaced in in the instances by
    # pydantic-core ValidatorIterator instance. In particular, this
    # affects tool_calls defined in ChatCompletionAssistantMessageParam
    # model:
    # see:
    #   - https://github.com/pydantic/pydantic/issues/9467
    # As a result, tool_calls from assistant messages are never
    # deserialized in the request object if the tool_calls iterator is
    # not consumed. This affect messages passed to the MistralTokenizer
    # since no chat template is applied and therefore the tools_calls
    # iterator is not directly consumed.
    # Issue is tracked on Pydantic side, with resolution planned for
    # v2.11 release. In the meantime, the official workaround is to
    # consume the iterator so the tool_calls are correctly deserialized
    # in the OpenAI ChatCompletionAssistantMessageParam object
    # https://github.com/pydantic/pydantic/issues/9467#issuecomment-2442097291 # noqa: E501
    # Official Pydantic Issues:
    #   - https://github.com/pydantic/pydantic/issues/9541
    # TODO: remove when pydantic v2.11 is released
    for i, message in enumerate(request.messages):
48
        if message.get("role") == "assistant":
49
50
51
52
53
54
55
56
57
58
59
60
            tool_calls_validator = message.get("tool_calls", ().__iter__())
            validated_tool_calls = []
            while True:
                try:
                    tool_call = next(tool_calls_validator)  # type: ignore
                    validated_tool_calls.append(tool_call)
                except StopIteration:
                    break

            request.messages[i]["tool_calls"] = validated_tool_calls


61
def truncate_tool_call_ids(request: "MistralChatCompletionRequest"):
62
63
    """Truncates tool call IDs for Mistral's ID requirements."""
    for i, message in enumerate(request.messages):
64
        if message.get("role") == "assistant":
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
            tool_calls = message.get("tool_calls", [])
            for tool_call in tool_calls:
                if len(tool_call["id"]) > 9:
                    logger.warning(
                        "Truncating tool call ID: %s to %s",
                        tool_call["id"],
                        tool_call["id"][-9:],
                    )
                    tool_call["id"] = tool_call["id"][-9:]

            request.messages[i]["tool_calls"] = tool_calls

        elif message.get("role") in {"tool_results", "tool"}:
            if "tool_call_id" in message:
                tool_call_id = message["tool_call_id"]

                if len(tool_call_id) > 9:
                    logger.warning(
                        "Truncating tool_call_id: %s to %s",
                        tool_call_id,
                        tool_call_id[-9:],
                    )
                    tool_call_id = tool_call_id[-9:]
                request.messages[i]["tool_call_id"] = tool_call_id


91
92
def _prepare_apply_chat_template_tools_and_messages(
    messages: list["ChatCompletionMessageParam"],
93
    tools: list[dict[str, Any]] | None = None,
94
95
    continue_final_message: bool = False,
    add_generation_prompt: bool = False,
96
) -> tuple[list["ChatCompletionMessageParam"], list[dict[str, Any]] | None]:
97
    if add_generation_prompt and continue_final_message:
98
        raise ValueError(
99
100
            "Cannot set both `add_generation_prompt` and "
            "`continue_final_message` to True."
101
        )
102

103
104
105
106
107
108
109
110
111
112
    last_message = cast(dict[str, Any], messages[-1])
    # add_generation_prompt is directly handled by the tokenizer but we
    # check if the user is trying to use it with a final assistant message
    # which is probably not what they want.
    # If add_generation_prompt is False, we don't need to check anything.
    if add_generation_prompt and last_message["role"] == "assistant":
        raise ValueError(
            "Cannot set `add_generation_prompt` to True when "
            "the last message is from the assistant. Consider "
            "using `continue_final_message` instead."
113
        )
114
115
116
117
    if continue_final_message and last_message["role"] != "assistant":
        raise ValueError(
            "Cannot set `continue_final_message` to True when "
            "the last message is not from the assistant."
118
        )
119

120
121
122
123
    # mistral-common requires AssistantMessage content to be string [1].
    #
    # [1]: https://github.com/mistralai/mistral-common/blob/f4a06998b75ed78bbf5aaf569590b772ea26c9f6/src/mistral_common/protocol/instruct/messages.py#L80
    for message in messages:
124
125
126
        # Remove reasoning_content as unsupported by Mistral
        _ = message.pop("reasoning_content", None)  # type: ignore

127
    # The Mistral client, in comparison to the OpenAI client, requires the
128
129
    # "parameters" dict and the "description" string to be present
    # even if they are empty.
130
131
    if tools:
        for function in [
132
            tool["function"] for tool in tools if tool["type"] == "function"
133
        ]:
134
135
            if function.get("parameters") is None:
                function["parameters"] = {}
136
137
            if function.get("description") is None:
                function["description"] = ""
138

139
    return messages, tools
140
141


142
143
144
def validate_request_params(request: "ChatCompletionRequest"):
    if request.chat_template is not None or request.chat_template_kwargs is not None:
        raise ValueError("chat_template is not supported for Mistral tokenizers.")
145
146


147
def _tekken_token_to_id(tokenizer: "Tekkenizer", t: str | bytes) -> int:
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
    from mistral_common.tokens.tokenizers.tekken import Tekkenizer

    assert isinstance(tokenizer, Tekkenizer), type(tokenizer)

    t_bytes = t.encode("utf-8") if not isinstance(t, bytes) else t
    shift = tokenizer.num_special_tokens
    try:
        return shift + tokenizer._tekken_token2id_nospecial[t_bytes]
    except KeyError:
        t_str = t_bytes.decode("utf-8")
        if t_str in tokenizer._special_tokens_reverse_vocab:
            return tokenizer._special_tokens_reverse_vocab[t_str]
        logger.warning(
            "Failed to convert token %s to id, replacing with <unk>", t_bytes
        )
        return tokenizer.unk_id

165

166
167
class MistralTokenizer(TokenizerBase):
    def __init__(self, tokenizer: "TransformersMistralTokenizer") -> None:
168
        from mistral_common.tokens.tokenizers.sentencepiece import (
169
170
            SentencePieceTokenizer,
        )
171
        from mistral_common.tokens.tokenizers.tekken import Tekkenizer
172

173
174
175
176
177
178
179
180
181
182
        self.transformers_tokenizer = tokenizer
        self.mistral = tokenizer.tokenizer
        self.instruct = self.mistral.instruct_tokenizer
        self.tokenizer = self.instruct.tokenizer

        _mistral_version_str = str(self.tokenizer.version.value)
        self.version: int = int(_mistral_version_str.split("v")[-1])

        self.is_tekken = isinstance(self.tokenizer, Tekkenizer)
        self.is_spm = isinstance(self.tokenizer, SentencePieceTokenizer)
183
        if not (self.is_tekken or self.is_spm):
184
185
186
187
188
189
190
191
192
193
194
195
            raise TypeError(f"Unsupported tokenizer: {type(self.tokenizer)}")

        # Reverse order to ensure that the lowest token id is kept.
        self._vocab_dict = {
            self.convert_ids_to_tokens([i], skip_special_tokens=False)[0]: i
            for i in range(self.vocab_size - 1, -1, -1)
        }
        # Sort the dict for convenience
        self._vocab_dict = dict(sorted(self._vocab_dict.items(), key=lambda x: x[1]))

        # Vocab sorted by token id.
        self._vocab = self.tokenizer._vocab
196
        self._max_token_id = self.vocab_size - 1
197
198

    @classmethod
199
    def from_pretrained(
200
        cls, path_or_repo_id: str, *, revision: str | None = None
201
    ) -> "MistralTokenizer":
202
203
        from transformers.tokenization_mistral_common import (
            MistralCommonTokenizer as TransformersMistralTokenizer,
204
205
        )

206
207
208
209
210
        str_revision = "main" if revision is None else revision
        return cls(
            TransformersMistralTokenizer.from_pretrained(
                path_or_repo_id, revision=str_revision
            )
211
        )
212

213
    # the following attributes are set to fit vLLM's design and are used
214
    # by the structured output backends.
215
    @property
216
    def all_special_tokens_extended(self) -> list[str]:
217
        return self.all_special_tokens
218
219

    @property
220
    def all_special_tokens(self) -> list[str]:
221
222
223
224
225
226
        from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy

        return [
            self.tokenizer.decode([i], special_token_policy=SpecialTokenPolicy.KEEP)
            for i in self.all_special_ids
        ]
227
228

    @property
229
    def all_special_ids(self) -> list[int]:
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
        from mistral_common.tokens.tokenizers.sentencepiece import (
            SentencePieceTokenizer,
        )
        from mistral_common.tokens.tokenizers.tekken import Tekkenizer

        if self.is_tekken:
            assert isinstance(self.tokenizer, Tekkenizer), type(self.tokenizer)
            special_ids = {t["rank"] for t in self.tokenizer._all_special_tokens}
        elif self.is_spm:
            assert isinstance(self.tokenizer, SentencePieceTokenizer), type(
                self.tokenizer
            )
            special_ids = self.tokenizer._control_tokens
        else:
            raise ValueError(f"Unknown tokenizer type: {type(self.tokenizer)}")
        return sorted(special_ids)
246
247
248
249
250
251
252
253
254

    @property
    def bos_token_id(self) -> int:
        return self.tokenizer.bos_id

    @property
    def eos_token_id(self) -> int:
        return self.tokenizer.eos_id

255
256
257
258
259
260
    @property
    def sep_token(self) -> str:
        raise NotImplementedError()

    @property
    def pad_token(self) -> str:
261
        return self.transformers_tokenizer.pad_token
262

263
264
265
266
267
268
    @property
    def is_fast(self) -> bool:
        return True

    @property
    def vocab_size(self) -> int:
269
        return self.transformers_tokenizer.vocab_size
270

271
272
273
274
    @property
    def max_token_id(self) -> int:
        return self._max_token_id

275
276
277
278
    @property
    def truncation_side(self) -> str:
        raise NotImplementedError()

279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
    def _is_special_token_id(self, token_id: int) -> bool:
        from mistral_common.tokens.tokenizers.sentencepiece import (
            SentencePieceTokenizer,
        )
        from mistral_common.tokens.tokenizers.tekken import Tekkenizer

        if self.is_spm:
            assert isinstance(self.tokenizer, SentencePieceTokenizer), type(
                self.tokenizer
            )
            return token_id in self.tokenizer._control_tokens
        if self.is_tekken:
            assert isinstance(self.tokenizer, Tekkenizer), type(self.tokenizer)
            return token_id < self.tokenizer.num_special_tokens
        else:
            raise ValueError(f"Unknown tokenizer type: {type(self.tokenizer)}")

296
297
298
    def __len__(self) -> int:
        return self.vocab_size

299
300
    def __call__(
        self,
301
302
        text: str | list[str] | list[int],
        text_pair: str | None = None,
303
304
        add_special_tokens: bool = False,
        truncation: bool = False,
305
        max_length: int | None = None,
306
    ):
307
308
309
310
311
312
313
314
315
316
317
        return self.transformers_tokenizer(
            text=text,
            text_pair=text_pair,
            add_special_tokens=add_special_tokens,
            truncation=truncation,
            max_length=max_length,
        )

    @property
    def vocab(self) -> list[str]:
        return self._vocab
318

319
    def get_vocab(self) -> dict[str, int]:
320
        return self._vocab_dict
321

322
    def get_added_vocab(self) -> dict[str, int]:
323
        # Mistral tokenizers have no added vocabulary
324
        return {}
325

326
327
    def encode_one(
        self,
328
        text: str,
329
        truncation: bool = False,
330
        max_length: int | None = None,
331
    ) -> list[int]:
332
        # Mistral Tokenizers should not add special tokens
333
334
335
        return self.transformers_tokenizer.encode(
            text, add_special_tokens=False, truncation=truncation, max_length=max_length
        )
336

337
338
339
    def encode(
        self,
        text: str,
340
341
342
        truncation: bool | None = None,
        max_length: int | None = None,
        add_special_tokens: bool | None = None,
343
    ) -> list[int]:
344
        if add_special_tokens is not None:
345
346
347
348
349
            return self.transformers_tokenizer.encode(
                text,
                truncation=truncation,
                max_length=max_length,
                add_special_tokens=add_special_tokens,
350
            )
351
        else:
352
353
354
355
356
357
            encoded = self.tokenizer.encode(text, bos=True, eos=False)

            if truncation is not False and max_length is not None:
                return encoded[:max_length]
            else:
                return encoded
358

359
360
361
    def apply_chat_template(
        self,
        messages: list["ChatCompletionMessageParam"],
362
        tools: list[dict[str, Any]] | None = None,
363
364
        **kwargs,
    ) -> list[int]:
365
366
367
368
369
370
371
372
373
        add_generation_prompt = kwargs.pop("add_generation_prompt", False)
        continue_final_message = kwargs.get("continue_final_message", False)
        padding = kwargs.get("padding", False)
        truncation = kwargs.get("truncation", False)
        max_length = kwargs.get("max_length")

        messages, tools = _prepare_apply_chat_template_tools_and_messages(
            messages, tools, continue_final_message, add_generation_prompt
        )
374

375
376
377
378
379
380
381
382
383
384
385
386
        return self.transformers_tokenizer.apply_chat_template(
            conversation=messages,
            tools=tools,
            continue_final_message=continue_final_message,
            tokenize=True,
            padding=padding,
            truncation=truncation,
            max_length=max_length,
            return_tensors=None,
            return_dict=False,
        )

387
    def decode(self, ids: list[int] | int, skip_special_tokens: bool = True) -> str:
388
389
390
        return self.transformers_tokenizer.decode(
            ids, skip_special_tokens=skip_special_tokens
        )
391

392
    def convert_tokens_to_string(self, tokens: list[str]) -> str:
393
394
395
396
397
398
399
400
        from mistral_common.tokens.tokenizers.base import (
            SpecialTokenPolicy,
            SpecialTokens,
        )
        from mistral_common.tokens.tokenizers.sentencepiece import (
            SentencePieceTokenizer,
        )
        from mistral_common.tokens.tokenizers.tekken import Tekkenizer
401

402
        to_decode_special_tokens = {SpecialTokens.tool_calls}
403
        if self.is_tekken:
404
            assert isinstance(self.tokenizer, Tekkenizer), type(self.tokenizer)
405
            tokens = [
406
407
                t
                for t in tokens
408
                if (t in to_decode_special_tokens or t not in self.all_special_tokens)
409
410
411
412
            ]

            if any(isinstance(t, bytes) for t in tokens):
                # we need to encode and decode all tokens again
413
414
415
416
                ids = [_tekken_token_to_id(self.tokenizer, t) for t in tokens]
                # We filtered unwanted special tokens before
                # so we can decode the rest.
                decoded = self.tokenizer.decode(ids, SpecialTokenPolicy.KEEP)
417
418
            else:
                decoded = "".join(tokens)
419
        else:
420
421
            # make sure certain special tokens like Tool calls are
            # not decoded
422
423
424
425
            assert isinstance(self.tokenizer, SentencePieceTokenizer), type(
                self.tokenizer
            )

426
            regular_tokens: list[str] = []
427
428
            decoded_list: list[str] = []
            decoded = ""
429
430

            for token in tokens:
431
                if token in to_decode_special_tokens:
432
433
                    if regular_tokens:
                        decoded_list.append(
434
                            self.tokenizer.decode(
435
                                regular_tokens, SpecialTokenPolicy.IGNORE
436
437
                            )
                        )
438
439
440
441
442
443
444
                        regular_tokens = []
                    decoded_list.append(token)
                else:
                    regular_tokens.append(token)

            if regular_tokens:
                decoded_list.append(
445
                    self.tokenizer.decode(regular_tokens, SpecialTokenPolicy.IGNORE)
446
447
                )
            decoded = "".join(decoded_list)
448
449

        return decoded
450
451

    def convert_ids_to_tokens(
452
        self,
453
        ids: list[int],
454
        skip_special_tokens: bool = True,
455
    ) -> list[str]:
456
457
458
459
        from mistral_common.tokens.tokenizers.base import (
            SpecialTokenPolicy,
            SpecialTokens,
        )
460
        from mistral_common.tokens.tokenizers.instruct import InstructTokenizerV13
461

462
463
        if not skip_special_tokens:
            return [self.tokenizer.id_to_piece(token_id) for token_id in ids]
464

465
466
467
468
469
470
471
472
        non_skip_special_tokens_ids = {
            self.tokenizer.get_control_token(SpecialTokens.tool_calls),
        }
        if isinstance(self.instruct, InstructTokenizerV13):
            if self.instruct.BEGIN_THINK:
                non_skip_special_tokens_ids.add(self.instruct.BEGIN_THINK)
            if self.instruct.END_THINK:
                non_skip_special_tokens_ids.add(self.instruct.END_THINK)
473

474
475
476
477
478
        ids_kept = [
            i
            for i in ids
            if i in non_skip_special_tokens_ids or not self._is_special_token_id(i)
        ]
479

480
481
        # We filtered unwanted special tokens so we can decode the rest.
        tokens = [self.tokenizer.id_to_piece(token_id) for token_id in ids_kept]
482

483
        if any("�" in t for t in tokens) and self.is_tekken:
484
485
            # if a decoded token contains the replacement character, then the
            # token has an incomplete UTF-8 character so we must use bytes
486
            # See: https://github.com/vllm-project/vllm/pull/8640
487
            #      https://github.com/vllm-project/vllm/pull/9625
488
489
            # if underlying tokenizer is sentencepiece, we just add "�".
            # We filtered unwanted special tokens so we can decode the rest.
490
            tokens = [
491
492
493
494
                self.tokenizer.id_to_byte_piece(token_id, SpecialTokenPolicy.KEEP)
                if token_id not in self.all_special_ids
                else self.tokenizer.decode([token_id], SpecialTokenPolicy.KEEP)
                for token_id in ids_kept
495
            ]
496

497
        return tokens