Skip to content
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 25 additions & 3 deletions crates/reasoning_parser/src/factory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@ use parking_lot::RwLock;

use crate::{
parsers::{
BaseReasoningParser, CohereCmdParser, DeepSeekR1Parser, Glm45Parser, InklingParser,
KimiK3Parser, KimiParser, MiniMaxParser, NanoV3Parser, PassthroughParser, Qwen3Parser,
QwenThinkingParser, Step3Parser,
BaseReasoningParser, CohereCmdParser, DeepSeekR1Parser, Gemma4Parser, Glm45Parser,
InklingParser, KimiK3Parser, KimiParser, MiniMaxParser, NanoV3Parser, PassthroughParser,
Qwen3Parser, QwenThinkingParser, Step3Parser,
},
traits::{ParserConfig, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE},
};
Expand Down Expand Up @@ -148,6 +148,9 @@ impl ParserFactory {

registry.register_parser("nano_v3", || Box::new(NanoV3Parser::new()));

// Gemma 4 channel-marker reasoning; markers are special tokens.
registry.register_parser("gemma4", || Box::new(Gemma4Parser::new()));

registry.register_parser("inkling", || Box::new(InklingParser::new()));

// standard think tokens, always_in_reasoning=false
Expand Down Expand Up @@ -236,6 +239,11 @@ impl ParserFactory {
registry.register_pattern("nemotron-3", "nano_v3");

// Inkling checkpoints use the model-family name in their ID or config.
// Gemma 4 only — earlier Gemma generations have no reasoning channel,
// so no bare "gemma" pattern.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Nit: The Gemma block was inserted between the pre-existing Inkling comment (line 241) and its register_pattern("inkling", ...) call (line 247), orphaning the comment from the code it describes. Lines 241–243 now read as a single comment block, which is confusing at first glance.

Swapping the order keeps each comment adjacent to its pattern:

Suggested change
// Inkling checkpoints use the model-family name in their ID or config.
// Gemma 4 only — earlier Gemma generations have no reasoning channel,
// so no bare "gemma" pattern.
// Gemma 4 only — earlier Gemma generations have no reasoning channel,
// so no bare "gemma" pattern.
registry.register_pattern("gemma-4", "gemma4");
registry.register_pattern("gemma4", "gemma4");
// Inkling checkpoints use the model-family name in their ID or config.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 264ee29 — the Gemma block now sits above the Inkling comment, keeping each comment adjacent to its pattern.

registry.register_pattern("gemma-4", "gemma4");
registry.register_pattern("gemma4", "gemma4");

registry.register_pattern("inkling", "inkling");

Self { registry }
Expand Down Expand Up @@ -348,6 +356,20 @@ mod tests {
assert_eq!(factory.create("kimi-chat").model_type(), "kimi");
}

#[test]
fn test_factory_creates_gemma4_for_gemma_4_only() {
let factory = ParserFactory::new();
assert_eq!(
factory.create("google/gemma-4-27b-it").model_type(),
"gemma4"
);
// Earlier generations must NOT match — they have no reasoning channel.
assert_eq!(
factory.create("google/gemma-3-27b-it").model_type(),
"passthrough"
);
}

#[test]
fn test_factory_creates_inkling() {
let factory = ParserFactory::new();
Expand Down
6 changes: 3 additions & 3 deletions crates/reasoning_parser/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,9 @@ pub mod traits;

pub use factory::{ParserFactory, ParserRegistry};
pub use parsers::{
BaseReasoningParser, CohereCmdParser, DeepSeekR1Parser, Glm45Parser, InklingParser,
KimiK3Parser, KimiParser, MiniMaxParser, NanoV3Parser, PassthroughParser, Qwen3Parser,
QwenThinkingParser, Step3Parser,
BaseReasoningParser, CohereCmdParser, DeepSeekR1Parser, Gemma4Parser, Glm45Parser,
InklingParser, KimiK3Parser, KimiParser, MiniMaxParser, NanoV3Parser, PassthroughParser,
Qwen3Parser, QwenThinkingParser, Step3Parser,
};
pub use traits::{
ParseError, ParserConfig, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE,
Expand Down
229 changes: 229 additions & 0 deletions crates/reasoning_parser/src/parsers/gemma4.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,229 @@
//! Gemma 4 reasoning parser.
//!
//! Gemma 4 delimits chain-of-thought with channel markers rather than think
//! tags: reasoning opens with `<|channel>` followed by a `thought\n` role
//! label and closes with `<channel|>` (format per the family's public
//! reference parsing utilities, which also document that some checkpoints
//! emit a bare `thought\n` label even with thinking disabled). Both markers
//! are special tokens, so detokenization must preserve them —
//! `requires_special_tokens` returns true.
//!
//! Delegates marker handling to [`BaseReasoningParser`] and adds the two
//! Gemma-specific behaviors: stripping the `thought\n` role label from the
//! start of reasoning content (with streaming hold-back while the label is
//! still a possible prefix), and stripping the spurious bare label from
//! non-streaming output when no markers are present.

use crate::{
parsers::BaseReasoningParser,
traits::{ParseError, ParserConfig, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE},
};

const THINK_START: &str = "<|channel>";
const THINK_END: &str = "<channel|>";
const THOUGHT_LABEL: &str = "thought\n";

pub struct Gemma4Parser {
base: BaseReasoningParser,
/// Streaming: whether the leading `thought\n` label decision was made.
label_handled: bool,
/// Streaming: reasoning held back while it is still a label prefix.
pending_reasoning: String,
}

impl Gemma4Parser {
pub fn new() -> Self {
let config = ParserConfig {
think_start_token: THINK_START.to_string(),
think_end_token: THINK_END.to_string(),
stream_reasoning: true,
max_buffer_size: DEFAULT_MAX_BUFFER_SIZE,
always_in_reasoning: false,
};
Self {
base: BaseReasoningParser::new(config).with_model_type("gemma4".to_string()),
label_handled: false,
pending_reasoning: String::new(),
}
}
}

impl Default for Gemma4Parser {
fn default() -> Self {
Self::new()
}
}

impl ReasoningParser for Gemma4Parser {
fn detect_and_parse_reasoning(&mut self, text: &str) -> Result<ParserResult, ParseError> {
// Thinking disabled, spurious role label: no markers anywhere, but
// the output leads with the bare label — strip it from content.
if !text.contains(THINK_START) && !text.contains(THINK_END) {
if let Some(rest) = text.strip_prefix(THOUGHT_LABEL) {
return Ok(ParserResult::normal(rest.to_string()));
}
return self.base.detect_and_parse_reasoning(text);
}

let mut result = self.base.detect_and_parse_reasoning(text)?;
if let Some(rest) = result.reasoning_text.strip_prefix(THOUGHT_LABEL) {
result.reasoning_text = rest.to_string();
}
Ok(result)
}

fn parse_reasoning_streaming_incremental(
&mut self,
text: &str,
) -> Result<ParserResult, ParseError> {
let mut result = self.base.parse_reasoning_streaming_incremental(text)?;
if self.label_handled {
return Ok(result);
}

if !result.reasoning_text.is_empty() {
self.pending_reasoning.push_str(&result.reasoning_text);
result.reasoning_text = String::new();
if let Some(rest) = self.pending_reasoning.strip_prefix(THOUGHT_LABEL) {
result.reasoning_text = rest.to_string();
self.pending_reasoning.clear();
self.label_handled = true;
} else if !THOUGHT_LABEL.starts_with(self.pending_reasoning.as_str()) {
// Diverged from the label: release everything held back.
result.reasoning_text = std::mem::take(&mut self.pending_reasoning);
self.label_handled = true;
}
// else: still a strict prefix of the label — keep holding.
}

// Reasoning closed while text was still held (e.g. a reasoning block
// shorter than the label): release it so nothing is swallowed.
if !result.normal_text.is_empty() && !self.pending_reasoning.is_empty() {
let held = std::mem::take(&mut self.pending_reasoning);
result.reasoning_text = format!("{held}{}", result.reasoning_text);
self.label_handled = true;
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

Ok(result)
}

fn reset(&mut self) {
self.base.reset();
self.label_handled = false;
self.pending_reasoning.clear();
}

fn model_type(&self) -> &str {
self.base.model_type()
}

fn requires_special_tokens(&self) -> bool {
// The channel markers are special tokens; stripping them during
// detokenization would leave the parser nothing to split on.
true
}

fn is_in_reasoning(&self) -> bool {
self.base.is_in_reasoning()
}

fn mark_reasoning_started(&mut self) {
self.base.mark_reasoning_started();
}

fn mark_think_start_stripped(&mut self) {
self.base.mark_think_start_stripped();
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn non_streaming_splits_and_strips_label() {
let mut parser = Gemma4Parser::new();
let result = parser
.detect_and_parse_reasoning("<|channel>thought\nlet me think<channel|>The answer is 4")
.unwrap();
assert_eq!(result.reasoning_text, "let me think");
assert_eq!(result.normal_text, "The answer is 4");
}

#[test]
fn non_streaming_strips_spurious_bare_label() {
let mut parser = Gemma4Parser::new();
let result = parser
.detect_and_parse_reasoning("thought\nThe answer is 4")
.unwrap();
assert_eq!(result.reasoning_text, "");
assert_eq!(result.normal_text, "The answer is 4");
}

#[test]
fn non_streaming_plain_output_untouched() {
let mut parser = Gemma4Parser::new();
let result = parser
.detect_and_parse_reasoning("The answer is 4")
.unwrap();
assert_eq!(result.normal_text, "The answer is 4");
assert_eq!(result.reasoning_text, "");
}

#[test]
fn streaming_strips_label_split_across_chunks() {
let mut parser = Gemma4Parser::new();
let mut reasoning = String::new();
let mut normal = String::new();
for chunk in [
"<|channel>",
"thou",
"ght\nstep one ",
"and two",
"<channel|>",
"answer",
] {
let r = parser.parse_reasoning_streaming_incremental(chunk).unwrap();
reasoning.push_str(&r.reasoning_text);
normal.push_str(&r.normal_text);
}
assert_eq!(reasoning, "step one and two");
assert_eq!(normal, "answer");
}

#[test]
fn streaming_releases_non_label_reasoning() {
let mut parser = Gemma4Parser::new();
let mut reasoning = String::new();
let mut normal = String::new();
// Reasoning that never carries the role label must not be swallowed.
for chunk in ["<|channel>", "no label here", "<channel|>", "done"] {
let r = parser.parse_reasoning_streaming_incremental(chunk).unwrap();
reasoning.push_str(&r.reasoning_text);
normal.push_str(&r.normal_text);
}
assert_eq!(reasoning, "no label here");
assert_eq!(normal, "done");
}

#[test]
fn streaming_releases_held_prefix_when_block_ends_early() {
let mut parser = Gemma4Parser::new();
let mut reasoning = String::new();
let mut normal = String::new();
// A reasoning block shorter than the label ("thou") must be released
// when the block closes, not dropped.
for chunk in ["<|channel>thou<channel|>", "answer"] {
let r = parser.parse_reasoning_streaming_incremental(chunk).unwrap();
reasoning.push_str(&r.reasoning_text);
normal.push_str(&r.normal_text);
}
assert_eq!(reasoning, "thou");
assert_eq!(normal, "answer");
}

#[test]
fn requires_special_tokens_is_true() {
assert!(Gemma4Parser::new().requires_special_tokens());
}
}
2 changes: 2 additions & 0 deletions crates/reasoning_parser/src/parsers/mod.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
pub mod base;
pub mod cohere_cmd;
pub mod deepseek_r1;
pub mod gemma4;
pub mod glm45;
pub mod inkling;
pub mod kimi;
Expand All @@ -14,6 +15,7 @@ pub mod step3;
pub use base::BaseReasoningParser;
pub use cohere_cmd::CohereCmdParser;
pub use deepseek_r1::DeepSeekR1Parser;
pub use gemma4::Gemma4Parser;
pub use glm45::Glm45Parser;
pub use inkling::InklingParser;
pub use kimi::KimiParser;
Expand Down
Loading