Zhou TaichangandClaude Opus 5
aff7d64948
Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
575 lines
24 KiB
Python
575 lines
24 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from collections.abc import Iterable, Sequence
|
|
from itertools import islice
|
|
from typing import TYPE_CHECKING
|
|
|
|
from vllm.entrypoints.generate.base.protocol import DeltaMessage
|
|
from vllm.logger import init_logger
|
|
from vllm.reasoning import ReasoningParser
|
|
from vllm.tokenizers import TokenizerLike
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
|
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
|
|
|
|
|
|
class StepAudioReasoningParser(ReasoningParser):
|
|
"""Reasoning parser for Step-Audio models.
|
|
|
|
Step-Audio supports two representations of thinking markers:
|
|
|
|
1. **Special tokens**: ``<|THINK_START|>`` and ``<|THINK_END|>``
|
|
(single-token IDs, e.g. 151669 and 151670).
|
|
|
|
2. **Text markers**: ```` and ``````
|
|
(multi-token sequences, e.g. ``</think`` → [522, 26865],
|
|
``>`` → [29]).
|
|
|
|
The chat template and typical model output use the **text form**
|
|
(multi-token). The parser therefore always uses **text-based**
|
|
matching for streaming extraction, but also checks single-token
|
|
IDs when available for methods that only receive token IDs
|
|
(``is_reasoning_end``, ``extract_content_ids``, etc.).
|
|
"""
|
|
|
|
# Text-form think markers (used in chat template & typical model output).
|
|
THINK_START_TEXT = "<think>"
|
|
THINK_END_TEXT = "</think>"
|
|
|
|
# Special-token-form think markers (single-token IDs in vocab).
|
|
THINK_START_SPECIAL = "<|THINK_START|>"
|
|
THINK_END_SPECIAL = "<|THINK_END|>"
|
|
|
|
def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
|
|
super().__init__(tokenizer, *args, **kwargs)
|
|
|
|
if not self.model_tokenizer:
|
|
raise ValueError(
|
|
"The model tokenizer must be passed to the ReasoningParser constructor during construction."
|
|
)
|
|
|
|
# --- Resolve think markers ----------------------------------------
|
|
# Special-token form (single tokens in vocab for fast ID checks).
|
|
self.think_start_special_id: int = self.vocab.get(self.THINK_START_SPECIAL, -1)
|
|
self.think_end_special_id: int = self.vocab.get(self.THINK_END_SPECIAL, -1)
|
|
|
|
# Text-form markers in vocab (usually NOT single tokens for
|
|
# Step-Audio, but checked for completeness).
|
|
self.think_start_text_id: int = self.vocab.get(self.THINK_START_TEXT, -1)
|
|
self.think_end_text_id: int = self.vocab.get(self.THINK_END_TEXT, -1)
|
|
|
|
# Set canonical token IDs for the ReasoningParser base contract.
|
|
# Prefer special-token IDs; fall back to text-form IDs.
|
|
self.think_start_token_id: int = (
|
|
self.think_start_special_id if self.think_start_special_id != -1 else self.think_start_text_id
|
|
)
|
|
self.think_end_token_id: int = (
|
|
self.think_end_special_id if self.think_end_special_id != -1 else self.think_end_text_id
|
|
)
|
|
|
|
# Whether single-token ID matching is possible for start/end.
|
|
self._has_start_token_id = self.think_start_token_id != -1
|
|
self._has_end_token_id = self.think_end_token_id != -1
|
|
|
|
# For text-based matching we must recognise BOTH forms.
|
|
# self.think_start_token / self.think_end_token are used by
|
|
# some helper methods and for the text-based path.
|
|
self.think_start_token = self.THINK_START_TEXT
|
|
self.think_end_token = self.THINK_END_TEXT
|
|
|
|
# Tracks whether we have already emitted the end of reasoning
|
|
# in streaming mode.
|
|
self._reasoning_ended = False
|
|
|
|
# Tracks whether reasoning JUST ended (within the last call),
|
|
# so we can strip a leading newline from the next content delta.
|
|
self._just_ended = False
|
|
|
|
# Buffer for text that might be the start of a multi-token
|
|
# think marker (e.g. "</" could be the beginning of "</think>").
|
|
# We hold this text until the next delta disambiguates it.
|
|
self._pending = ""
|
|
|
|
# Complete think marker strings used for ambiguous-prefix detection.
|
|
# When streaming, the tail of delta_text might be the start of one
|
|
# of these markers (e.g. "</" could be the beginning of ATHINK_END_AT).
|
|
# We buffer such ambiguous text to avoid leaking partial markers.
|
|
self._all_markers = (
|
|
self.THINK_START_TEXT,
|
|
self.THINK_END_TEXT,
|
|
self.THINK_START_SPECIAL,
|
|
self.THINK_END_SPECIAL,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Token-ID helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _has_end_token_in_ids(self, token_ids: Sequence[int]) -> bool:
|
|
"""Check for end marker in token IDs (single-token forms only)."""
|
|
if self.think_end_special_id != -1 and self.think_end_special_id in token_ids:
|
|
return True
|
|
if self.think_end_text_id != -1 and self.think_end_text_id in token_ids:
|
|
return True
|
|
return False
|
|
|
|
def _has_start_token_in_ids(self, token_ids: Sequence[int]) -> bool:
|
|
"""Check for start marker in token IDs (single-token forms only)."""
|
|
if self.think_start_special_id != -1 and self.think_start_special_id in token_ids:
|
|
return True
|
|
if self.think_start_text_id != -1 and self.think_start_text_id in token_ids:
|
|
return True
|
|
return False
|
|
|
|
# ------------------------------------------------------------------
|
|
# Text-based helpers (handles both text-form and special-form markers)
|
|
# ------------------------------------------------------------------
|
|
|
|
def _has_end_token_in_text(self, text: str) -> bool:
|
|
return self.THINK_END_TEXT in text or self.THINK_END_SPECIAL in text
|
|
|
|
def _has_start_token_in_text(self, text: str) -> bool:
|
|
return self.THINK_START_TEXT in text or self.THINK_START_SPECIAL in text
|
|
|
|
def _find_end_token_in_text(self, text: str) -> tuple[int, int]:
|
|
"""Find the first end token in text. Returns (start, end) indices
|
|
such that text[start:end] is the end token. Returns (-1, -1) if
|
|
not found."""
|
|
idx = text.find(self.THINK_END_TEXT)
|
|
if idx != -1:
|
|
return idx, idx + len(self.THINK_END_TEXT)
|
|
idx = text.find(self.THINK_END_SPECIAL)
|
|
if idx != -1:
|
|
return idx, idx + len(self.THINK_END_SPECIAL)
|
|
return -1, -1
|
|
|
|
def _find_start_token_in_text(self, text: str) -> tuple[int, int]:
|
|
"""Find the first start token in text. Returns (start, end) indices
|
|
such that text[start:end] is the start token. Returns (-1, -1) if
|
|
not found."""
|
|
idx = text.find(self.THINK_START_TEXT)
|
|
if idx != -1:
|
|
return idx, idx + len(self.THINK_START_TEXT)
|
|
idx = text.find(self.THINK_START_SPECIAL)
|
|
if idx != -1:
|
|
return idx, idx + len(self.THINK_START_SPECIAL)
|
|
return -1, -1
|
|
|
|
def _decode_ids(self, token_ids: Sequence[int]) -> str:
|
|
"""Decode a sequence of token IDs to text using the model tokenizer."""
|
|
return self.model_tokenizer.decode(token_ids, skip_special_tokens=False)
|
|
|
|
# ------------------------------------------------------------------
|
|
# ReasoningParser interface
|
|
# ------------------------------------------------------------------
|
|
|
|
def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
|
|
"""Check if reasoning has ended in the given token sequence.
|
|
|
|
When called with **prompt** token IDs (by the serving layer), the
|
|
prompt may contain think markers from previous assistant turns.
|
|
In multi-turn conversations the prompt can include both start
|
|
and end markers, with the *last* marker being a start marker
|
|
(from the generation prompt). In that case reasoning is NOT
|
|
ended — the model is about to generate inside a new think block.
|
|
|
|
We therefore find the **last** think marker (start or end) in
|
|
the decoded text and only return True if it is an end marker.
|
|
"""
|
|
try:
|
|
text = self._decode_ids(input_ids)
|
|
except Exception:
|
|
# Fall back to token-ID-only check.
|
|
result = self._has_end_token_in_ids(input_ids)
|
|
logger.debug(
|
|
"StepAudio is_reasoning_end: decode failed, fallback result=%s, num_ids=%d",
|
|
result,
|
|
len(input_ids),
|
|
)
|
|
return result
|
|
|
|
# Find the position of the LAST think marker (start or end).
|
|
last_start = max(
|
|
text.rfind(self.THINK_START_TEXT),
|
|
text.rfind(self.THINK_START_SPECIAL),
|
|
)
|
|
last_end = max(
|
|
text.rfind(self.THINK_END_TEXT),
|
|
text.rfind(self.THINK_END_SPECIAL),
|
|
)
|
|
|
|
if last_end == -1:
|
|
# No end marker at all → reasoning not ended.
|
|
logger.debug(
|
|
"StepAudio is_reasoning_end: no end marker found, text_tail=%r (len=%d), returning False",
|
|
text[-100:] if len(text) > 100 else text,
|
|
len(text),
|
|
)
|
|
return False
|
|
|
|
if last_start > last_end:
|
|
# Last marker is a START marker (e.g. the generation prompt's
|
|
# ``). Reasoning is still active.
|
|
logger.debug(
|
|
"StepAudio is_reasoning_end: last_start=%d > last_end=%d, reasoning still active, returning False",
|
|
last_start,
|
|
last_end,
|
|
)
|
|
return False
|
|
|
|
# Last marker is an END marker → reasoning has ended.
|
|
logger.debug(
|
|
"StepAudio is_reasoning_end: last_end=%d >= last_start=%d, reasoning ended, returning True",
|
|
last_end,
|
|
last_start,
|
|
)
|
|
return True
|
|
|
|
def is_reasoning_end_streaming(self, input_ids: Sequence[int], delta_ids: Iterable[int]) -> bool:
|
|
delta_list = tuple(delta_ids)
|
|
# Fast check: single-token markers in delta.
|
|
if self._has_end_token_in_ids(delta_list):
|
|
return True
|
|
# Fallback: decode delta + tail of input_ids and check text.
|
|
# The end marker may span the boundary between previous tokens
|
|
# and new delta tokens when it is encoded as multiple tokens
|
|
# (e.g. "</think" = [522, 26865]).
|
|
try:
|
|
delta_text = self._decode_ids(delta_list)
|
|
if self._has_end_token_in_text(delta_text):
|
|
return True
|
|
overlap_ids = list(input_ids[-8:]) + list(delta_list)
|
|
overlap_text = self._decode_ids(overlap_ids)
|
|
return self._has_end_token_in_text(overlap_text)
|
|
except Exception:
|
|
return False
|
|
|
|
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
|
|
# Fast check: single end-token in IDs.
|
|
if self._has_end_token_id:
|
|
# Check for special-token or text-form single token.
|
|
for end_id in (self.think_end_special_id, self.think_end_text_id):
|
|
if end_id != -1 and end_id in islice(input_ids, 0, max(0, len(input_ids) - 1)):
|
|
idx = input_ids.index(end_id)
|
|
return input_ids[idx + 1 :]
|
|
|
|
# Fallback: decode to text and find multi-token end marker.
|
|
try:
|
|
text = self._decode_ids(input_ids)
|
|
end_pos, end_len = self._find_end_token_in_text(text)
|
|
if end_pos == -1:
|
|
return []
|
|
content_start = end_len
|
|
if content_start < len(text) and text[content_start] == "\n":
|
|
content_start += 1
|
|
if content_start >= len(text):
|
|
return []
|
|
for i in range(len(input_ids)):
|
|
prefix_text = self._decode_ids(input_ids[: i + 1])
|
|
if len(prefix_text) > content_start:
|
|
return input_ids[i:]
|
|
return []
|
|
except Exception:
|
|
return []
|
|
|
|
def extract_reasoning(
|
|
self,
|
|
model_output: str,
|
|
request: "ChatCompletionRequest | ResponsesRequest",
|
|
) -> tuple[str | None, str | None]:
|
|
logger.debug(
|
|
"StepAudio extract_reasoning: model_output=%r (len=%d)",
|
|
model_output[:200],
|
|
len(model_output),
|
|
)
|
|
# Strip leading start token if present (text or special form).
|
|
for start_tok in (self.THINK_START_TEXT, self.THINK_START_SPECIAL):
|
|
if model_output.startswith(start_tok):
|
|
model_output = model_output[len(start_tok) :]
|
|
logger.debug("StepAudio extract_reasoning: stripped start_tok=%r", start_tok)
|
|
break
|
|
|
|
# Find the end token (text or special form).
|
|
end_pos, end_len = self._find_end_token_in_text(model_output)
|
|
if end_pos == -1:
|
|
# No end token found — everything is reasoning.
|
|
logger.debug(
|
|
"StepAudio extract_reasoning: no end token, reasoning=%r (len=%d), content=None",
|
|
model_output[:200],
|
|
len(model_output),
|
|
)
|
|
return model_output or None, None
|
|
|
|
reasoning = model_output[:end_pos]
|
|
content = model_output[end_len:]
|
|
# Strip leading newline that models often emit right after end token.
|
|
content = content.lstrip("\n")
|
|
logger.debug(
|
|
"StepAudio extract_reasoning: end_pos=%d, reasoning=%r (len=%d), content=%r (len=%d)",
|
|
end_pos,
|
|
reasoning[:200],
|
|
len(reasoning),
|
|
content[:200],
|
|
len(content),
|
|
)
|
|
return reasoning or None, content or None
|
|
|
|
def extract_reasoning_streaming(
|
|
self,
|
|
previous_text: str,
|
|
current_text: str,
|
|
delta_text: str,
|
|
previous_token_ids: Sequence[int],
|
|
current_token_ids: Sequence[int],
|
|
delta_token_ids: Sequence[int],
|
|
) -> DeltaMessage | None:
|
|
# Step-Audio typically uses multi-token text-form markers
|
|
# (e.g. "</think" → [522, 26865] + ">" → [29]).
|
|
# We always use the text-based path because it handles both
|
|
# single-token special markers and multi-token text markers.
|
|
# The text-based path is also efficient since previous_text
|
|
# and delta_text are already available.
|
|
return self._extract_streaming_text(
|
|
previous_text,
|
|
delta_text,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Streaming: text-based (handles both text and special token forms)
|
|
# ------------------------------------------------------------------
|
|
|
|
def _ends_with_marker_prefix(self, text: str) -> str:
|
|
"""If text ends with a prefix of any think marker, return the
|
|
ambiguous suffix that should be buffered.
|
|
|
|
For example:
|
|
- text="reasoning</" → marker ATHINK_END_AT starts with "</" → return "</"
|
|
- text="reasoning<" → markers start with "<" → return "<"
|
|
- text="reasoning<th" → ATHINK_START starts with "<th" → return "<th"
|
|
(but "</think" also starts with "<", so we buffer from the longest)
|
|
- text="reasoning</thi" → ATHINK_END_AT starts with "</thi" → return "</thi"
|
|
"""
|
|
if not text:
|
|
return ""
|
|
|
|
# Find the longest suffix of text that is a prefix of any marker.
|
|
best = ""
|
|
for marker in self._all_markers:
|
|
for i in range(1, len(marker) + 1):
|
|
prefix = marker[:i]
|
|
if text.endswith(prefix) and len(prefix) > len(best):
|
|
best = prefix
|
|
return best
|
|
|
|
# Counter for debug logging (to avoid flooding logs)
|
|
_stream_call_count: int = 0
|
|
|
|
def _extract_streaming_text(
|
|
self,
|
|
previous_text: str,
|
|
delta_text: str,
|
|
) -> DeltaMessage | None:
|
|
# Reset _just_ended flag at the start of each call.
|
|
# It will be set to True if this call ends reasoning.
|
|
just_ended_before = self._just_ended
|
|
self._just_ended = False
|
|
|
|
# Debug logging for first few calls
|
|
StepAudioReasoningParser._stream_call_count += 1
|
|
call_num = StepAudioReasoningParser._stream_call_count
|
|
if call_num <= 20:
|
|
logger.debug(
|
|
"StepAudio stream #%d: prev_text=%r (len=%d), delta_text=%r (len=%d), _reasoning_ended=%s, _pending=%r",
|
|
call_num,
|
|
previous_text[:100],
|
|
len(previous_text),
|
|
delta_text[:100],
|
|
len(delta_text),
|
|
self._reasoning_ended,
|
|
self._pending[:50] if self._pending else "",
|
|
)
|
|
|
|
# Prepend any pending text from a previous call.
|
|
# The pending text is already part of previous_text (the framework
|
|
# accumulates all generated text), so we must also trim it from
|
|
# previous_text to avoid double-counting when we form `combined`.
|
|
if self._pending:
|
|
n = len(self._pending)
|
|
if previous_text.endswith(self._pending):
|
|
previous_text = previous_text[:-n]
|
|
delta_text = self._pending + delta_text
|
|
self._pending = ""
|
|
|
|
combined = previous_text + delta_text
|
|
has_end = self._has_end_token_in_text(combined)
|
|
has_end_in_previous = self._has_end_token_in_text(previous_text)
|
|
has_end_in_delta = self._has_end_token_in_text(delta_text)
|
|
|
|
if call_num <= 20:
|
|
logger.debug(
|
|
"StepAudio stream #%d: has_end=%s, has_end_in_previous=%s, has_end_in_delta=%s, combined=%r (len=%d)",
|
|
call_num,
|
|
has_end,
|
|
has_end_in_previous,
|
|
has_end_in_delta,
|
|
combined[:150],
|
|
len(combined),
|
|
)
|
|
|
|
if has_end_in_previous or self._reasoning_ended:
|
|
# Already past reasoning — everything is content.
|
|
self._reasoning_ended = True
|
|
# Strip at most one leading newline right after the end marker.
|
|
if delta_text.startswith("\n") and just_ended_before:
|
|
return DeltaMessage(content=delta_text[1:] or None)
|
|
return DeltaMessage(content=delta_text)
|
|
|
|
if has_end_in_delta:
|
|
# End token fully contained in the delta — split reasoning
|
|
# and content.
|
|
end_index, end_len = self._find_end_token_in_text(delta_text)
|
|
reasoning = delta_text[:end_index]
|
|
content = delta_text[end_len:]
|
|
# Strip leading newline after end token.
|
|
if content.startswith("\n"):
|
|
content = content[1:]
|
|
self._reasoning_ended = True
|
|
self._just_ended = True
|
|
return DeltaMessage(
|
|
reasoning=reasoning or None,
|
|
content=content or None,
|
|
)
|
|
|
|
if has_end and not has_end_in_previous and not has_end_in_delta:
|
|
# End token spans the boundary between previous_text and
|
|
# delta_text (e.g. previous ends with "</" and delta
|
|
# starts with "think>..."). Find the split point in the
|
|
# combined text.
|
|
end_index, end_len = self._find_end_token_in_text(combined)
|
|
marker_start_in_prev = end_index - len(previous_text)
|
|
if marker_start_in_prev > 0:
|
|
# Some of previous_text's trailing chars are part of
|
|
# the marker — they were already emitted as reasoning
|
|
# in a previous call. The delta begins the content
|
|
# after the marker.
|
|
content_after = combined[end_len:]
|
|
if content_after.startswith("\n"):
|
|
content_after = content_after[1:]
|
|
self._reasoning_ended = True
|
|
self._just_ended = True
|
|
return DeltaMessage(content=content_after or None)
|
|
else:
|
|
# Edge case: the entire marker is in previous_text
|
|
# but was not detected because of how text was split
|
|
# across streaming chunks.
|
|
self._reasoning_ended = True
|
|
self._just_ended = True
|
|
full_content = combined[end_len:]
|
|
if full_content.startswith("\n"):
|
|
full_content = full_content[1:]
|
|
return DeltaMessage(content=full_content or None)
|
|
|
|
# No end token seen yet — everything is reasoning.
|
|
# Skip the start token prefix if present in delta.
|
|
start_idx, start_len = self._find_start_token_in_text(delta_text)
|
|
if start_idx != -1 and delta_text.strip() == delta_text[start_idx:start_len]:
|
|
# Delta is only the start token — skip it entirely.
|
|
return None
|
|
if start_idx == 0:
|
|
# Delta starts with the start token — skip past it.
|
|
after_start = delta_text[start_len:]
|
|
return DeltaMessage(reasoning=after_start or None)
|
|
|
|
# Check if the delta ends with a partial marker prefix (e.g. "</"
|
|
# which could be the beginning of the end marker). Buffer the
|
|
# ambiguous text to avoid leaking partial markers into reasoning.
|
|
ambiguous_suffix = self._ends_with_marker_prefix(delta_text)
|
|
if ambiguous_suffix:
|
|
if ambiguous_suffix == delta_text:
|
|
# Entire delta is a possible marker prefix — buffer it all.
|
|
self._pending = delta_text
|
|
return None
|
|
else:
|
|
# Emit the non-ambiguous part as reasoning, buffer the rest.
|
|
to_emit = delta_text[: -len(ambiguous_suffix)]
|
|
self._pending = ambiguous_suffix
|
|
return DeltaMessage(reasoning=to_emit) if to_emit else None
|
|
|
|
return DeltaMessage(reasoning=delta_text)
|
|
|
|
def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
|
|
"""Count tokens within thinking spans."""
|
|
# Fast path: if single-token markers are present, count using IDs.
|
|
if self._has_start_token_id and self._has_end_token_id:
|
|
count = 0
|
|
depth = 0
|
|
for token_id in token_ids:
|
|
if token_id == self.think_start_token_id:
|
|
depth += 1
|
|
continue
|
|
if token_id == self.think_end_token_id:
|
|
if depth > 0:
|
|
depth -= 1
|
|
continue
|
|
if depth > 0:
|
|
count += 1
|
|
# Only trust the fast-path result if we actually found markers.
|
|
if depth == 0 and count > 0:
|
|
return count
|
|
|
|
# Text-based fallback: handles both text-form multi-token markers
|
|
# and special-token single-token markers in decoded text.
|
|
try:
|
|
text = self._decode_ids(token_ids)
|
|
if not text:
|
|
return 0
|
|
n_tokens = len(token_ids)
|
|
if n_tokens == 0:
|
|
return 0
|
|
|
|
total_chars = len(text)
|
|
if total_chars == 0:
|
|
return 0
|
|
chars_per_token = total_chars / n_tokens
|
|
|
|
count = 0
|
|
depth = 0
|
|
i = 0
|
|
while i < len(text):
|
|
found_start = False
|
|
for start_tok in (self.THINK_START_TEXT, self.THINK_START_SPECIAL):
|
|
if text[i:].startswith(start_tok):
|
|
depth += 1
|
|
i += len(start_tok)
|
|
found_start = True
|
|
break
|
|
if found_start:
|
|
continue
|
|
|
|
found_end = False
|
|
for end_tok in (self.THINK_END_TEXT, self.THINK_END_SPECIAL):
|
|
if text[i:].startswith(end_tok):
|
|
if depth > 0:
|
|
depth -= 1
|
|
i += len(end_tok)
|
|
found_end = True
|
|
break
|
|
if found_end:
|
|
continue
|
|
|
|
if depth > 0:
|
|
count += 1
|
|
i += 1
|
|
|
|
# Scale per-character count to approximate token count.
|
|
if chars_per_token > 1:
|
|
count = max(1, round(count / chars_per_token))
|
|
|
|
return count
|
|
except Exception:
|
|
return 0
|