-
Notifications
You must be signed in to change notification settings - Fork 142
feat(reasoning): add Gemma 4 channel-marker reasoning parser #2164
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
pallasathena92
wants to merge
3
commits into
main
Choose a base branch
from
feat/gemma4-reasoning
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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; | ||
| } | ||
|
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()); | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
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:
There was a problem hiding this comment.
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.