| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import copy |
| 4 | from typing import Any |
| 5 | |
| 6 | from transformers import TokenizersBackend |
| 7 | |
| 8 | from vllm.entrypoints.chat_utils import ChatCompletionMessageParam |
| 9 | |
| 10 | from .deepseek_v32_encoding import encode_messages |
| 11 | from .hf import HfTokenizer, get_cached_tokenizer |
| 12 | from .protocol import TokenizerLike |
| 13 | |
| 14 | |
| 15 | def get_deepseek_v32_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: |
| 16 | """Wraps a tokenizer to use the custom DeepSeek V3.2 chat template encoding.""" |
| 17 | dsv32_tokenizer = copy.copy(tokenizer) |
| 18 | |
| 19 | added_vocab = tokenizer.get_added_vocab() |
| 20 | |
| 21 | class _DeepseekV32Tokenizer(tokenizer.__class__): # type: ignore |
| 22 | def apply_chat_template( |
| 23 | self, |
| 24 | messages: list["ChatCompletionMessageParam"], |
| 25 | tools: list[dict[str, Any]] | None = None, |
| 26 | **kwargs, |
| 27 | ) -> str | list[int]: |
| 28 | thinking = kwargs.get("thinking", False) |
| 29 | enable_thinking = kwargs.get("enable_thinking", False) |
| 30 | thinking = thinking or enable_thinking |
| 31 | thinking_mode = "thinking" |
| 32 | if not thinking: |
| 33 | thinking_mode = "chat" |
| 34 | conversation = kwargs.get("conversation", messages) |
| 35 | messages = conversation.copy() |
| 36 | if tools is not None and len(tools) > 0: |
| 37 | messages.insert(0, {"role": "system"}) |
| 38 | messages[0]["tools"] = tools # type: ignore[typeddict-unknown-key] |
| 39 | |
| 40 | # Historical reasoning content is dropped when a new user message |
| 41 | # is introduced |
| 42 | drop_thinking = messages[-1]["role"] == "user" |
| 43 | |
| 44 | encode_config = dict( |
| 45 | thinking_mode=thinking_mode, drop_thinking=drop_thinking |
| 46 | ) |
| 47 | |
| 48 | prompt_str = encode_messages(messages, **encode_config) # type: ignore |
| 49 | |
| 50 | if kwargs.get("tokenize", True): |
| 51 | tokenizer_kwargs = { |
| 52 | k: kwargs[k] for k in ("truncation", "max_length") if k in kwargs |
| 53 | } |
| 54 | return self.encode( |
| 55 | prompt_str, |
| 56 | add_special_tokens=False, |
| 57 | **tokenizer_kwargs, |
| 58 | ) |
| 59 | |
| 60 | return prompt_str |
| 61 | |
| 62 | def num_special_tokens_to_add(self) -> int: |
| 63 | return len(self.encode("")) |
| 64 | |
| 65 | def get_added_vocab(self) -> dict[str, int]: |
| 66 | return added_vocab.copy() |
| 67 | |
| 68 | def __reduce__(self): |
| 69 | return get_deepseek_v32_tokenizer, (tokenizer,) |
| 70 | |
| 71 | _DeepseekV32Tokenizer.__name__ = f"DSV32{tokenizer.__class__.__name__}" |
| 72 | |
| 73 | dsv32_tokenizer.__class__ = _DeepseekV32Tokenizer |
| 74 | return dsv32_tokenizer |
| 75 | |
| 76 | |
| 77 | class DeepseekV32Tokenizer(TokenizerLike): |
| 78 | @classmethod |
| 79 | def from_pretrained(cls, *args, **kwargs) -> HfTokenizer: |
| 80 | tokenizer = TokenizersBackend.from_pretrained(*args, **kwargs) |
| 81 | return get_cached_tokenizer(get_deepseek_v32_tokenizer(tokenizer)) |