257 lines
8.2 KiB
Python
257 lines
8.2 KiB
Python
"""Normalize model reasoning and final-response channels across LLM backends."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import logging
|
|
from typing import Any, Protocol, Sequence
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_QWEN3_REASONING_PARSER = "qwen3"
|
|
_NO_REASONING_PARSER = "none"
|
|
_ARCHITECTURE_REASONING_PARSERS = {
|
|
"Qwen3_5ForConditionalGeneration": _QWEN3_REASONING_PARSER,
|
|
}
|
|
|
|
|
|
class ReasoningProfile(Protocol):
|
|
"""Describe the profile fields needed to select a reasoning parser."""
|
|
|
|
architecture: str
|
|
reasoning_parser: str
|
|
|
|
|
|
class TokenDecoder(Protocol):
|
|
"""Describe the tokenizer operations used for token-aware extraction."""
|
|
|
|
def get_vocab(self) -> dict[str, int]:
|
|
"""Return the token-to-id vocabulary."""
|
|
|
|
def decode(self, token_ids: Sequence[int], **kwargs: Any) -> str:
|
|
"""Decode generated token IDs into text."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ReasoningOutput:
|
|
"""Hold a clean final response and its separately extracted reasoning."""
|
|
|
|
response: str
|
|
reasoning: str
|
|
reasoning_tokens: int
|
|
parser: str
|
|
|
|
|
|
class ReasoningOutputParser(Protocol):
|
|
"""Define backend-neutral reasoning extraction from one generated candidate."""
|
|
|
|
parser_name: str
|
|
requires_boundary_tokens: bool
|
|
|
|
def extract(
|
|
self,
|
|
text: str,
|
|
token_ids: Sequence[int],
|
|
*,
|
|
native_reasoning: str | None = None,
|
|
) -> ReasoningOutput:
|
|
"""Separate reasoning from the final response."""
|
|
|
|
|
|
def reasoning_parser_name(profile: ReasoningProfile) -> str:
|
|
"""Return the configured model-output parser, with architecture fallback."""
|
|
configured = profile.reasoning_parser.strip().lower()
|
|
if configured:
|
|
return configured
|
|
return _ARCHITECTURE_REASONING_PARSERS.get(
|
|
profile.architecture,
|
|
_NO_REASONING_PARSER,
|
|
)
|
|
|
|
|
|
def reasoning_chat_template_kwargs(
|
|
profile: ReasoningProfile,
|
|
enable_reasoning: bool = True,
|
|
) -> dict[str, bool]:
|
|
"""Return chat-template settings that agree with the requested reasoning state."""
|
|
if reasoning_parser_name(profile) == _QWEN3_REASONING_PARSER:
|
|
return {"enable_thinking": enable_reasoning}
|
|
return {}
|
|
|
|
|
|
class PassthroughReasoningParser:
|
|
"""Return ordinary non-reasoning model output unchanged."""
|
|
|
|
parser_name = _NO_REASONING_PARSER
|
|
requires_boundary_tokens = False
|
|
|
|
def extract(
|
|
self,
|
|
text: str,
|
|
token_ids: Sequence[int],
|
|
*,
|
|
native_reasoning: str | None = None,
|
|
) -> ReasoningOutput:
|
|
"""Use a native reasoning channel when an engine provides one."""
|
|
del token_ids
|
|
return ReasoningOutput(
|
|
response=text.strip(),
|
|
reasoning=(native_reasoning or "").strip(),
|
|
reasoning_tokens=0,
|
|
parser="native" if native_reasoning is not None else self.parser_name,
|
|
)
|
|
|
|
|
|
class DisabledReasoningParser:
|
|
"""Treat generated text as a direct response when reasoning is disabled."""
|
|
|
|
parser_name = "disabled"
|
|
requires_boundary_tokens = False
|
|
|
|
def extract(
|
|
self,
|
|
text: str,
|
|
token_ids: Sequence[int],
|
|
*,
|
|
native_reasoning: str | None = None,
|
|
) -> ReasoningOutput:
|
|
"""Return only the direct response and discard unexpected reasoning metadata."""
|
|
del token_ids, native_reasoning
|
|
return ReasoningOutput(
|
|
response=text.strip(),
|
|
reasoning="",
|
|
reasoning_tokens=0,
|
|
parser=self.parser_name,
|
|
)
|
|
|
|
|
|
class Qwen3ReasoningParser:
|
|
"""Split Qwen3 reasoning delimited by its chat-template thinking tokens."""
|
|
|
|
parser_name = _QWEN3_REASONING_PARSER
|
|
requires_boundary_tokens = True
|
|
_start_token = "<think>"
|
|
_end_token = "</think>"
|
|
|
|
def __init__(self, tokenizer: TokenDecoder) -> None:
|
|
"""Retain the model tokenizer so special boundary tokens survive decoding."""
|
|
self._tokenizer = tokenizer
|
|
|
|
def _decode(self, token_ids: Sequence[int]) -> str:
|
|
"""Decode one channel while removing model control tokens."""
|
|
return str(
|
|
self._tokenizer.decode(
|
|
list(token_ids),
|
|
skip_special_tokens=True,
|
|
clean_up_tokenization_spaces=False,
|
|
)
|
|
).strip()
|
|
|
|
def _boundary_ids(self) -> tuple[int | None, int | None]:
|
|
"""Return exact reasoning boundary IDs when the tokenizer exposes them."""
|
|
vocabulary = self._tokenizer.get_vocab()
|
|
return vocabulary.get(self._start_token), vocabulary.get(self._end_token)
|
|
|
|
def _extract_from_tokens(
|
|
self,
|
|
token_ids: Sequence[int],
|
|
) -> ReasoningOutput | None:
|
|
"""Split at the exact end token without mistaking user-visible text for markup."""
|
|
start_id, end_id = self._boundary_ids()
|
|
if end_id is None:
|
|
return None
|
|
generated_ids = list(token_ids)
|
|
try:
|
|
boundary_index = generated_ids.index(end_id)
|
|
except ValueError:
|
|
reasoning_ids = generated_ids
|
|
response_ids: list[int] = []
|
|
else:
|
|
reasoning_ids = generated_ids[:boundary_index]
|
|
response_ids = generated_ids[boundary_index + 1 :]
|
|
while response_ids and response_ids[0] == end_id:
|
|
response_ids = response_ids[1:]
|
|
if start_id is not None and start_id in reasoning_ids:
|
|
reasoning_ids = reasoning_ids[reasoning_ids.index(start_id) + 1 :]
|
|
return ReasoningOutput(
|
|
response=self._decode(response_ids),
|
|
reasoning=self._decode(reasoning_ids),
|
|
reasoning_tokens=len(reasoning_ids),
|
|
parser=self.parser_name,
|
|
)
|
|
|
|
def _extract_from_text(self, text: str) -> ReasoningOutput:
|
|
"""Fall back to Qwen's documented delimiters when token IDs are unavailable."""
|
|
raw_text = text.strip()
|
|
if self._end_token in raw_text:
|
|
reasoning, response = raw_text.split(self._end_token, maxsplit=1)
|
|
else:
|
|
reasoning, response = raw_text, ""
|
|
if self._start_token in reasoning:
|
|
reasoning = reasoning.split(self._start_token, maxsplit=1)[1]
|
|
return ReasoningOutput(
|
|
response=response.strip(),
|
|
reasoning=reasoning.strip(),
|
|
reasoning_tokens=0,
|
|
parser=self.parser_name,
|
|
)
|
|
|
|
def extract(
|
|
self,
|
|
text: str,
|
|
token_ids: Sequence[int],
|
|
*,
|
|
native_reasoning: str | None = None,
|
|
) -> ReasoningOutput:
|
|
"""Prefer an engine-native channel, then exact tokens, then text delimiters."""
|
|
if native_reasoning is not None:
|
|
return ReasoningOutput(
|
|
response=text.strip(),
|
|
reasoning=native_reasoning.strip(),
|
|
reasoning_tokens=0,
|
|
parser="native",
|
|
)
|
|
token_result = self._extract_from_tokens(token_ids) if token_ids else None
|
|
return token_result or self._extract_from_text(text)
|
|
|
|
|
|
def create_reasoning_parser(
|
|
profile: ReasoningProfile,
|
|
tokenizer: TokenDecoder,
|
|
) -> ReasoningOutputParser:
|
|
"""Create the model-aware output parser shared by every inference backend."""
|
|
parser_name = reasoning_parser_name(profile)
|
|
logger.debug(
|
|
"Selected Modal LLM reasoning parser=%s architecture=%s.",
|
|
parser_name,
|
|
profile.architecture or "unspecified",
|
|
)
|
|
if parser_name == _NO_REASONING_PARSER:
|
|
return PassthroughReasoningParser()
|
|
if parser_name == _QWEN3_REASONING_PARSER:
|
|
return Qwen3ReasoningParser(tokenizer)
|
|
raise ValueError(f"Unsupported Modal LLM reasoning parser {parser_name!r}.")
|
|
|
|
|
|
def reasoning_parser_for_request(
|
|
configured_parser: ReasoningOutputParser,
|
|
enable_reasoning: bool,
|
|
) -> ReasoningOutputParser:
|
|
"""Return the configured parser or a direct-response parser for this request."""
|
|
if enable_reasoning:
|
|
return configured_parser
|
|
return DisabledReasoningParser()
|
|
|
|
|
|
__all__ = [
|
|
"DisabledReasoningParser",
|
|
"PassthroughReasoningParser",
|
|
"Qwen3ReasoningParser",
|
|
"ReasoningOutput",
|
|
"ReasoningOutputParser",
|
|
"create_reasoning_parser",
|
|
"reasoning_chat_template_kwargs",
|
|
"reasoning_parser_for_request",
|
|
"reasoning_parser_name",
|
|
]
|