From 57ea047acf1256322c508d1c791da5089443d43e Mon Sep 17 00:00:00 2001 From: Mike Qiu Date: Thu, 11 Jun 2026 05:03:19 +0800 Subject: [PATCH] Step audio R1 reasoning parser (#2846) Signed-off-by: Mike_Qiu Signed-off-by: Yueqian Lin Signed-off-by: Mike Qiu Co-authored-by: Mike_Qiu Co-authored-by: Claude Opus 4.6 Co-authored-by: Yueqian Lin Co-authored-by: Hongsheng Liu --- pyproject.toml | 4 + tests/reasoning/__init__.py | 0 .../test_step_audio_reasoning_parser.py | 918 ++++++++++++++++++ vllm_omni/engine/arg_utils.py | 3 + vllm_omni/engine/stage_engine_core_proc.py | 14 + vllm_omni/entrypoints/openai/api_server.py | 9 + vllm_omni/reasoning/__init__.py | 24 + .../reasoning/step_audio_reasoning_parser.py | 574 +++++++++++ 8 files changed, 1546 insertions(+) create mode 100644 tests/reasoning/__init__.py create mode 100644 tests/reasoning/test_step_audio_reasoning_parser.py create mode 100644 vllm_omni/reasoning/__init__.py create mode 100644 vllm_omni/reasoning/step_audio_reasoning_parser.py diff --git a/pyproject.toml b/pyproject.toml index 09b8e23ca..757f038af 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -292,5 +292,9 @@ fullset = "fullset" Vai = "Vai" tockens = "tockens" CANN = "CANN" +# Step-Audio reasoning parser streams `` tokens split across chunks +# (e.g. ""). The fragment "thi" appears literally in its +# parser logic and tests and is not a typo of "the"/"this". +thi = "thi" ND = "ND" byt5 = "byt5" diff --git a/tests/reasoning/__init__.py b/tests/reasoning/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/reasoning/test_step_audio_reasoning_parser.py b/tests/reasoning/test_step_audio_reasoning_parser.py new file mode 100644 index 000000000..7eb9fa17a --- /dev/null +++ b/tests/reasoning/test_step_audio_reasoning_parser.py @@ -0,0 +1,918 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Tests for StepAudioReasoningParser. + +Step-Audio supports two representations of thinking markers: + +1. **Special tokens**: ``<|THINK_START|>`` and ``<|THINK_END|>`` + (single-token IDs in vocab). + +2. **Text markers**: ```` and ```` + (multi-token sequences for , e.g. [522, 26865]). + +These tests use standalone mock objects so they can run without +installing the full vLLM/torch stack. +""" + +import sys +from types import ModuleType +from unittest.mock import MagicMock + +import pytest + +# --------------------------------------------------------------------------- +# Bootstrap minimal vLLM stubs so the parser can be imported +# --------------------------------------------------------------------------- + + +def _install_vllm_stubs() -> dict[str, ModuleType | None]: + """Temporarily install minimal stub modules for the vllm imports the + parser needs, so it can be imported without the full vLLM/torch stack. + + Returns a snapshot of the original ``sys.modules`` entries that were + replaced. The caller MUST pass it to ``_restore_vllm_stubs`` once the + parser is imported. These stubs are process-global: if they leak, they + shadow the real vllm modules for every other test in the same worker + (and any forked server subprocess), e.g. replacing + ``vllm.entrypoints.openai.engine.protocol`` with a stub that lacks + ``ErrorResponse`` and breaking unrelated server-startup tests. + """ + stubs: dict[str, ModuleType] = {} + + # DeltaMessage stub + delta_mod = ModuleType("vllm.entrypoints.openai.engine.protocol") + + class DeltaMessage: + __slots__ = ("reasoning", "content") + + def __init__(self, reasoning=None, content=None): + self.reasoning = reasoning + self.content = content + + delta_mod.DeltaMessage = DeltaMessage + stubs["vllm.entrypoints.openai.engine.protocol"] = delta_mod + + # ReasoningParser stub + reasoning_mod = ModuleType("vllm.reasoning") + + class ReasoningParser: + def __init__(self, tokenizer, *args, **kwargs): + self.model_tokenizer = tokenizer + + @property + def vocab(self): + if not self.model_tokenizer: + return {} + return self.model_tokenizer.get_vocab() + + class _ReasoningParserManager: + reasoning_parsers = {} + lazy_parsers = {} + + @classmethod + def register_lazy_module(cls, name, module_path, class_name): + cls.lazy_parsers[name] = (module_path, class_name) + + reasoning_mod.ReasoningParser = ReasoningParser + reasoning_mod.ReasoningParserManager = _ReasoningParserManager + stubs["vllm.reasoning"] = reasoning_mod + stubs["vllm.reasoning.abs_reasoning_parsers"] = reasoning_mod + + # TokenizerLike stub + tokenizers_mod = ModuleType("vllm.tokenizers") + tokenizers_mod.TokenizerLike = object + stubs["vllm.tokenizers"] = tokenizers_mod + + # ChatCompletionRequest / ResponsesRequest stubs + chat_mod = ModuleType("vllm.entrypoints.openai.chat_completion.protocol") + chat_mod.ChatCompletionRequest = MagicMock + stubs["vllm.entrypoints.openai.chat_completion.protocol"] = chat_mod + responses_mod = ModuleType("vllm.entrypoints.openai.responses.protocol") + responses_mod.ResponsesRequest = MagicMock + stubs["vllm.entrypoints.openai.responses.protocol"] = responses_mod + + # Parent packages: only stub the ones not already importable, so we never + # clobber a real (installed) vllm package tree. + for mod_path in [ + "vllm", + "vllm.entrypoints", + "vllm.entrypoints.openai", + "vllm.entrypoints.openai.engine", + "vllm.entrypoints.openai.chat_completion", + "vllm.entrypoints.openai.responses", + ]: + if mod_path not in sys.modules and mod_path not in stubs: + stubs[mod_path] = ModuleType(mod_path) + + saved = {name: sys.modules.get(name) for name in stubs} + sys.modules.update(stubs) + return saved + + +def _restore_vllm_stubs(saved: dict[str, ModuleType | None]) -> None: + """Undo ``_install_vllm_stubs`` so ``sys.modules`` is left untouched.""" + for name, original in saved.items(): + if original is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = original + + +# Install stubs only for the duration of the import. The parser binds these +# stub classes into its own module namespace at import time, so it keeps using +# them at runtime even after sys.modules is restored below. +_saved_modules = _install_vllm_stubs() +try: + from vllm_omni.reasoning.step_audio_reasoning_parser import ( # noqa: E402 + StepAudioReasoningParser, + ) +finally: + _restore_vllm_stubs(_saved_modules) + +# Short aliases for marker constants +TS = StepAudioReasoningParser.THINK_START_TEXT # +TE = StepAudioReasoningParser.THINK_END_TEXT # +TSS = StepAudioReasoningParser.THINK_START_SPECIAL # <|THINK_START|> +TES = StepAudioReasoningParser.THINK_END_SPECIAL # <|THINK_END|> + + +# --------------------------------------------------------------------------- +# Mock tokenizers +# --------------------------------------------------------------------------- + + +class SingleTokenTokenizer: + """Tokenizer where both ```` and ```` are single tokens. + + This exercises the fast (token-ID) path. + Also includes special-token forms ``<|THINK_START|>`` and ``<|THINK_END|>``. + """ + + def __init__(self): + self._vocab = { + TS: 100, # + TE: 101, # + TSS: 151669, # <|THINK_START|> + TES: 151670, # <|THINK_END|> + "Hello": 1, + " world": 2, + "reasoning": 3, + " content": 4, + "\n": 5, + "": 7, + "some text": 8, + } + self._reverse = {v: k for k, v in self._vocab.items()} + + def get_vocab(self): + return dict(self._vocab) + + def decode(self, ids, skip_special_tokens=False): + return "".join(self._reverse.get(i, "") for i in ids) + + def tokenize(self, text): + if not text: + return [] + return [text] + + +class MultiTokenTokenizer: + """Tokenizer where ```` is NOT a single token. + + ```` is in the vocab (ID 100), but ```` is NOT. + Instead, ```` encodes as [""] → [522, 26865]. + + Both special tokens ``<|THINK_START|>`` and ``<|THINK_END|>`` ARE + present as single tokens. + + This exercises the text-based path (with special-token fallback). + """ + + def __init__(self): + self._vocab = { + TS: 100, # - single token + # is intentionally NOT in the vocab + "": 26865, + TSS: 151669, # <|THINK_START|> + TES: 151670, # <|THINK_END|> + "Hello": 1, + " world": 2, + "reasoning": 3, + " content": 4, + "\n": 5, + "some text": 8, + } + self._reverse = {v: k for k, v in self._vocab.items()} + + def get_vocab(self): + return dict(self._vocab) + + def decode(self, ids, skip_special_tokens=False): + return "".join(self._reverse.get(i, "") for i in ids) + + def tokenize(self, text): + if not text: + return [] + return [text] + + +class TextOnlyTokenizer: + """Tokenizer with only text-form markers (no special tokens at all). + + Both ```` and ```` are single tokens. + ``<|THINK_START|>`` and ``<|THINK_END|>`` are NOT in the vocab. + + This tests the pure text-based fallback paths. + """ + + def __init__(self): + self._vocab = { + TS: 100, # + TE: 101, # + "Hello": 1, + " world": 2, + "reasoning": 3, + " content": 4, + "\n": 5, + "": 7, + } + self._reverse = {v: k for k, v in self._vocab.items()} + + def get_vocab(self): + return dict(self._vocab) + + def decode(self, ids, skip_special_tokens=False): + return "".join(self._reverse.get(i, "") for i in ids) + + def tokenize(self, text): + if not text: + return [] + return [text] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _stream_extract(parser, deltas_and_ids): + """Run streaming extraction over a list of (delta_text, delta_ids) tuples. + Returns (reasoning_text, content_text).""" + prev_text = "" + prev_ids = [] + all_r = [] + all_c = [] + + for delta_text, delta_ids in deltas_and_ids: + cur_text = prev_text + delta_text + cur_ids = prev_ids + delta_ids + result = parser.extract_reasoning_streaming(prev_text, cur_text, delta_text, prev_ids, cur_ids, delta_ids) + if result is not None: + if result.reasoning: + all_r.append(result.reasoning) + if result.content: + all_c.append(result.content) + prev_text = cur_text + prev_ids = cur_ids + + return "".join(all_r), "".join(all_c) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def single_token_parser(): + return StepAudioReasoningParser(SingleTokenTokenizer()) + + +@pytest.fixture +def multi_token_parser(): + return StepAudioReasoningParser(MultiTokenTokenizer()) + + +@pytest.fixture +def text_only_parser(): + return StepAudioReasoningParser(TextOnlyTokenizer()) + + +@pytest.fixture +def request_obj(): + return MagicMock(messages=[], model="test-model") + + +# --------------------------------------------------------------------------- +# Initialization tests +# --------------------------------------------------------------------------- + + +class TestInit: + def test_single_token_ids(self, single_token_parser): + """Text-form markers are in vocab as single tokens.""" + assert single_token_parser.think_start_text_id == 100 + assert single_token_parser.think_end_text_id == 101 + + def test_single_token_special_ids(self, single_token_parser): + """Special-token markers are also in vocab.""" + assert single_token_parser.think_start_special_id == 151669 + assert single_token_parser.think_end_special_id == 151670 + + def test_single_token_has_ids(self, single_token_parser): + assert single_token_parser._has_start_token_id is True + assert single_token_parser._has_end_token_id is True + + def test_multi_token_text_ids(self, multi_token_parser): + """ is a single token; is NOT.""" + assert multi_token_parser.think_start_text_id == 100 + assert multi_token_parser.think_end_text_id == -1 + + def test_multi_token_special_ids(self, multi_token_parser): + """Special tokens are still in vocab.""" + assert multi_token_parser.think_start_special_id == 151669 + assert multi_token_parser.think_end_special_id == 151670 + + def test_multi_token_has_ids(self, multi_token_parser): + """Even without as single token, special tokens provide IDs.""" + assert multi_token_parser._has_start_token_id is True + assert multi_token_parser._has_end_token_id is True + + def test_text_only_ids(self, text_only_parser): + """Only text-form markers, no special tokens.""" + assert text_only_parser.think_start_text_id == 100 + assert text_only_parser.think_end_text_id == 101 + assert text_only_parser.think_start_special_id == -1 + assert text_only_parser.think_end_special_id == -1 + assert text_only_parser._has_start_token_id is True + assert text_only_parser._has_end_token_id is True + + def test_missing_tokenizer_raises(self): + with pytest.raises(ValueError, match="model tokenizer must be passed"): + StepAudioReasoningParser(None) + + +# --------------------------------------------------------------------------- +# extract_reasoning (non-streaming) +# --------------------------------------------------------------------------- + + +class TestExtractReasoning: + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_basic_extraction(self, request, parser_name, request_obj): + parser = request.getfixturevalue(parser_name) + output = f"some reasoning{TE}Hello world" + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning == "some reasoning" + assert content == "Hello world" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_no_end_token(self, request, parser_name, request_obj): + parser = request.getfixturevalue(parser_name) + output = "still thinking..." + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning == "still thinking..." + assert content is None + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_with_start_token_prefix(self, request, parser_name, request_obj): + parser = request.getfixturevalue(parser_name) + output = f"{TS}my reasoning{TE}answer" + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning == "my reasoning" + assert content == "answer" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_strips_newline_after_end(self, request, parser_name, request_obj): + parser = request.getfixturevalue(parser_name) + output = f"reasoning{TE}\nHello" + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning == "reasoning" + assert content == "Hello" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_empty_reasoning(self, request, parser_name, request_obj): + parser = request.getfixturevalue(parser_name) + output = f"{TE}Hello" + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning is None + assert content == "Hello" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + ], + ) + def test_special_end_token(self, request, parser_name, request_obj): + """Test with <|THINK_END|> special token form.""" + parser = request.getfixturevalue(parser_name) + output = f"reasoning{TES}answer" + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning == "reasoning" + assert content == "answer" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + ], + ) + def test_special_start_token(self, request, parser_name, request_obj): + """Test with <|THINK_START|> special token form.""" + parser = request.getfixturevalue(parser_name) + output = f"{TSS}reasoning{TE}answer" + reasoning, content = parser.extract_reasoning(output, request_obj) + assert reasoning == "reasoning" + assert content == "answer" + + +# --------------------------------------------------------------------------- +# is_reasoning_end +# --------------------------------------------------------------------------- + + +class TestIsReasoningEnd: + def test_single_token_text_end(self, single_token_parser): + assert single_token_parser.is_reasoning_end([1, 2, 101]) is True + + def test_single_token_special_end(self, single_token_parser): + assert single_token_parser.is_reasoning_end([1, 2, 151670]) is True + + def test_single_token_no_end(self, single_token_parser): + assert single_token_parser.is_reasoning_end([1, 2, 3]) is False + + def test_multi_token_text_end_decoded(self, multi_token_parser): + """ = [522, 26865] detected via decode.""" + assert multi_token_parser.is_reasoning_end([1, 522, 26865]) is True + + def test_multi_token_special_end(self, multi_token_parser): + """<|THINK_END|> detected as single token ID.""" + assert multi_token_parser.is_reasoning_end([1, 151670]) is True + + def test_multi_token_no_end(self, multi_token_parser): + assert multi_token_parser.is_reasoning_end([1, 2, 3]) is False + + def test_multi_token_partial_marker(self, multi_token_parser): + """Only should NOT count.""" + assert multi_token_parser.is_reasoning_end([1, 522, 3]) is False + + def test_multi_token_empty_input(self, multi_token_parser): + assert multi_token_parser.is_reasoning_end([]) is False + + def test_text_only_end(self, text_only_parser): + assert text_only_parser.is_reasoning_end([1, 2, 101]) is True + + def test_text_only_no_end(self, text_only_parser): + assert text_only_parser.is_reasoning_end([1, 2, 3]) is False + + def test_start_after_end_returns_false(self, single_token_parser): + """Multi-turn prompt: start marker AFTER end marker → reasoning + is still active (the generation prompt's start marker).""" + # IDs: [start=100, end=101, start=100] → last marker is start + assert single_token_parser.is_reasoning_end([100, 101, 100]) is False + + def test_end_is_last_marker_returns_true(self, single_token_parser): + """End marker is the last marker → reasoning has ended.""" + # IDs: [start=100, end=101] → last marker is end + assert single_token_parser.is_reasoning_end([100, 101]) is True + + def test_start_only_no_end(self, single_token_parser): + """Only a start marker, no end marker → reasoning is active.""" + assert single_token_parser.is_reasoning_end([100, 1, 2]) is False + + def test_multi_turn_prompt_start_after_end(self, multi_token_parser): + """Simulated multi-turn prompt with previous assistant response. + + Prompt contains: ... " + " " + ", 151669=<|THINK_START|>] + ids = [151670, 151669] + assert multi_token_parser.is_reasoning_end(ids) is False + + +# --------------------------------------------------------------------------- +# is_reasoning_end_streaming +# --------------------------------------------------------------------------- + + +class TestIsReasoningEndStreaming: + def test_single_token_delta_has_end(self, single_token_parser): + assert single_token_parser.is_reasoning_end_streaming([1, 2], [101]) is True + + def test_single_token_delta_no_end(self, single_token_parser): + assert single_token_parser.is_reasoning_end_streaming([1, 2], [3]) is False + + def test_multi_token_delta_has_both_parts(self, multi_token_parser): + """Full in delta → [522, 26865].""" + assert multi_token_parser.is_reasoning_end_streaming([1], [522, 26865]) is True + + def test_multi_token_marker_spans_boundary(self, multi_token_parser): + """ in delta.""" + assert multi_token_parser.is_reasoning_end_streaming([1, 522], [26865]) is True + + def test_multi_token_no_end(self, multi_token_parser): + assert multi_token_parser.is_reasoning_end_streaming([1, 2], [3]) is False + + +# --------------------------------------------------------------------------- +# extract_reasoning_streaming +# --------------------------------------------------------------------------- + + +class TestExtractReasoningStreaming: + """Test streaming extraction for both single-token and multi-token modes.""" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_simple_streaming(self, request_, parser_name): + parser = request_.getfixturevalue(parser_name) + # Simulate: reasoning final + if parser_name == "multi_token_parser": + deltas = [TS, "reasoning", TE, " content"] + token_ids_per_delta = [ + [100], # + [3], # reasoning + [522, 26865], # + [4], # content + ] + else: + deltas = [TS, "reasoning", TE, " content"] + token_ids_per_delta = [ + [100], # + [3], # reasoning + [101], # + [4], # content + ] + + r, c = _stream_extract(parser, list(zip(deltas, token_ids_per_delta))) + assert r == "reasoning" + assert c == " content" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + "text_only_parser", + ], + ) + def test_no_end_token_streaming(self, request_, parser_name): + parser = request_.getfixturevalue(parser_name) + deltas_and_ids = [ + (TS, [100]), + ("thinking", [3]), + (" more", [2]), + ] + r, c = _stream_extract(parser, deltas_and_ids) + assert r == "thinking more" + assert c == "" + + def test_multi_token_end_split_across_chunks(self, multi_token_parser): + """The key test: is split across streaming chunks. + + Previous chunk ends with "reasoning\\nfinal answer". This is the real-world scenario when + ```` is tokenized as two separate tokens. + """ + deltas_and_ids = [ + ("reasoning", [3]), + ("\nHello", [26865, 1]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "reasoning" + assert c == "Hello" + + def test_multi_token_start_token_in_delta(self, multi_token_parser): + """Test that start token is skipped in streaming output.""" + deltas_and_ids = [ + (TS, [100]), + ("reasoning", [3]), + (TE, [522, 26865]), + ("Hello", [1]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "reasoning" + assert c == "Hello" + + def test_multi_token_three_way_split(self, multi_token_parser): + """ split across 3 chunks: .""" + deltas_and_ids = [ + ("some reasoning", [3]), + ("", [8]), + ("\nHello", [1]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "some reasoning" + assert c == "Hello" + + def test_special_end_token_in_stream(self, multi_token_parser): + """<|THINK_END|> special token in streaming (decoded as text).""" + deltas_and_ids = [ + ("reasoning", [3]), + (TES, [151670]), + ("answer", [1]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "reasoning" + assert c == "answer" + + def test_no_start_token_in_stream(self, multi_token_parser): + """Model starts directly with reasoning (no prefix).""" + deltas_and_ids = [ + ("reasoning", [3]), + (TE, [522, 26865]), + ("answer", [1]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "reasoning" + assert c == "answer" + + def test_empty_reasoning(self, single_token_parser): + """Start immediately followed by end → empty reasoning.""" + deltas_and_ids = [ + (TS, [100]), + (TE, [101]), + ("Hello", [1]), + ] + r, c = _stream_extract(single_token_parser, deltas_and_ids) + assert r in ("", None) + assert c == "Hello" + + def test_end_marker_no_content_after(self, multi_token_parser): + """End marker with no content following it.""" + deltas_and_ids = [ + ("reasoning", [3]), + (TE, [522, 26865]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "reasoning" + assert c == "" + + def test_delta_is_just_start_token(self, multi_token_parser): + """Delta consisting only of the start token should be skipped.""" + deltas_and_ids = [ + (TS, [100]), + ] + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "" + assert c == "" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + ], + ) + def test_only_first_newline_stripped_after_end(self, request_, parser_name): + """Only the first newline after should be stripped.""" + parser = request_.getfixturevalue(parser_name) + if parser_name == "multi_token_parser": + deltas_and_ids = [ + ("reasoning", [3]), + (TE, [522, 26865]), + ("\n\nHello", [5, 1]), + ] + else: + deltas_and_ids = [ + ("reasoning", [3]), + (TE, [101]), + ("\n\nHello", [5, 1]), + ] + r, c = _stream_extract(parser, deltas_and_ids) + assert r == "reasoning" + assert c == "\nHello" + + @pytest.mark.parametrize( + "parser_name", + [ + "single_token_parser", + "multi_token_parser", + ], + ) + def test_content_after_end_multiple_deltas(self, request_, parser_name): + """Multiple content deltas after end marker.""" + parser = request_.getfixturevalue(parser_name) + if parser_name == "multi_token_parser": + deltas_and_ids = [ + (TS, [100]), + ("reasoning", [3]), + (TE, [522, 26865]), + ("Hello", [1]), + (" world", [2]), + ] + else: + deltas_and_ids = [ + (TS, [100]), + ("reasoning", [3]), + (TE, [101]), + ("Hello", [1]), + (" world", [2]), + ] + r, c = _stream_extract(parser, deltas_and_ids) + assert r == "reasoning" + assert c == "Hello world" + + def test_pending_buffer_resolved_on_non_marker(self, multi_token_parser): + """When buffered text turns out not to be a marker prefix, it is + emitted as reasoning on the next call.""" + # "reason<" is buffered because "<" could start a marker. + # Next delta ">" doesn't form a marker, so the buffer is + # flushed as reasoning. + parser = multi_token_parser + deltas_and_ids = [ + (TS, [100]), + ("reason", [3]), + ("<", [8]), # "<" is buffered (prefix of markers) + ("> more", [8]), # ">" doesn't complete any marker + ] + r, c = _stream_extract(parser, deltas_and_ids) + assert r == "reason> more" + assert c == "" + + def test_multi_token_end_with_reasoning_before_buffer(self, multi_token_parser): + """Reasoning text before a buffered prefix should be emitted.""" + # "reasoning<" → "reasoning" emitted, "<" buffered + # "\nanswer" → buffer resolves as end marker + parser = multi_token_parser + deltas_and_ids = [ + ("reasoning", [3]), + ("<", [8]), + (f"/{TE[2:]}\nanswer", [522, 26865, 1]), + ] + r, c = _stream_extract(parser, deltas_and_ids) + assert r == "reasoning" + assert c == "answer" + + def test_multi_token_five_way_split(self, multi_token_parser): + """ split across 5 tiny chunks.""" + # Split: each chunk is one character of + chunks = list(TE) + deltas_and_ids = [ + ("reasoning", [3]), + ] + for ch in chunks: + deltas_and_ids.append((ch, [8])) + deltas_and_ids.append(("\nHello", [1])) + r, c = _stream_extract(multi_token_parser, deltas_and_ids) + assert r == "reasoning" + assert c == "Hello" + + def test_single_token_parser_with_pending(self, single_token_parser): + """Single-token parser: ambiguous '<' prefix buffered then resolved.""" + # Even with single-token markers, the text-based path buffers + # '<' because it could be a prefix of a marker. + parser = single_token_parser + deltas_and_ids = [ + (TS, [100]), + ("reason", [3]), + ("<", [8]), + ("> not a marker", [8]), + ] + r, c = _stream_extract(parser, deltas_and_ids) + assert r == "reason> not a marker" + assert c == "" + + +# --------------------------------------------------------------------------- + + +class TestCountReasoningTokens: + def test_single_token_basic(self, single_token_parser): + # IDs: [100=, 3=reasoning, 101=, 4=content] + assert single_token_parser.count_reasoning_tokens([100, 3, 101, 4]) == 1 + + def test_single_token_no_end(self, single_token_parser): + # All tokens after start are reasoning + assert single_token_parser.count_reasoning_tokens([100, 3, 2]) == 2 + + def test_multi_token_basic(self, multi_token_parser): + # encodes as [522, 26865] + # Full sequence: reasoning content + # IDs: [100, 3, 522, 26865, 4] + result = multi_token_parser.count_reasoning_tokens([100, 3, 522, 26865, 4]) + assert result > 0 + + def test_multi_token_no_end(self, multi_token_parser): + # No end token → all tokens inside are counted. + result = multi_token_parser.count_reasoning_tokens([100, 3, 2]) + assert result > 0 + + def test_text_only_basic(self, text_only_parser): + assert text_only_parser.count_reasoning_tokens([100, 3, 101, 4]) == 1 + + def test_text_only_no_end(self, text_only_parser): + assert text_only_parser.count_reasoning_tokens([100, 3, 2]) == 2 + + +# --------------------------------------------------------------------------- +# extract_content_ids +# --------------------------------------------------------------------------- + + +class TestExtractContentIds: + def test_single_token_basic(self, single_token_parser): + # [1, 101=, 2, 3] → content starts after ID 101 + result = single_token_parser.extract_content_ids([1, 101, 2, 3]) + assert result == [2, 3] + + def test_single_token_no_end(self, single_token_parser): + result = single_token_parser.extract_content_ids([1, 2, 3]) + assert result == [] + + def test_multi_token_basic(self, multi_token_parser): + # = [522, 26865] + # [3, 522, 26865, 4] → content after + result = multi_token_parser.extract_content_ids([3, 522, 26865, 4]) + assert 4 in result + + def test_multi_token_no_end(self, multi_token_parser): + result = multi_token_parser.extract_content_ids([1, 2, 3]) + assert result == [] + + def test_text_only_basic(self, text_only_parser): + result = text_only_parser.extract_content_ids([1, 101, 2, 3]) + assert result == [2, 3] + + def test_text_only_no_end(self, text_only_parser): + result = text_only_parser.extract_content_ids([1, 2, 3]) + assert result == [] + + +# --------------------------------------------------------------------------- +# Fixture request workaround (pytest "request" fixture shadowed by parametrize) +# --------------------------------------------------------------------------- + + +@pytest.fixture +def request_(request): + """Proxy for the pytest ``request`` fixture to avoid name collisions + with the ``request_obj`` fixture used in extract_reasoning tests.""" + return request diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index ac52de953..148412012 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -84,6 +84,9 @@ def register_omni_models_to_vllm(): if arch not in supported_archs: ModelRegistry.register_model(arch, f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}:{cls_name}") + # Register omni-specific reasoning parsers (e.g., step_audio). + import vllm_omni.reasoning # noqa: F401 + @dataclass class OmniEngineArgs(EngineArgs): diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 32c92ff68..eb989db0e 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -89,6 +89,20 @@ class StageEngineCoreProc(EngineCoreProc): signal_callback: SignalCallback | None = None maybe_register_config_serialize_by_value() + # Register vllm-omni reasoning parsers (e.g. step_audio) in this + # subprocess so they are available when the engine core resolves + # ``--reasoning-parser``. The main process already registered them + # at import time, but the forked subprocess starts with a fresh + # ReasoningParserManager. + try: + import vllm_omni.reasoning # noqa: F401 + except ImportError: + logger.warning( + "Failed to import vllm_omni.reasoning in subprocess; " + "custom reasoning parsers (e.g. step_audio) will not be " + "available." + ) + engine_core: StageEngineCoreProc | None = None coord_client: OmniCoordClientForStage | None = None try: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 67a639b15..0164e584a 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -425,6 +425,15 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) + # In vllm-omni's multi-process architecture, create_engine_config() runs + # in the subprocess (StageEngineCoreProc), not in the API server process. + # Propagate reasoning_parser from CLI args into structured_outputs_config + # here so that OmniOpenAIServingChat receives the correct value. + # (Upstream vLLM does this in EngineArgs.create_engine_config().) + if hasattr(args, "reasoning_parser") and args.reasoning_parser: + if not args.structured_outputs_config.reasoning_parser: + args.structured_outputs_config.reasoning_parser = args.reasoning_parser + # Load logging config for uvicorn if specified log_config = get_uvicorn_log_config(args) if log_config is not None: diff --git a/vllm_omni/reasoning/__init__.py b/vllm_omni/reasoning/__init__.py new file mode 100644 index 000000000..37cdcf0fd --- /dev/null +++ b/vllm_omni/reasoning/__init__.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from vllm.reasoning import ReasoningParserManager + +# Step-Audio specific reasoning parser. +# Registered at import time so it is available before the engine +# resolves ``--reasoning-parser step_audio``. + +_OMNI_REASONING_PARSERS = { + "step_audio": ( + "vllm_omni.reasoning.step_audio_reasoning_parser", + "StepAudioReasoningParser", + ), +} + + +def register_omni_reasoning_parsers(): + """Register vllm-omni reasoning parsers with vLLM's ReasoningParserManager.""" + for name, (module_path, class_name) in _OMNI_REASONING_PARSERS.items(): + ReasoningParserManager.register_lazy_module(name, module_path, class_name) + + +register_omni_reasoning_parsers() diff --git a/vllm_omni/reasoning/step_audio_reasoning_parser.py b/vllm_omni/reasoning/step_audio_reasoning_parser.py new file mode 100644 index 000000000..81ef63496 --- /dev/null +++ b/vllm_omni/reasoning/step_audio_reasoning_parser.py @@ -0,0 +1,574 @@ +# 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.openai.engine.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. ```` → [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_END_TEXT = "" + + # 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. ""). + # 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. " 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. " 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. "" → [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 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 "..."). 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. " 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