model: support Command A plus (#26106)
Co-authored-by: Yihao Wang <42559837+AgainstEntropy@users.noreply.github.com>
This commit is contained in:
zijiexia
@@ -597,6 +597,245 @@ class _PoolsideV1Detector(Qwen3Detector):
|
||||
self.reasoning_default = "explicit_enable_thinking"
|
||||
|
||||
|
||||
class CohereCommand4Detector(BaseReasoningFormatDetector):
|
||||
"""Detector for Cohere Command4 / Command-A family (incl. cohere2_moe and
|
||||
cohere2_vision Command-A-Plus).
|
||||
|
||||
Generated format (the assistant prefix in the chat template already emits
|
||||
``<|START_THINKING|>`` when ``reasoning=True``, so the *generated* text
|
||||
typically begins inside the thinking block):
|
||||
|
||||
thinking_content<|END_THINKING|><|START_TEXT|>final_answer<|END_TEXT|>
|
||||
|
||||
When ``reasoning=False`` the chat template emits both START/END_THINKING
|
||||
in the prefix and the generated text is just::
|
||||
|
||||
<|START_TEXT|>final_answer<|END_TEXT|>
|
||||
|
||||
This detector returns:
|
||||
- ``reasoning_text`` = the thinking block (between START_THINKING and
|
||||
END_THINKING, with the START tag stripped if the model echoed it).
|
||||
- ``normal_text`` = the content between ``<|START_TEXT|>`` and
|
||||
``<|END_TEXT|>``, with both markers stripped. If no ``<|START_TEXT|>``
|
||||
appears (the model exhausted max_new_tokens still inside thinking),
|
||||
``normal_text`` is the empty string.
|
||||
|
||||
Matches the public token names from the model's
|
||||
``special_tokens_map.json`` (``<|START_THINKING|>`` etc.).
|
||||
"""
|
||||
|
||||
TEXT_START_TOKEN = "<|START_TEXT|>"
|
||||
TEXT_END_TOKEN = "<|END_TEXT|>"
|
||||
# When the model decides to call tools instead of producing a final text
|
||||
# block, it emits an action block instead of a text block. The reasoning
|
||||
# parser must leave that block intact so the downstream tool-call parser
|
||||
# can pick it up.
|
||||
ACTION_START_TOKEN = "<|START_ACTION|>"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stream_reasoning: bool = True,
|
||||
force_reasoning: bool = True,
|
||||
continue_final_message: bool = False,
|
||||
previous_content: str = "",
|
||||
):
|
||||
# The chat template puts <|START_THINKING|> in the assistant prefix
|
||||
# when reasoning is enabled, so the *generated* text usually starts
|
||||
# already inside thinking. ``force_reasoning=True`` makes the base
|
||||
# detector treat the leading bytes as reasoning even though the
|
||||
# generated stream typically does not echo <|START_THINKING|>.
|
||||
super().__init__(
|
||||
think_start_token="<|START_THINKING|>",
|
||||
think_end_token="<|END_THINKING|>",
|
||||
force_reasoning=force_reasoning,
|
||||
stream_reasoning=stream_reasoning,
|
||||
continue_final_message=continue_final_message,
|
||||
previous_content=previous_content,
|
||||
)
|
||||
# Streaming state machine. The model emits, in order:
|
||||
# 1. reasoning (between START_THINKING [in prefix] and END_THINKING)
|
||||
# 2. either ``<|START_TEXT|>...<|END_TEXT|>`` (final answer) or
|
||||
# ``<|START_ACTION|>...<|END_ACTION|>`` (tool calls) -- never both.
|
||||
# When ``reasoning=False`` the chat template emits both START/END
|
||||
# thinking in the prefix and step 1 is empty; the generated stream
|
||||
# then starts directly with the text or action block.
|
||||
self._reasoning_done = False
|
||||
self._saw_text_start = False
|
||||
self._saw_text_end = False
|
||||
self._in_action_mode = False
|
||||
|
||||
@classmethod
|
||||
def _strip_text_markers(cls, raw: str) -> str:
|
||||
"""Extract the substring between ``<|START_TEXT|>`` and
|
||||
``<|END_TEXT|>``. If ``<|START_TEXT|>`` is absent but a
|
||||
``<|START_ACTION|>`` block is present, the model produced a tool
|
||||
call instead of a text answer -- return the raw text untouched so
|
||||
the downstream tool-call parser can pick up the action block. If
|
||||
neither marker is present (ran out of tokens still inside
|
||||
thinking) return ``""``. If ``<|END_TEXT|>`` is absent (stop token
|
||||
or max_new_tokens cut the stream off inside the text block) return
|
||||
everything after ``<|START_TEXT|>``.
|
||||
"""
|
||||
if not raw:
|
||||
return ""
|
||||
s = raw.find(cls.TEXT_START_TOKEN)
|
||||
if s == -1:
|
||||
if cls.ACTION_START_TOKEN in raw:
|
||||
return raw
|
||||
return ""
|
||||
s += len(cls.TEXT_START_TOKEN)
|
||||
tail = raw[s:]
|
||||
e = tail.find(cls.TEXT_END_TOKEN)
|
||||
if e == -1:
|
||||
return tail
|
||||
return tail[:e]
|
||||
|
||||
def detect_and_parse(self, text: str) -> StreamingParseResult:
|
||||
# Direct parse: split on the (single) ``<|END_THINKING|>`` token if
|
||||
# present. Anything before is reasoning, anything after is the
|
||||
# final-text block. If no END_THINKING but a START_TEXT exists,
|
||||
# we're in the reasoning=False case (chat template emitted both
|
||||
# START/END thinking in the prefix; the model only generated the
|
||||
# text block). Otherwise the model exhausted tokens still thinking
|
||||
# and ``normal_text`` ends up empty -- matching the convention of
|
||||
# the other detectors in this module (DeepSeekR1, Qwen3, ...). The
|
||||
# empty content is propagated as ``message.content = None`` by
|
||||
# serving_chat, and downstream code is expected to treat that as
|
||||
# "no answer" rather than falling back to ``reasoning_content``.
|
||||
end_think_idx = text.find(self.think_end_token)
|
||||
text_start_idx = text.find(self.TEXT_START_TOKEN)
|
||||
action_start_idx = text.find(self.ACTION_START_TOKEN)
|
||||
if end_think_idx != -1:
|
||||
reasoning = text[:end_think_idx]
|
||||
rest = text[end_think_idx + len(self.think_end_token) :]
|
||||
elif text_start_idx != -1:
|
||||
reasoning = text[:text_start_idx]
|
||||
rest = text[text_start_idx:]
|
||||
elif action_start_idx != -1:
|
||||
# reasoning=False + tool call: chat template emitted both
|
||||
# START/END thinking in the prefix, the model only generated
|
||||
# an action block. Treat the prefix before the action block as
|
||||
# (probably empty) reasoning so the action block reaches the
|
||||
# tool-call parser intact.
|
||||
reasoning = text[:action_start_idx]
|
||||
rest = text[action_start_idx:]
|
||||
else:
|
||||
reasoning = text
|
||||
rest = ""
|
||||
|
||||
# Some checkpoints echo the START_THINKING token even though the
|
||||
# chat template put it in the prefix; drop it if so.
|
||||
think_start_text = self.think_start_token + self.think_start_self_label
|
||||
if reasoning.startswith(think_start_text):
|
||||
reasoning = reasoning[len(think_start_text) :]
|
||||
|
||||
return StreamingParseResult(
|
||||
normal_text=self._strip_text_markers(rest),
|
||||
reasoning_text=reasoning,
|
||||
)
|
||||
|
||||
def parse_streaming_increment(self, new_text: str) -> StreamingParseResult:
|
||||
"""Streaming parse. Custom state machine -- we don't reuse the base
|
||||
class because Cohere's "reasoning=False" path (the model emits no
|
||||
``<|END_THINKING|>``, just goes straight to a text or action block)
|
||||
is fundamentally incompatible with the base detector's
|
||||
``force_reasoning`` semantics."""
|
||||
self._buffer += new_text
|
||||
buf = self._buffer
|
||||
|
||||
if not self._reasoning_done:
|
||||
# Look for any marker that ends reasoning: an explicit
|
||||
# END_THINKING, or an implicit transition via the start of the
|
||||
# final-text or action block (reasoning=False case).
|
||||
markers = (
|
||||
(self.think_end_token, "think_end"),
|
||||
(self.TEXT_START_TOKEN, "text"),
|
||||
(self.ACTION_START_TOKEN, "action"),
|
||||
)
|
||||
first_pos = None
|
||||
first_marker = None
|
||||
first_kind = None
|
||||
for marker_text, kind in markers:
|
||||
p = buf.find(marker_text)
|
||||
if p != -1 and (first_pos is None or p < first_pos):
|
||||
first_pos, first_marker, first_kind = p, marker_text, kind
|
||||
if first_pos is None:
|
||||
# No marker seen yet. Stream the reasoning prefix, but keep
|
||||
# enough tail in the buffer to recognise a marker split
|
||||
# across chunk boundaries.
|
||||
if not self.stream_reasoning:
|
||||
return StreamingParseResult()
|
||||
max_keep = max(len(m) for m, _ in markers) - 1
|
||||
if len(buf) > max_keep:
|
||||
head = buf[:-max_keep]
|
||||
self._buffer = buf[-max_keep:]
|
||||
return StreamingParseResult(reasoning_text=head)
|
||||
return StreamingParseResult()
|
||||
|
||||
reasoning_chunk = buf[:first_pos]
|
||||
if first_kind == "think_end":
|
||||
self._buffer = buf[first_pos + len(first_marker) :]
|
||||
else:
|
||||
# Implicit reasoning-end: leave the start-of-block marker in
|
||||
# the buffer for the post-thinking branch below to consume.
|
||||
self._buffer = buf[first_pos:]
|
||||
self._reasoning_done = True
|
||||
if reasoning_chunk:
|
||||
return StreamingParseResult(reasoning_text=reasoning_chunk)
|
||||
buf = self._buffer
|
||||
|
||||
# Reasoning is closed. Decide between text-stripping and
|
||||
# action-passthrough on first sight of a marker.
|
||||
if self._in_action_mode:
|
||||
if not buf:
|
||||
return StreamingParseResult()
|
||||
self._buffer = ""
|
||||
return StreamingParseResult(normal_text=buf)
|
||||
|
||||
if not self._saw_text_start:
|
||||
s_text = buf.find(self.TEXT_START_TOKEN)
|
||||
s_action = buf.find(self.ACTION_START_TOKEN)
|
||||
picks = [
|
||||
(p, k) for p, k in ((s_text, "text"), (s_action, "action")) if p != -1
|
||||
]
|
||||
if not picks:
|
||||
max_keep = (
|
||||
max(len(self.TEXT_START_TOKEN), len(self.ACTION_START_TOKEN)) - 1
|
||||
)
|
||||
if len(buf) > max_keep:
|
||||
self._buffer = buf[-max_keep:]
|
||||
return StreamingParseResult()
|
||||
picks.sort()
|
||||
first_pos, first_kind = picks[0]
|
||||
if first_kind == "action":
|
||||
self._in_action_mode = True
|
||||
out_normal = buf[first_pos:]
|
||||
self._buffer = ""
|
||||
return StreamingParseResult(normal_text=out_normal)
|
||||
# Found <|START_TEXT|>. Drop everything up to and including the
|
||||
# marker -- text content streams next.
|
||||
self._buffer = buf[first_pos + len(self.TEXT_START_TOKEN) :]
|
||||
self._saw_text_start = True
|
||||
buf = self._buffer
|
||||
|
||||
if self._saw_text_start and not self._saw_text_end:
|
||||
e = buf.find(self.TEXT_END_TOKEN)
|
||||
if e == -1:
|
||||
# Emit everything except a possible partial END_TEXT tail.
|
||||
keep = len(self.TEXT_END_TOKEN) - 1
|
||||
if len(buf) > keep:
|
||||
out_normal = buf[:-keep]
|
||||
self._buffer = buf[-keep:]
|
||||
return StreamingParseResult(normal_text=out_normal)
|
||||
return StreamingParseResult()
|
||||
out_normal = buf[:e]
|
||||
self._buffer = buf[e + len(self.TEXT_END_TOKEN) :]
|
||||
self._saw_text_end = True
|
||||
return StreamingParseResult(normal_text=out_normal)
|
||||
|
||||
return StreamingParseResult()
|
||||
|
||||
|
||||
class ReasoningParser:
|
||||
"""
|
||||
Parser that handles both streaming and non-streaming scenarios for extracting
|
||||
@@ -629,6 +868,7 @@ class ReasoningParser:
|
||||
"nemotron_3": Nemotron3Detector,
|
||||
"interns1": Qwen3Detector,
|
||||
"gemma4": Gemma4Detector,
|
||||
"cohere_command4": CohereCommand4Detector,
|
||||
}
|
||||
|
||||
def __init__(
|
||||
|
||||
Reference in New Issue
Block a user