main
py 196 lines 7.59 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from abc import abstractmethod
5 from collections.abc import Iterable, Sequence
6 from itertools import islice
7 from typing import TYPE_CHECKING
8
9 from vllm.entrypoints.generate.base.protocol import DeltaMessage
10 from vllm.reasoning.abs_reasoning_parsers import ReasoningParser
11 from vllm.tokenizers import TokenizerLike
12
13 if TYPE_CHECKING:
14 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
15 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
16
17
18 class BaseThinkingReasoningParser(ReasoningParser):
19 """Base class for reasoning parsers that use thinking tokens.
20
21 This class provides common functionality for parsers that use start and end
22 tokens to delimit reasoning content (
23 e.g., <think>...</think>, <seed:think>...</seed:think>).
24
25 Subclasses must implement the start and end tokens via abstract
26 properties.
27 """
28
29 @property
30 @abstractmethod
31 def start_token(self) -> str:
32 """The token that starts reasoning content."""
33 raise NotImplementedError
34
35 @property
36 @abstractmethod
37 def end_token(self) -> str:
38 """The token that ends reasoning content."""
39 raise NotImplementedError
40
41 @property
42 def reasoning_start_str(self) -> str:
43 return self.start_token
44
45 @property
46 def reasoning_end_str(self) -> str:
47 return self.end_token
48
49 def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
50 super().__init__(tokenizer, *args, **kwargs)
51
52 if not self.model_tokenizer:
53 raise ValueError(
54 "The model tokenizer must be passed to the ReasoningParser "
55 "constructor during construction."
56 )
57
58 if not self.start_token or not self.end_token:
59 raise ValueError("start_token and end_token must be defined in subclasses")
60
61 start_token_id = self.vocab.get(self.start_token)
62 end_token_id = self.vocab.get(self.end_token)
63 if start_token_id is None or end_token_id is None:
64 raise RuntimeError(
65 f"{self.__class__.__name__} reasoning parser could not locate "
66 "think start/end tokens in the tokenizer!"
67 )
68 self.start_token_id: int = start_token_id
69 self.end_token_id: int = end_token_id
70
71 def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
72 start_token_id = self.start_token_id
73 end_token_id = self.end_token_id
74
75 for i in range(len(input_ids) - 1, -1, -1):
76 if input_ids[i] == start_token_id:
77 return False
78 if input_ids[i] == end_token_id:
79 return True
80 return False
81
82 def is_reasoning_end_streaming(
83 self, input_ids: Sequence[int], delta_ids: Iterable[int]
84 ) -> bool:
85 end_token_id = self.end_token_id
86 return end_token_id in delta_ids
87
88 def extract_content_ids(self, input_ids: list[int]) -> list[int]:
89 """Extract the content after the end tokens."""
90 if self.end_token_id not in islice(input_ids, 0, max(0, len(input_ids) - 1)):
91 return []
92 else:
93 return input_ids[input_ids.index(self.end_token_id) + 1 :]
94
95 def extract_reasoning_streaming(
96 self,
97 previous_text: str,
98 current_text: str,
99 delta_text: str,
100 previous_token_ids: Sequence[int],
101 current_token_ids: Sequence[int],
102 delta_token_ids: Sequence[int],
103 ) -> DeltaMessage | None:
104 """Extract reasoning content from a delta message.
105 Handles streaming output where previous + delta = current.
106 Uses token IDs for faster processing.
107 """
108 # Skip single special tokens
109 if len(delta_token_ids) == 1 and (
110 delta_token_ids[0] in [self.start_token_id, self.end_token_id]
111 ):
112 return None
113
114 # Check if start token is present in previous or delta.
115 # Keep compatibility with models that don't generate start tokens.
116 if self.start_token_id in previous_token_ids:
117 if self.end_token_id in delta_token_ids:
118 # start token in previous, end token in delta,
119 # extract reasoning content
120 end_index = delta_text.find(self.end_token)
121 reasoning = delta_text[:end_index]
122 content = delta_text[end_index + len(self.end_token) :]
123 return DeltaMessage(
124 reasoning=reasoning, content=content if content else None
125 )
126 elif self.end_token_id in previous_token_ids:
127 # start token in previous, end token in previous,
128 # reasoning content continues
129 return DeltaMessage(content=delta_text)
130 else:
131 # start token in previous, no end token in previous or delta,
132 # reasoning content continues
133 return DeltaMessage(reasoning=delta_text)
134 elif self.start_token_id in delta_token_ids:
135 if self.end_token_id in delta_token_ids:
136 # start token in delta, end token in delta,
137 # extract reasoning content
138 start_index = delta_text.find(self.start_token)
139 end_index = delta_text.find(self.end_token)
140 reasoning = delta_text[start_index + len(self.start_token) : end_index]
141 content = delta_text[end_index + len(self.end_token) :]
142 return DeltaMessage(
143 reasoning=reasoning, content=content if content else None
144 )
145 else:
146 # start token in delta, no end token in delta,
147 # reasoning content continues
148 return DeltaMessage(reasoning=delta_text)
149 else:
150 # not find thinking start token
151 return DeltaMessage(content=delta_text)
152
153 def extract_reasoning(
154 self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest"
155 ) -> tuple[str | None, str | None]:
156 """Extract reasoning content from the model output.
157
158 This is the base implementation that works for most models.
159 Subclasses can override this method for specific behavior.
160 """
161 # Check if the start token is present in the model output, remove it
162 # if it is present.
163 model_output_parts = model_output.partition(self.start_token)
164 model_output = (
165 model_output_parts[2] if model_output_parts[1] else model_output_parts[0]
166 )
167
168 # For models that may not generate start token,
169 # assume the reasoning content is always at the start.
170 if self.end_token not in model_output:
171 return model_output, None
172 else:
173 reasoning, _, content = model_output.partition(self.end_token)
174 # If generation stops right after end-of-think, return null content
175 final_content = content or None
176 return reasoning, final_content
177
178 def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
179 """Count tokens that fall within start/end thinking markers.
180
181 Uses a depth counter so nested spans are handled safely and stray end
182 tokens do not drive the counter negative.
183 """
184 count = 0
185 depth = 0
186 for token_id in token_ids:
187 if token_id == self.start_token_id:
188 depth += 1
189 continue
190 if token_id == self.end_token_id:
191 if depth > 0:
192 depth -= 1
193 continue
194 if depth > 0:
195 count += 1
196 return count