main
py 96 lines 3.63 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Iterable, Sequence
5 from typing import TYPE_CHECKING
6
7 from transformers import PreTrainedTokenizerBase
8
9 from vllm.reasoning import ReasoningParser
10 from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
11
12 from .identity_reasoning_parser import IdentityReasoningParser
13
14 if TYPE_CHECKING:
15 from vllm.entrypoints.generate.base.protocol import DeltaMessage
16 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
17 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
18
19
20 class DeepSeekV3ReasoningParser(ReasoningParser):
21 """V3 parser that delegates to either DeepSeekR1ReasoningParser or
22 IdentityReasoningParser based on `thinking` and `separate_reasoning`.
23 """
24
25 def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs):
26 super().__init__(tokenizer, *args, **kwargs)
27
28 chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
29 thinking = bool(chat_kwargs.get("thinking", False))
30 enable_thinking = bool(chat_kwargs.get("enable_thinking", False))
31 thinking = thinking or enable_thinking
32
33 self._parser: ReasoningParser
34 if thinking:
35 self._parser = DeepSeekR1ReasoningParser(tokenizer, *args, **kwargs)
36 else:
37 self._parser = IdentityReasoningParser(tokenizer, *args, **kwargs)
38
39 @property
40 def reasoning_start_str(self) -> str | None:
41 return self._parser.reasoning_start_str
42
43 @property
44 def reasoning_end_str(self) -> str | None:
45 return self._parser.reasoning_end_str
46
47 def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
48 return self._parser.is_reasoning_end(input_ids)
49
50 def is_reasoning_end_streaming(
51 self, input_ids: Sequence[int], delta_ids: Iterable[int]
52 ) -> bool:
53 return self._parser.is_reasoning_end_streaming(input_ids, delta_ids)
54
55 def extract_content_ids(self, input_ids: list[int]) -> list[int]:
56 return self._parser.extract_content_ids(input_ids)
57
58 def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
59 return self._parser.count_reasoning_tokens(token_ids)
60
61 def extract_reasoning(
62 self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest"
63 ) -> tuple[str | None, str | None]:
64 return self._parser.extract_reasoning(model_output, request)
65
66 def extract_reasoning_streaming(
67 self,
68 previous_text: str,
69 current_text: str,
70 delta_text: str,
71 previous_token_ids: Sequence[int],
72 current_token_ids: Sequence[int],
73 delta_token_ids: Sequence[int],
74 ) -> "DeltaMessage | None":
75 return self._parser.extract_reasoning_streaming(
76 previous_text,
77 current_text,
78 delta_text,
79 previous_token_ids,
80 current_token_ids,
81 delta_token_ids,
82 )
83
84
85 class DeepSeekV3ReasoningWithThinkingParser(DeepSeekV3ReasoningParser):
86 """DeepSeekV3ReasoningParser that defaults to thinking mode."""
87
88 def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs):
89 chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
90 thinking = chat_kwargs.get("thinking", None)
91 enable_thinking = chat_kwargs.get("enable_thinking", None)
92 if thinking is None and enable_thinking is None:
93 chat_kwargs["thinking"] = True
94 chat_kwargs["enable_thinking"] = True
95 kwargs["chat_template_kwargs"] = chat_kwargs
96 super().__init__(tokenizer, *args, **kwargs)