| 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 vllm.entrypoints.generate.base.protocol import DeltaMessage |
| 8 | from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser |
| 9 | |
| 10 | if TYPE_CHECKING: |
| 11 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 12 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 13 | |
| 14 | |
| 15 | class MiniMaxM3ReasoningParser(BaseThinkingReasoningParser): |
| 16 | """Reasoning parser for MiniMax M3 explicit thinking blocks. |
| 17 | |
| 18 | MiniMax M3 emits reasoning as: |
| 19 | |
| 20 | <mm:think>reasoning text</mm:think>assistant content |
| 21 | |
| 22 | The M3 tokenizer exposes both markers as complete vocabulary entries, but |
| 23 | generated marker text may be tokenized into smaller pieces. The streaming |
| 24 | parser therefore uses text markers for extraction instead of relying on the |
| 25 | single vocabulary IDs. The chat template may also prefill the start marker |
| 26 | when ``thinking_mode="enabled"``, so generated text can begin directly |
| 27 | inside a reasoning block without emitting ``<mm:think>`` again. |
| 28 | """ |
| 29 | |
| 30 | @property |
| 31 | def start_token(self) -> str: |
| 32 | return "<mm:think>" |
| 33 | |
| 34 | @property |
| 35 | def end_token(self) -> str: |
| 36 | return "</mm:think>" |
| 37 | |
| 38 | def __init__(self, tokenizer, *args, **kwargs): |
| 39 | super().__init__(tokenizer, *args, **kwargs) |
| 40 | self._start_token_ids = self._encode_marker(self.start_token) |
| 41 | self._end_token_ids = self._encode_marker(self.end_token) |
| 42 | chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} |
| 43 | self._initial_in_reasoning = chat_kwargs.get("thinking_mode") == "enabled" |
| 44 | self._reasoning_ended_streaming = False |
| 45 | self._reasoning_active_streaming = self._initial_in_reasoning |
| 46 | self._pending_marker_streaming = False |
| 47 | self._last_streaming_delta_token_ids: tuple[int, ...] | None = None |
| 48 | self._last_streaming_content_token_ids: list[int] | None = None |
| 49 | |
| 50 | def _encode_text(self, text: str) -> list[int]: |
| 51 | try: |
| 52 | return list(self.model_tokenizer.encode(text, add_special_tokens=False)) |
| 53 | except TypeError: |
| 54 | return list(self.model_tokenizer.encode(text)) |
| 55 | |
| 56 | def _encode_marker(self, marker: str) -> tuple[int, ...]: |
| 57 | return tuple(self._encode_text(marker)) |
| 58 | |
| 59 | def _decode_text(self, token_ids: Sequence[int]) -> str: |
| 60 | try: |
| 61 | return self.model_tokenizer.decode( |
| 62 | list(token_ids), skip_special_tokens=False |
| 63 | ) |
| 64 | except TypeError: |
| 65 | return self.model_tokenizer.decode(list(token_ids)) |
| 66 | |
| 67 | def _content_suffix_token_ids( |
| 68 | self, |
| 69 | delta_text: str, |
| 70 | delta_token_ids: Sequence[int], |
| 71 | content: str | None, |
| 72 | ) -> list[int]: |
| 73 | if content is None: |
| 74 | return [] |
| 75 | if content == delta_text: |
| 76 | return list(delta_token_ids) |
| 77 | if delta_text.endswith(content): |
| 78 | prefix_text = delta_text[: len(delta_text) - len(content)] |
| 79 | for index in range(len(delta_token_ids) + 1): |
| 80 | if self._decode_text(delta_token_ids[:index]) == prefix_text: |
| 81 | return list(delta_token_ids[index:]) |
| 82 | return self._encode_text(content) |
| 83 | |
| 84 | @staticmethod |
| 85 | def _contains_token_sequence( |
| 86 | token_ids: Sequence[int], marker_ids: Sequence[int] |
| 87 | ) -> bool: |
| 88 | if not marker_ids or len(marker_ids) > len(token_ids): |
| 89 | return False |
| 90 | marker_len = len(marker_ids) |
| 91 | return any( |
| 92 | tuple(token_ids[i : i + marker_len]) == tuple(marker_ids) |
| 93 | for i in range(len(token_ids) - marker_len + 1) |
| 94 | ) |
| 95 | |
| 96 | @staticmethod |
| 97 | def _rfind_token_sequence( |
| 98 | token_ids: Sequence[int], marker_ids: Sequence[int] |
| 99 | ) -> int: |
| 100 | if not marker_ids or len(marker_ids) > len(token_ids): |
| 101 | return -1 |
| 102 | marker_len = len(marker_ids) |
| 103 | for i in range(len(token_ids) - marker_len, -1, -1): |
| 104 | if tuple(token_ids[i : i + marker_len]) == tuple(marker_ids): |
| 105 | return i |
| 106 | return -1 |
| 107 | |
| 108 | @staticmethod |
| 109 | def _ends_with_token_sequence_prefix( |
| 110 | token_ids: Sequence[int], marker_ids: Sequence[int] |
| 111 | ) -> bool: |
| 112 | if not marker_ids: |
| 113 | return False |
| 114 | max_len = min(len(token_ids), len(marker_ids) - 1) |
| 115 | for prefix_len in range(max_len, 0, -1): |
| 116 | if tuple(token_ids[-prefix_len:]) == tuple(marker_ids[:prefix_len]): |
| 117 | return True |
| 118 | return False |
| 119 | |
| 120 | @staticmethod |
| 121 | def _strip_partial_marker_suffix(text: str, marker: str) -> str: |
| 122 | max_len = min(len(text), len(marker) - 1) |
| 123 | for suffix_len in range(max_len, 0, -1): |
| 124 | if marker.startswith(text[-suffix_len:]): |
| 125 | return text[:-suffix_len] |
| 126 | return text |
| 127 | |
| 128 | @staticmethod |
| 129 | def _visible_delta(previous: str | None, current: str | None) -> str | None: |
| 130 | if not current: |
| 131 | return None |
| 132 | if not previous: |
| 133 | return current |
| 134 | if current.startswith(previous): |
| 135 | delta = current[len(previous) :] |
| 136 | return delta or None |
| 137 | return current |
| 138 | |
| 139 | def _visible_segments(self, text: str) -> tuple[str | None, str | None]: |
| 140 | if not text: |
| 141 | return None, None |
| 142 | |
| 143 | if not self._initial_in_reasoning: |
| 144 | if self.end_token.startswith(text) and len(text) < len(self.end_token): |
| 145 | return None, None |
| 146 | if text.startswith(self.end_token): |
| 147 | text = text[len(self.end_token) :] |
| 148 | if not text: |
| 149 | return None, None |
| 150 | |
| 151 | if self._initial_in_reasoning and self.start_token not in text: |
| 152 | reasoning, end, content = text.partition(self.end_token) |
| 153 | if end: |
| 154 | return reasoning or None, content or None |
| 155 | reasoning = self._strip_partial_marker_suffix(reasoning, self.end_token) |
| 156 | return reasoning or None, None |
| 157 | |
| 158 | if self.start_token not in text: |
| 159 | content = self._strip_partial_marker_suffix(text, self.start_token) |
| 160 | return None, content or None |
| 161 | |
| 162 | content_before, _, after_start = text.partition(self.start_token) |
| 163 | reasoning, end, content_after = after_start.partition(self.end_token) |
| 164 | if end: |
| 165 | return reasoning or None, (content_before + content_after) or None |
| 166 | |
| 167 | reasoning = self._strip_partial_marker_suffix(reasoning, self.end_token) |
| 168 | return reasoning or None, content_before or None |
| 169 | |
| 170 | def extract_reasoning( |
| 171 | self, |
| 172 | model_output: str, |
| 173 | request: "ChatCompletionRequest | ResponsesRequest", |
| 174 | ) -> tuple[str | None, str | None]: |
| 175 | # MiniMax M3 can start a response with a stray closer. Drop that first |
| 176 | # token only; later unmatched closers stay visible as content. |
| 177 | if not self._initial_in_reasoning and model_output.startswith(self.end_token): |
| 178 | content = model_output[len(self.end_token) :] |
| 179 | return None, content or None |
| 180 | |
| 181 | if self._initial_in_reasoning and self.start_token not in model_output: |
| 182 | reasoning, end, content = model_output.partition(self.end_token) |
| 183 | if not end: |
| 184 | return model_output, None |
| 185 | return reasoning, content or None |
| 186 | |
| 187 | if self.start_token not in model_output: |
| 188 | return None, model_output |
| 189 | |
| 190 | content_before, _, after_start = model_output.partition(self.start_token) |
| 191 | reasoning, end, content_after = after_start.partition(self.end_token) |
| 192 | if not end: |
| 193 | return reasoning, content_before or None |
| 194 | |
| 195 | return reasoning, (content_before + content_after) or None |
| 196 | |
| 197 | def is_reasoning_end_streaming( |
| 198 | self, input_ids: Sequence[int], delta_ids: Iterable[int] |
| 199 | ) -> bool: |
| 200 | if self._reasoning_ended_streaming: |
| 201 | return True |
| 202 | |
| 203 | if self._reasoning_active_streaming or self._pending_marker_streaming: |
| 204 | return False |
| 205 | |
| 206 | delta_ids = tuple(delta_ids) |
| 207 | if self._contains_token_sequence(delta_ids, self._end_token_ids): |
| 208 | return True |
| 209 | if self._contains_token_sequence(input_ids, self._end_token_ids): |
| 210 | return True |
| 211 | if self._initial_in_reasoning: |
| 212 | return False |
| 213 | if self._ends_with_token_sequence_prefix(input_ids, self._start_token_ids): |
| 214 | return False |
| 215 | if self._ends_with_token_sequence_prefix(input_ids, self._end_token_ids): |
| 216 | return False |
| 217 | if not self._contains_token_sequence(input_ids, self._start_token_ids): |
| 218 | return bool(input_ids) |
| 219 | return False |
| 220 | |
| 221 | def extract_content_ids(self, input_ids: list[int]) -> list[int]: |
| 222 | if ( |
| 223 | self._last_streaming_delta_token_ids == tuple(input_ids) |
| 224 | and self._last_streaming_content_token_ids is not None |
| 225 | ): |
| 226 | content_ids = self._last_streaming_content_token_ids |
| 227 | self._last_streaming_delta_token_ids = None |
| 228 | self._last_streaming_content_token_ids = None |
| 229 | return list(content_ids) |
| 230 | |
| 231 | end_index = self._rfind_token_sequence(input_ids, self._end_token_ids) |
| 232 | if end_index >= 0: |
| 233 | return input_ids[end_index + len(self._end_token_ids) :] |
| 234 | |
| 235 | has_start = self._contains_token_sequence(input_ids, self._start_token_ids) |
| 236 | if self._initial_in_reasoning and not has_start: |
| 237 | return [] |
| 238 | |
| 239 | if not has_start: |
| 240 | return input_ids |
| 241 | return [] |
| 242 | |
| 243 | def extract_reasoning_streaming( |
| 244 | self, |
| 245 | previous_text: str, |
| 246 | current_text: str, |
| 247 | delta_text: str, |
| 248 | previous_token_ids: Sequence[int], |
| 249 | current_token_ids: Sequence[int], |
| 250 | delta_token_ids: Sequence[int], |
| 251 | ) -> DeltaMessage | None: |
| 252 | if not delta_text: |
| 253 | return None |
| 254 | |
| 255 | if not previous_text: |
| 256 | self._reasoning_ended_streaming = False |
| 257 | self._reasoning_active_streaming = self._initial_in_reasoning |
| 258 | self._pending_marker_streaming = False |
| 259 | self._last_streaming_delta_token_ids = None |
| 260 | self._last_streaming_content_token_ids = None |
| 261 | previous_reasoning, previous_content = self._visible_segments(previous_text) |
| 262 | current_reasoning, current_content = self._visible_segments(current_text) |
| 263 | if self.end_token in current_text or current_content is not None: |
| 264 | self._reasoning_ended_streaming = True |
| 265 | self._reasoning_active_streaming = False |
| 266 | self._pending_marker_streaming = False |
| 267 | else: |
| 268 | self._last_streaming_delta_token_ids = None |
| 269 | self._last_streaming_content_token_ids = None |
| 270 | self._reasoning_active_streaming = ( |
| 271 | self._initial_in_reasoning |
| 272 | or self.start_token in current_text |
| 273 | or current_reasoning is not None |
| 274 | ) |
| 275 | self._pending_marker_streaming = not self._reasoning_active_streaming and ( |
| 276 | self.start_token.startswith(current_text) |
| 277 | or self.end_token.startswith(current_text) |
| 278 | ) |
| 279 | reasoning = self._visible_delta(previous_reasoning, current_reasoning) |
| 280 | content = self._visible_delta(previous_content, current_content) |
| 281 | if self._reasoning_ended_streaming: |
| 282 | self._last_streaming_delta_token_ids = tuple(delta_token_ids) |
| 283 | self._last_streaming_content_token_ids = self._content_suffix_token_ids( |
| 284 | delta_text, delta_token_ids, content |
| 285 | ) |
| 286 | if reasoning is None and content is None: |
| 287 | return None |
| 288 | return DeltaMessage(reasoning=reasoning, content=content) |
| 289 | |
| 290 | def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: |
| 291 | count = 0 |
| 292 | depth = 1 if self._initial_in_reasoning else 0 |
| 293 | i = 0 |
| 294 | while i < len(token_ids): |
| 295 | if tuple(token_ids[i : i + len(self._start_token_ids)]) == ( |
| 296 | self._start_token_ids |
| 297 | ): |
| 298 | depth += 1 |
| 299 | i += len(self._start_token_ids) |
| 300 | continue |
| 301 | if tuple(token_ids[i : i + len(self._end_token_ids)]) == ( |
| 302 | self._end_token_ids |
| 303 | ): |
| 304 | if depth > 0: |
| 305 | depth -= 1 |
| 306 | i += len(self._end_token_ids) |
| 307 | continue |
| 308 | if depth > 0: |
| 309 | count += 1 |
| 310 | i += 1 |
| 311 | return count |
| 312 | |
| 313 | def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: |
| 314 | start_index = self._rfind_token_sequence(input_ids, self._start_token_ids) |
| 315 | end_index = self._rfind_token_sequence(input_ids, self._end_token_ids) |
| 316 | if end_index < 0: |
| 317 | return False |
| 318 | if start_index < 0: |
| 319 | return True |
| 320 | return end_index > start_index |