def get_deepseek_v41_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer:
"""Wrap an HF tokenizer with the V4.1 prompt encoder."""
wrapped = copy.copy(tokenizer)
added_vocab = tokenizer.get_added_vocab()
class _DeepseekV41Tokenizer(tokenizer.__class__): # type: ignore
def apply_chat_template(
self,
messages: list[ChatCompletionMessageParam],
tools: list[dict[str, Any]] | None = None,
**kwargs,
) -> str | list[int]:
# The generic renderer's conversation has already flattened text
# parts with '\n'. V4.1 encodes the original messages with '\n\n'.
conversation = _normalize_messages(messages)
if tools:
system = next((m for m in conversation if m["role"] == "system"), None)
if system is None:
system = {"role": "system", "content": ""}
conversation.insert(0, system)
system["tools"] = tools
thinking = bool(kwargs.get("thinking") or kwargs.get("enable_thinking"))
if "thinking" not in kwargs and "enable_thinking" not in kwargs:
thinking = True
effort = kwargs.get("reasoning_effort")
if effort == "none":
thinking = False
effort = None
if effort is None:
effort = "high"
if not (
(type(effort) is int and 1 <= effort <= 100)
or (isinstance(effort, str) and effort in REASONING_EFFORT_MAPPINGS)
):
raise ValueError(
"DeepSeek V4.1 reasoning_effort must be low, high, xhigh, max, "
"or an integer within [1, 100] in chat_template_kwargs"
)
prompt = encode_messages(
conversation,
thinking_mode="thinking" if thinking else "chat",
drop_thinking=kwargs.get("drop_thinking", True),
reasoning_effort=effort,
)
if kwargs.get("tokenize", True):
tokenizer_kwargs = {
key: kwargs[key]
for key in ("truncation", "max_length")
if key in kwargs
}
return self.encode(prompt, add_special_tokens=False, **tokenizer_kwargs)
return prompt
def num_special_tokens_to_add(self) -> int:
return len(self.encode(""))
def get_added_vocab(self) -> dict[str, int]:
return added_vocab.copy()
def __reduce__(self):
return get_deepseek_v41_tokenizer, (tokenizer,)
_DeepseekV41Tokenizer.__name__ = f"DSV41{tokenizer.__class__.__name__}"
wrapped.__class__ = _DeepseekV41Tokenizer
return wrapped