mirror of
https://github.com/vllm-project/vllm.git
synced 2026-08-17 11:10:16 +00:00
[Bugfix] Parse MiniMax M3 streaming reasoning by text markers (#45718)
Signed-off-by: test test <[email protected]>
This commit is contained in:
Rui Yin
@@ -83,6 +83,20 @@ class MiniMaxM3Tokenizer:
|
||||
return "".join(tokens)
|
||||
|
||||
|
||||
class SplitMiniMaxM3Tokenizer(MiniMaxM3Tokenizer):
|
||||
"""Tokenizer that exposes marker vocab entries but encodes them as text."""
|
||||
|
||||
def tokenize(self, text: str) -> list[str]:
|
||||
return list(text)
|
||||
|
||||
|
||||
class RuntimeSplitMiniMaxM3Tokenizer(MiniMaxM3Tokenizer):
|
||||
"""Tokenizer whose runtime output splits markers despite atomic encodes."""
|
||||
|
||||
def encode_runtime(self, text: str) -> list[int]:
|
||||
return [self._add_token(token) for token in list(text)]
|
||||
|
||||
|
||||
def make_parser(
|
||||
chat_template_kwargs: dict[str, str] | None = None,
|
||||
) -> tuple[MiniMaxM3ReasoningParser, MiniMaxM3Tokenizer]:
|
||||
@@ -105,7 +119,8 @@ def run_streaming(
|
||||
reasoning_end_states: list[bool] = []
|
||||
|
||||
for chunk in chunks:
|
||||
delta_token_ids = tokenizer.encode(chunk, add_special_tokens=False)
|
||||
encode_runtime = getattr(tokenizer, "encode_runtime", tokenizer.encode)
|
||||
delta_token_ids = encode_runtime(chunk)
|
||||
current_text = previous_text + chunk
|
||||
current_token_ids = previous_token_ids + delta_token_ids
|
||||
delta = parser.extract_reasoning_streaming(
|
||||
@@ -288,6 +303,132 @@ def test_streaming_plain_content_ends_reasoning_phase():
|
||||
assert end_states == [True, True]
|
||||
|
||||
|
||||
def test_streaming_split_marker_tokens_are_not_returned():
|
||||
tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
|
||||
parser = MiniMaxM3ReasoningParser(tokenizer)
|
||||
|
||||
reasoning, content, end_states = run_streaming(
|
||||
parser,
|
||||
tokenizer,
|
||||
["<mm:think>", "Reasoning", " content", "</mm:think>", "content"],
|
||||
)
|
||||
|
||||
assert reasoning == "Reasoning content"
|
||||
assert content == "content"
|
||||
assert end_states == [False, False, False, True, True]
|
||||
|
||||
|
||||
def test_streaming_split_marker_text_drives_end_state():
|
||||
tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
|
||||
parser = MiniMaxM3ReasoningParser(tokenizer)
|
||||
previous_text = ""
|
||||
previous_token_ids: list[int] = []
|
||||
|
||||
for chunk in ["<mm:think>", "Reasoning", " content", "</mm:think>"]:
|
||||
delta_token_ids = tokenizer.encode_runtime(chunk)
|
||||
current_text = previous_text + chunk
|
||||
current_token_ids = previous_token_ids + delta_token_ids
|
||||
parser.extract_reasoning_streaming(
|
||||
previous_text=previous_text,
|
||||
current_text=current_text,
|
||||
delta_text=chunk,
|
||||
previous_token_ids=previous_token_ids,
|
||||
current_token_ids=current_token_ids,
|
||||
delta_token_ids=delta_token_ids,
|
||||
)
|
||||
previous_text = current_text
|
||||
previous_token_ids = current_token_ids
|
||||
|
||||
assert parser.is_reasoning_end_streaming(previous_token_ids, []) is True
|
||||
|
||||
|
||||
def test_streaming_split_end_marker_content_ids_are_stripped():
|
||||
tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
|
||||
parser = MiniMaxM3ReasoningParser(tokenizer)
|
||||
previous_text = "<mm:think>Reasoning"
|
||||
previous_token_ids = tokenizer.encode_runtime(previous_text)
|
||||
delta_text = "</mm:think>content"
|
||||
delta_token_ids = tokenizer.encode_runtime(delta_text)
|
||||
current_token_ids = previous_token_ids + delta_token_ids
|
||||
|
||||
parser.extract_reasoning_streaming(
|
||||
previous_text=previous_text,
|
||||
current_text=previous_text + delta_text,
|
||||
delta_text=delta_text,
|
||||
previous_token_ids=previous_token_ids,
|
||||
current_token_ids=current_token_ids,
|
||||
delta_token_ids=delta_token_ids,
|
||||
)
|
||||
|
||||
assert parser.is_reasoning_end_streaming(current_token_ids, delta_token_ids)
|
||||
assert tokenizer.decode(parser.extract_content_ids(delta_token_ids)) == "content"
|
||||
|
||||
|
||||
def test_streaming_split_marker_tokens_enabled_mode():
|
||||
tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
|
||||
parser = MiniMaxM3ReasoningParser(
|
||||
tokenizer, chat_template_kwargs={"thinking_mode": "enabled"}
|
||||
)
|
||||
|
||||
reasoning, content, end_states = run_streaming(
|
||||
parser,
|
||||
tokenizer,
|
||||
["Reasoning", " content", "</mm:think>", "content"],
|
||||
)
|
||||
|
||||
assert reasoning == "Reasoning content"
|
||||
assert content == "content"
|
||||
assert end_states == [False, False, True, True]
|
||||
|
||||
|
||||
def test_streaming_split_marker_text_across_deltas():
|
||||
tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
|
||||
parser = MiniMaxM3ReasoningParser(tokenizer)
|
||||
|
||||
reasoning, content, end_states = run_streaming(
|
||||
parser,
|
||||
tokenizer,
|
||||
["<mm:", "think>", "Reasoning", " content", "</mm:", "think>", "content"],
|
||||
)
|
||||
|
||||
assert reasoning == "Reasoning content"
|
||||
assert content == "content"
|
||||
assert end_states == [False, False, False, False, False, True, True]
|
||||
|
||||
|
||||
def test_streaming_split_leading_end_marker_text_across_deltas():
|
||||
tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
|
||||
parser = MiniMaxM3ReasoningParser(tokenizer)
|
||||
|
||||
reasoning, content, end_states = run_streaming(
|
||||
parser,
|
||||
tokenizer,
|
||||
["</mm:", "think>", "content"],
|
||||
)
|
||||
|
||||
assert reasoning is None
|
||||
assert content == "content"
|
||||
assert end_states == [False, True, True]
|
||||
|
||||
|
||||
def test_token_id_helpers_with_split_marker_tokens():
|
||||
tokenizer = SplitMiniMaxM3Tokenizer()
|
||||
parser = MiniMaxM3ReasoningParser(tokenizer)
|
||||
output_ids = tokenizer.encode(
|
||||
"<mm:think>abc</mm:think>def", add_special_tokens=False
|
||||
)
|
||||
open_reasoning_ids = tokenizer.encode("<mm:think>abc", add_special_tokens=False)
|
||||
content_ids = tokenizer.encode("plain", add_special_tokens=False)
|
||||
|
||||
assert parser.is_reasoning_end(output_ids)
|
||||
assert not parser.is_reasoning_end(open_reasoning_ids)
|
||||
assert not parser.is_reasoning_end(content_ids)
|
||||
assert tokenizer.decode(parser.extract_content_ids(output_ids)) == "def"
|
||||
assert parser.extract_content_ids(open_reasoning_ids) == []
|
||||
assert parser.extract_content_ids(content_ids) == content_ids
|
||||
assert parser.count_reasoning_tokens(output_ids) == len(tokenizer.encode("abc"))
|
||||
|
||||
|
||||
def test_token_id_helpers():
|
||||
parser, tokenizer = make_parser()
|
||||
output_ids = tokenizer.encode(
|
||||
|
||||
@@ -19,10 +19,12 @@ class MiniMaxM3ReasoningParser(BaseThinkingReasoningParser):
|
||||
|
||||
<mm:think>reasoning text</mm:think>assistant content
|
||||
|
||||
The M3 tokenizer exposes both markers as complete vocabulary tokens. The
|
||||
chat template may also prefill the start marker when
|
||||
``thinking_mode="enabled"``, so generated text can begin directly inside a
|
||||
reasoning block without emitting ``<mm:think>`` again.
|
||||
The M3 tokenizer exposes both markers as complete vocabulary entries, but
|
||||
generated marker text may be tokenized into smaller pieces. The streaming
|
||||
parser therefore uses text markers for extraction instead of relying on the
|
||||
single vocabulary IDs. The chat template may also prefill the start marker
|
||||
when ``thinking_mode="enabled"``, so generated text can begin directly
|
||||
inside a reasoning block without emitting ``<mm:think>`` again.
|
||||
"""
|
||||
|
||||
@property
|
||||
@@ -35,9 +37,135 @@ class MiniMaxM3ReasoningParser(BaseThinkingReasoningParser):
|
||||
|
||||
def __init__(self, tokenizer, *args, **kwargs):
|
||||
super().__init__(tokenizer, *args, **kwargs)
|
||||
self._start_token_ids = self._encode_marker(self.start_token)
|
||||
self._end_token_ids = self._encode_marker(self.end_token)
|
||||
chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
|
||||
self._initial_in_reasoning = chat_kwargs.get("thinking_mode") == "enabled"
|
||||
self._at_response_start = True
|
||||
self._reasoning_ended_streaming = False
|
||||
self._reasoning_active_streaming = self._initial_in_reasoning
|
||||
self._pending_marker_streaming = False
|
||||
self._last_streaming_delta_token_ids: tuple[int, ...] | None = None
|
||||
self._last_streaming_content_token_ids: list[int] | None = None
|
||||
|
||||
def _encode_text(self, text: str) -> list[int]:
|
||||
try:
|
||||
return list(self.model_tokenizer.encode(text, add_special_tokens=False))
|
||||
except TypeError:
|
||||
return list(self.model_tokenizer.encode(text))
|
||||
|
||||
def _encode_marker(self, marker: str) -> tuple[int, ...]:
|
||||
return tuple(self._encode_text(marker))
|
||||
|
||||
def _decode_text(self, token_ids: Sequence[int]) -> str:
|
||||
try:
|
||||
return self.model_tokenizer.decode(
|
||||
list(token_ids), skip_special_tokens=False
|
||||
)
|
||||
except TypeError:
|
||||
return self.model_tokenizer.decode(list(token_ids))
|
||||
|
||||
def _content_suffix_token_ids(
|
||||
self,
|
||||
delta_text: str,
|
||||
delta_token_ids: Sequence[int],
|
||||
content: str | None,
|
||||
) -> list[int]:
|
||||
if content is None:
|
||||
return []
|
||||
if content == delta_text:
|
||||
return list(delta_token_ids)
|
||||
if delta_text.endswith(content):
|
||||
prefix_text = delta_text[: len(delta_text) - len(content)]
|
||||
for index in range(len(delta_token_ids) + 1):
|
||||
if self._decode_text(delta_token_ids[:index]) == prefix_text:
|
||||
return list(delta_token_ids[index:])
|
||||
return self._encode_text(content)
|
||||
|
||||
@staticmethod
|
||||
def _contains_token_sequence(
|
||||
token_ids: Sequence[int], marker_ids: Sequence[int]
|
||||
) -> bool:
|
||||
if not marker_ids or len(marker_ids) > len(token_ids):
|
||||
return False
|
||||
marker_len = len(marker_ids)
|
||||
return any(
|
||||
tuple(token_ids[i : i + marker_len]) == tuple(marker_ids)
|
||||
for i in range(len(token_ids) - marker_len + 1)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _rfind_token_sequence(
|
||||
token_ids: Sequence[int], marker_ids: Sequence[int]
|
||||
) -> int:
|
||||
if not marker_ids or len(marker_ids) > len(token_ids):
|
||||
return -1
|
||||
marker_len = len(marker_ids)
|
||||
for i in range(len(token_ids) - marker_len, -1, -1):
|
||||
if tuple(token_ids[i : i + marker_len]) == tuple(marker_ids):
|
||||
return i
|
||||
return -1
|
||||
|
||||
@staticmethod
|
||||
def _ends_with_token_sequence_prefix(
|
||||
token_ids: Sequence[int], marker_ids: Sequence[int]
|
||||
) -> bool:
|
||||
if not marker_ids:
|
||||
return False
|
||||
max_len = min(len(token_ids), len(marker_ids) - 1)
|
||||
for prefix_len in range(max_len, 0, -1):
|
||||
if tuple(token_ids[-prefix_len:]) == tuple(marker_ids[:prefix_len]):
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _strip_partial_marker_suffix(text: str, marker: str) -> str:
|
||||
max_len = min(len(text), len(marker) - 1)
|
||||
for suffix_len in range(max_len, 0, -1):
|
||||
if marker.startswith(text[-suffix_len:]):
|
||||
return text[:-suffix_len]
|
||||
return text
|
||||
|
||||
@staticmethod
|
||||
def _visible_delta(previous: str | None, current: str | None) -> str | None:
|
||||
if not current:
|
||||
return None
|
||||
if not previous:
|
||||
return current
|
||||
if current.startswith(previous):
|
||||
delta = current[len(previous) :]
|
||||
return delta or None
|
||||
return current
|
||||
|
||||
def _visible_segments(self, text: str) -> tuple[str | None, str | None]:
|
||||
if not text:
|
||||
return None, None
|
||||
|
||||
if not self._initial_in_reasoning:
|
||||
if self.end_token.startswith(text) and len(text) < len(self.end_token):
|
||||
return None, None
|
||||
if text.startswith(self.end_token):
|
||||
text = text[len(self.end_token) :]
|
||||
if not text:
|
||||
return None, None
|
||||
|
||||
if self._initial_in_reasoning and self.start_token not in text:
|
||||
reasoning, end, content = text.partition(self.end_token)
|
||||
if end:
|
||||
return reasoning or None, content or None
|
||||
reasoning = self._strip_partial_marker_suffix(reasoning, self.end_token)
|
||||
return reasoning or None, None
|
||||
|
||||
if self.start_token not in text:
|
||||
content = self._strip_partial_marker_suffix(text, self.start_token)
|
||||
return None, content or None
|
||||
|
||||
content_before, _, after_start = text.partition(self.start_token)
|
||||
reasoning, end, content_after = after_start.partition(self.end_token)
|
||||
if end:
|
||||
return reasoning or None, (content_before + content_after) or None
|
||||
|
||||
reasoning = self._strip_partial_marker_suffix(reasoning, self.end_token)
|
||||
return reasoning or None, content_before or None
|
||||
|
||||
def extract_reasoning(
|
||||
self,
|
||||
@@ -69,26 +197,46 @@ class MiniMaxM3ReasoningParser(BaseThinkingReasoningParser):
|
||||
def is_reasoning_end_streaming(
|
||||
self, input_ids: Sequence[int], delta_ids: Iterable[int]
|
||||
) -> bool:
|
||||
delta_ids = tuple(delta_ids)
|
||||
if self.end_token_id in delta_ids:
|
||||
if self._reasoning_ended_streaming:
|
||||
return True
|
||||
if self.end_token_id in input_ids:
|
||||
|
||||
if self._reasoning_active_streaming or self._pending_marker_streaming:
|
||||
return False
|
||||
|
||||
delta_ids = tuple(delta_ids)
|
||||
if self._contains_token_sequence(delta_ids, self._end_token_ids):
|
||||
return True
|
||||
if self._contains_token_sequence(input_ids, self._end_token_ids):
|
||||
return True
|
||||
if self._initial_in_reasoning:
|
||||
return False
|
||||
if self.start_token_id not in input_ids:
|
||||
if self._ends_with_token_sequence_prefix(input_ids, self._start_token_ids):
|
||||
return False
|
||||
if self._ends_with_token_sequence_prefix(input_ids, self._end_token_ids):
|
||||
return False
|
||||
if not self._contains_token_sequence(input_ids, self._start_token_ids):
|
||||
return bool(input_ids)
|
||||
return False
|
||||
|
||||
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
|
||||
if self.end_token_id in input_ids:
|
||||
end_index = len(input_ids) - 1 - input_ids[::-1].index(self.end_token_id)
|
||||
return input_ids[end_index + 1 :]
|
||||
if (
|
||||
self._last_streaming_delta_token_ids == tuple(input_ids)
|
||||
and self._last_streaming_content_token_ids is not None
|
||||
):
|
||||
content_ids = self._last_streaming_content_token_ids
|
||||
self._last_streaming_delta_token_ids = None
|
||||
self._last_streaming_content_token_ids = None
|
||||
return list(content_ids)
|
||||
|
||||
if self._initial_in_reasoning and self.start_token_id not in input_ids:
|
||||
end_index = self._rfind_token_sequence(input_ids, self._end_token_ids)
|
||||
if end_index >= 0:
|
||||
return input_ids[end_index + len(self._end_token_ids) :]
|
||||
|
||||
has_start = self._contains_token_sequence(input_ids, self._start_token_ids)
|
||||
if self._initial_in_reasoning and not has_start:
|
||||
return []
|
||||
|
||||
if self.start_token_id not in input_ids:
|
||||
if not has_start:
|
||||
return input_ids
|
||||
return []
|
||||
|
||||
@@ -104,68 +252,69 @@ class MiniMaxM3ReasoningParser(BaseThinkingReasoningParser):
|
||||
if not delta_text:
|
||||
return None
|
||||
|
||||
if self._at_response_start and not self._initial_in_reasoning:
|
||||
# Apply the leading-closer tolerance once. Later unmatched closers
|
||||
# stay visible as content.
|
||||
self._at_response_start = False
|
||||
if delta_text.startswith(self.end_token):
|
||||
delta_text = delta_text[len(self.end_token) :]
|
||||
if not delta_text:
|
||||
return None
|
||||
if delta_token_ids and delta_token_ids[0] == self.end_token_id:
|
||||
delta_token_ids = delta_token_ids[1:]
|
||||
|
||||
if self.end_token_id in previous_token_ids:
|
||||
return DeltaMessage(content=delta_text)
|
||||
|
||||
if (
|
||||
self._initial_in_reasoning
|
||||
and self.start_token_id not in previous_token_ids
|
||||
and self.start_token_id not in delta_token_ids
|
||||
):
|
||||
if self.end_token_id in delta_token_ids:
|
||||
reasoning, _, content = delta_text.partition(self.end_token)
|
||||
return DeltaMessage(
|
||||
reasoning=reasoning or None,
|
||||
content=content or None,
|
||||
)
|
||||
return DeltaMessage(reasoning=delta_text)
|
||||
|
||||
if (
|
||||
self.start_token_id not in previous_token_ids
|
||||
and self.start_token_id not in delta_token_ids
|
||||
):
|
||||
return DeltaMessage(content=delta_text)
|
||||
|
||||
if self.end_token_id in delta_token_ids:
|
||||
reasoning_text, _, content = delta_text.partition(self.end_token)
|
||||
if self.start_token_id in delta_token_ids:
|
||||
_, _, reasoning_text = reasoning_text.partition(self.start_token)
|
||||
return DeltaMessage(
|
||||
reasoning=reasoning_text or None,
|
||||
content=content or None,
|
||||
if not previous_text:
|
||||
self._reasoning_ended_streaming = False
|
||||
self._reasoning_active_streaming = self._initial_in_reasoning
|
||||
self._pending_marker_streaming = False
|
||||
self._last_streaming_delta_token_ids = None
|
||||
self._last_streaming_content_token_ids = None
|
||||
previous_reasoning, previous_content = self._visible_segments(previous_text)
|
||||
current_reasoning, current_content = self._visible_segments(current_text)
|
||||
if self.end_token in current_text or current_content is not None:
|
||||
self._reasoning_ended_streaming = True
|
||||
self._reasoning_active_streaming = False
|
||||
self._pending_marker_streaming = False
|
||||
else:
|
||||
self._last_streaming_delta_token_ids = None
|
||||
self._last_streaming_content_token_ids = None
|
||||
self._reasoning_active_streaming = (
|
||||
self._initial_in_reasoning
|
||||
or self.start_token in current_text
|
||||
or current_reasoning is not None
|
||||
)
|
||||
|
||||
if self.start_token_id in delta_token_ids:
|
||||
_, _, reasoning = delta_text.partition(self.start_token)
|
||||
return DeltaMessage(reasoning=reasoning) if reasoning else None
|
||||
|
||||
return DeltaMessage(reasoning=delta_text)
|
||||
self._pending_marker_streaming = not self._reasoning_active_streaming and (
|
||||
self.start_token.startswith(current_text)
|
||||
or self.end_token.startswith(current_text)
|
||||
)
|
||||
reasoning = self._visible_delta(previous_reasoning, current_reasoning)
|
||||
content = self._visible_delta(previous_content, current_content)
|
||||
if self._reasoning_ended_streaming:
|
||||
self._last_streaming_delta_token_ids = tuple(delta_token_ids)
|
||||
self._last_streaming_content_token_ids = self._content_suffix_token_ids(
|
||||
delta_text, delta_token_ids, content
|
||||
)
|
||||
if reasoning is None and content is None:
|
||||
return None
|
||||
return DeltaMessage(reasoning=reasoning, content=content)
|
||||
|
||||
def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
|
||||
if not self._initial_in_reasoning:
|
||||
return super().count_reasoning_tokens(token_ids)
|
||||
|
||||
count = 0
|
||||
depth = 1
|
||||
for token_id in token_ids:
|
||||
if token_id == self.start_token_id:
|
||||
depth = 1 if self._initial_in_reasoning else 0
|
||||
i = 0
|
||||
while i < len(token_ids):
|
||||
if tuple(token_ids[i : i + len(self._start_token_ids)]) == (
|
||||
self._start_token_ids
|
||||
):
|
||||
depth += 1
|
||||
i += len(self._start_token_ids)
|
||||
continue
|
||||
if token_id == self.end_token_id:
|
||||
if tuple(token_ids[i : i + len(self._end_token_ids)]) == (
|
||||
self._end_token_ids
|
||||
):
|
||||
if depth > 0:
|
||||
depth -= 1
|
||||
i += len(self._end_token_ids)
|
||||
continue
|
||||
if depth > 0:
|
||||
count += 1
|
||||
i += 1
|
||||
return count
|
||||
|
||||
def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
|
||||
start_index = self._rfind_token_sequence(input_ids, self._start_token_ids)
|
||||
end_index = self._rfind_token_sequence(input_ids, self._end_token_ids)
|
||||
if end_index < 0:
|
||||
return False
|
||||
if start_index < 0:
|
||||
return True
|
||||
return end_index > start_index
|
||||
|
||||
Reference in New Issue
Block a user