| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import copy |
| 5 | from typing import Any |
| 6 | |
| 7 | from transformers import TokenizersBackend |
| 8 | |
| 9 | from vllm.entrypoints.chat_utils import ChatCompletionMessageParam |
| 10 | |
| 11 | from .deepseek_v41_encoding import ( |
| 12 | IMAGE_PLACEHOLDER, |
| 13 | REASONING_EFFORT_MAPPINGS, |
| 14 | encode_messages, |
| 15 | ) |
| 16 | from .hf import HfTokenizer, get_cached_tokenizer |
| 17 | from .protocol import TokenizerLike |
| 18 | |
| 19 | |
| 20 | def _normalize_messages( |
| 21 | messages: list[ChatCompletionMessageParam], |
| 22 | ) -> list[dict[str, Any]]: |
| 23 | result = [dict(message) for message in copy.deepcopy(messages)] |
| 24 | for message in result: |
| 25 | role = message.get("role") |
| 26 | if role not in ("system", "developer", "user", "assistant", "tool"): |
| 27 | raise ValueError(f"Invalid role: {role}") |
| 28 | if "reasoning" in message: |
| 29 | message["reasoning_content"] = message["reasoning"] |
| 30 | content = message.get("content") |
| 31 | if isinstance(content, list): |
| 32 | parts = [] |
| 33 | for block in content: |
| 34 | part_type = block.get("type") |
| 35 | if part_type in ("text", "input_text", "output_text"): |
| 36 | parts.append(block.get("text", "")) |
| 37 | elif part_type in ("image_url", "input_image", "image_pil"): |
| 38 | parts.append(IMAGE_PLACEHOLDER) |
| 39 | else: |
| 40 | raise ValueError( |
| 41 | "DeepSeek V4.1 supports text and image content only; " |
| 42 | f"got {part_type!r}" |
| 43 | ) |
| 44 | message["content"] = "\n\n".join(parts) |
| 45 | return result |
| 46 | |
| 47 | |
| 48 | def get_deepseek_v41_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: |
| 49 | """Wrap an HF tokenizer with the V4.1 prompt encoder.""" |
| 50 | wrapped = copy.copy(tokenizer) |
| 51 | added_vocab = tokenizer.get_added_vocab() |
| 52 | |
| 53 | class _DeepseekV41Tokenizer(tokenizer.__class__): # type: ignore |
| 54 | def apply_chat_template( |
| 55 | self, |
| 56 | messages: list[ChatCompletionMessageParam], |
| 57 | tools: list[dict[str, Any]] | None = None, |
| 58 | **kwargs, |
| 59 | ) -> str | list[int]: |
| 60 | # The generic renderer's conversation has already flattened text |
| 61 | # parts with '\n'. V4.1 encodes the original messages with '\n\n'. |
| 62 | conversation = _normalize_messages(messages) |
| 63 | if tools: |
| 64 | system = next((m for m in conversation if m["role"] == "system"), None) |
| 65 | if system is None: |
| 66 | system = {"role": "system", "content": ""} |
| 67 | conversation.insert(0, system) |
| 68 | system["tools"] = tools |
| 69 | |
| 70 | thinking = bool(kwargs.get("thinking") or kwargs.get("enable_thinking")) |
| 71 | if "thinking" not in kwargs and "enable_thinking" not in kwargs: |
| 72 | thinking = True |
| 73 | effort = kwargs.get("reasoning_effort") |
| 74 | if effort == "none": |
| 75 | thinking = False |
| 76 | effort = None |
| 77 | if effort is None: |
| 78 | effort = "high" |
| 79 | if not ( |
| 80 | (type(effort) is int and 1 <= effort <= 100) |
| 81 | or (isinstance(effort, str) and effort in REASONING_EFFORT_MAPPINGS) |
| 82 | ): |
| 83 | raise ValueError( |
| 84 | "DeepSeek V4.1 reasoning_effort must be low, high, xhigh, max, " |
| 85 | "or an integer within [1, 100] in chat_template_kwargs" |
| 86 | ) |
| 87 | |
| 88 | prompt = encode_messages( |
| 89 | conversation, |
| 90 | thinking_mode="thinking" if thinking else "chat", |
| 91 | drop_thinking=kwargs.get("drop_thinking", True), |
| 92 | reasoning_effort=effort, |
| 93 | ) |
| 94 | if kwargs.get("tokenize", True): |
| 95 | tokenizer_kwargs = { |
| 96 | key: kwargs[key] |
| 97 | for key in ("truncation", "max_length") |
| 98 | if key in kwargs |
| 99 | } |
| 100 | return self.encode(prompt, add_special_tokens=False, **tokenizer_kwargs) |
| 101 | return prompt |
| 102 | |
| 103 | def num_special_tokens_to_add(self) -> int: |
| 104 | return len(self.encode("")) |
| 105 | |
| 106 | def get_added_vocab(self) -> dict[str, int]: |
| 107 | return added_vocab.copy() |
| 108 | |
| 109 | def __reduce__(self): |
| 110 | return get_deepseek_v41_tokenizer, (tokenizer,) |
| 111 | |
| 112 | _DeepseekV41Tokenizer.__name__ = f"DSV41{tokenizer.__class__.__name__}" |
| 113 | wrapped.__class__ = _DeepseekV41Tokenizer |
| 114 | return wrapped |
| 115 | |
| 116 | |
| 117 | class DeepseekV41Tokenizer(TokenizerLike): |
| 118 | @classmethod |
| 119 | def from_pretrained(cls, *args, **kwargs) -> HfTokenizer: |
| 120 | tokenizer = TokenizersBackend.from_pretrained(*args, **kwargs) |
| 121 | return get_cached_tokenizer(get_deepseek_v41_tokenizer(tokenizer)) |