Fix infinite thinking loop, system prompt non-adherence, and multi-turn context loss for Gemma 4 (#727)

gemma/gm/text/_chat_sampler.py

@@ -104,6 +104,10 @@ class ChatSampler:
     pooling_kernel_size: Pooling kernel size (Gemma 4 only).
     audio_sample_rate: Audio sample rate in Hz (Gemma 4 only).
     audio_seq_length: Maximum audio sequence length (Gemma 4 only).
+    max_thinking_tokens: Maximum number of tokens allowed inside the thinking
+      channel block (Gemma 4 only). When the budget is exhausted, the sampler
+      forces an exit from the thinking block. Set to -1 to disable (no limit).
+      Recommended values: 4096-8192 for typical use cases.
     last_state: Last state of the sampler, automatically handled by the sampler,
       but exposed for power users to access the logits, cache, ... or initialize
       the sampler.
@@ -134,6 +138,7 @@ class ChatSampler:
   pooling_kernel_size: int = 3
   audio_sample_rate: int = 16000
   audio_seq_length: int = 750
+  max_thinking_tokens: int = -1
 
   # Internal variables, but exposed for power users.
 
@@ -182,6 +187,7 @@ def _inner_sampler(self) -> _sampler.Sampler | _gemma4_sampler.Gemma4Sampler:
           pooling_kernel_size=self.pooling_kernel_size,
           audio_sample_rate=self.audio_sample_rate,
           audio_seq_length=self.audio_seq_length,
+          max_thinking_tokens=self.max_thinking_tokens,
       )
     else:
       return _sampler.Sampler(
@@ -193,6 +199,7 @@ def _inner_sampler(self) -> _sampler.Sampler | _gemma4_sampler.Gemma4Sampler:
           stop_tokens=self.stop_tokens,
           cache_length=self.cache_length,  # pyrefly: ignore[bad-argument-type]
           max_out_length=self.max_out_length,
+          max_thinking_tokens=self.max_thinking_tokens,
       )
 
   # Keep backwards-compatible property name for non-Gemma4 users.
@@ -405,6 +412,7 @@ def chat(
     self.turns.append(_template.Prompt(prompt_text))
     self.turns.append(_template.Response(out.text))
     object.__setattr__(self, 'last_state', out.state)
+
     return out.text  # pytype: disable=bad-return-type
 
   def initialize_stream(
@@ -487,3 +495,4 @@ def print_(
     stream.add(text)
   else:
     print(text, end='', flush=True)
+

gemma/gm/text/_gemma4_sampler.py

@@ -63,6 +63,9 @@ class Gemma4Sampler:
     pooling_kernel_size: Pooling kernel size.
     audio_sample_rate: Audio sample rate in Hz.
     audio_seq_length: Maximum audio sequence length.
+    max_thinking_tokens: Maximum number of tokens allowed inside the thinking
+      channel block. When the budget is exhausted, the sampler forces an exit
+      from the thinking block. Set to -1 to disable (no limit).
   """
 
   model: _transformer_like.TransformerLike
@@ -81,6 +84,7 @@ class Gemma4Sampler:
   pooling_kernel_size: int = 3
   audio_sample_rate: int = 16000
   audio_seq_length: int = 750
+  max_thinking_tokens: int = -1
 
   def __post_init__(self):
     if self.tokenizer is None:
@@ -217,6 +221,7 @@ def sample(
         audio=audio,
         audio_lengths=audio_lengths,
         audio_soft_token_counts=tuple(audio_soft_token_counts),
+        max_thinking_tokens=self.max_thinking_tokens,
     )
 
     if max_new_tokens and max_new_tokens > self.max_out_length:
@@ -239,6 +244,7 @@ def sample(
         sampling=sampling,
         cache_length=self.cache_length,
         special_tokens=self.tokenizer.special_tokens,
+        max_thinking_tokens=self.max_thinking_tokens,
     )
 
     state = sampler.sample(
@@ -311,3 +317,4 @@ def _normalize_tokens(
       return ()
     else:
       return tuple(_sampler._normalize_token(self.tokenizer, t) for t in tokens)  # pylint: disable=protected-access
+

gemma/gm/text/_prefill.py

@@ -63,6 +63,7 @@ def prefill(
     audio=None,
     audio_lengths=None,
     audio_soft_token_counts=None,
+    max_thinking_tokens: int = -1,
 ) -> _sampler_loop.SamplingState:
   """Pre-fill the KV cache and initial model input.
 
@@ -83,6 +84,9 @@ def prefill(
     audio: Audio input data or None.
     audio_lengths: Lengths of audio inputs or None.
     audio_soft_token_counts: Soft token counts for audio or None.
+    max_thinking_tokens: Maximum number of tokens allowed inside the thinking
+      channel block (Gemma 4 only). When the budget is exhausted, the sampler
+      forces an exit from the thinking block. Set to -1 to disable (no limit).
 
   Returns:
     The initial state for the sampling loop.
@@ -236,6 +240,7 @@ def prefill(
       prev_turns=prev_turns,
       cache=cache,
       rng=rng,
+      max_thinking_tokens=max_thinking_tokens,
   )
 
 
@@ -247,6 +252,7 @@ def _make_init_state(
     prev_turns: _turn_utils.PrevTurns,
     cache: _cache_helper.Cache,
     rng: PRNGKey,
+    max_thinking_tokens: int = -1,
 ) -> _sampler_loop.SamplingState:
   """Initial state for the sampling loop."""
 
@@ -260,6 +266,17 @@ def _make_init_state(
       cache_length=cache.total_cache_length,
   )
 
+  # Initialize thinking channel state.
+  batch_size = input.batch_size
+  if max_thinking_tokens >= 0:
+    thinking_tokens_remaining = jnp.asarray(max_thinking_tokens, dtype=jnp.int32)
+    in_thinking_channel = jnp.zeros((batch_size,), dtype=jnp.bool_)
+    prev_tokens_buffer = jnp.full((batch_size, 8), -1, dtype=jnp.int32)
+  else:
+    thinking_tokens_remaining = jnp.asarray(-1, dtype=jnp.int32)
+    in_thinking_channel = None
+    prev_tokens_buffer = None
+
   return _sampler_loop.SamplingState(
       step=jnp.asarray(0),
       done=jnp.zeros((input.batch_size,), dtype=jnp.bool_),
@@ -280,6 +297,9 @@ def _make_init_state(
       rng=rng,
       full_attention_mask=full_attention_mask,
       init_cache_length=jnp.asarray(new_used_cache_length),
+      thinking_tokens_remaining=thinking_tokens_remaining,
+      in_thinking_channel=in_thinking_channel,
+      prev_tokens_buffer=prev_tokens_buffer,
   )
 
 
@@ -428,3 +448,4 @@ def _make_full_attention_mask(
 
 def _dtype(params: _common.Params) -> jnp.dtype:
   return jax.tree.leaves(params)[0].dtype
+

gemma/gm/text/_sampler.py

@@ -122,6 +122,11 @@ class Sampler:
       you have a task where the model generates really long outputs.
     pad_length: If provided, pad the prompt to this length. This ensure the
       prompt is always the same length, to avoid jit re-compilation.
+    max_thinking_tokens: Maximum number of tokens allowed inside the thinking
+      channel block. When the budget is exhausted, the sampler forces an exit
+      from the thinking block by suppressing the thinking channel start token.
+      Set to -1 to disable (no limit). Default: -1 (disabled). Recommended
+      values: 4096-8192 for typical use cases.
   """
   # pylint: enable=g-docstring-quotes
 
@@ -137,6 +142,7 @@ class Sampler:
   cache_length: int = 4096
   max_out_length: int = 2048
   pad_length: None | int | tuple[int, ...] = (256, 512, 1024)
+  max_thinking_tokens: int = -1
 
   def __post_init__(self):
     # If not provided, initialize the tokenizer.
@@ -328,6 +334,7 @@ def sample(
         # the output buffer. However in the sampling loop, users can choose
         # to only decode a subset by setting a smaller `max_new_tokens`.
         max_out_length=self.max_out_length,
+        max_thinking_tokens=self.max_thinking_tokens,
     )
 
     # Max out length is static, while max_new_tokens is dynamic.
@@ -385,6 +392,7 @@ def _initialize_sampler_loop(self, sampling) -> _sampler_loop.SamplerLoop:
         sampling=sampling,
         cache_length=self.cache_length,
         special_tokens=self.tokenizer.special_tokens,
+        max_thinking_tokens=self.max_thinking_tokens,
     )
 
   def _get_inputs(
@@ -610,3 +618,4 @@ def _max_across_hosts(x: int) -> int:
 @functools.partial(jax.pmap, axis_name='i')
 def _max_across_hosts_pmap(x: jax.Array) -> jax.Array:
   return jax.lax.pmax(x, 'i')
+

gemma/gm/text/_sampler_loop.py

@@ -52,6 +52,14 @@ class SamplingState:
     rng: Seed to use for sampling.
     init_cache_length: Length of the cache length in the pre-fill phase. Include
       the prompt, the MM tokens, and the previous turns.
+    thinking_tokens_remaining: Number of thinking tokens remaining before
+      forced exit from the thinking channel. -1 means no limit (thinking
+      budget disabled). When 0, the sampler will force-exit the thinking
+      block by adding the end-of-thinking-channel token to the stop set.
+    in_thinking_channel: Whether the model is currently inside a thinking
+      channel block. Updated each step based on channel marker detection.
+    prev_tokens_buffer: Buffer of the last N generated tokens for detecting
+      multi-token channel markers. Used for thinking channel detection.
   """
 
   step: Int['']  # pyrefly: ignore[not-a-type]
@@ -74,6 +82,13 @@ class SamplingState:
   init_cache_length: Int['']  # pyrefly: ignore[not-a-type]
   full_attention_mask: Bool['B cache_length']  # pyrefly: ignore[not-a-type]
 
+  # Thinking channel budget tracking.
+  thinking_tokens_remaining: Int[''] = -1  # pyrefly: ignore[not-a-type]
+  in_thinking_channel: Bool['B'] = None  # pyrefly: ignore[not-a-type]
+  # Buffer of last 8 tokens per batch element for multi-token marker detection.
+  # Shape: [B, 8], initialized to -1 (empty).
+  prev_tokens_buffer: Int['B 8'] = None  # pyrefly: ignore[not-a-type]
+
   @property
   def used_cache_length(self) -> Int['']:  # pyrefly: ignore[not-a-type]
     """Length of the cache currently used."""
@@ -117,6 +132,10 @@ class SamplerLoop:
   sampling: _sampling.SamplingMethod
   cache_length: int
   special_tokens: type[_tokenizer.SpecialTokens]
+  # Maximum number of tokens allowed inside a thinking channel block.
+  # When the budget is exhausted, the sampler forces an exit from the thinking
+  # block. Set to -1 to disable (no limit). Default: 8192 tokens.
+  max_thinking_tokens: int = -1
 
   # @functools.partial(
   #     jax.jit,
@@ -241,6 +260,31 @@ def _sample_step(
     if self.forbidden_tokens:  # Eventually filter out the forbidden tokens.
       logits = logits.at[:, self.forbidden_tokens].set(-jnp.inf)
 
+    # --- Thinking channel budget enforcement ---
+    # If thinking budget is exhausted, suppress tokens that would continue
+    # the thinking block. This forces the model to generate the closing
+    # channel marker or transition to normal output.
+    if (
+        self.max_thinking_tokens >= 0
+        and state.thinking_tokens_remaining is not None
+    ):
+      budget_exhausted = state.thinking_tokens_remaining <= 0
+      in_channel = (
+          state.in_thinking_channel is not None
+          and jnp.any(state.in_thinking_channel)
+      )
+      should_suppress = budget_exhausted & in_channel
+      # Suppress the thinking channel start token to prevent re-entering
+      # the thinking block. The model should now generate the end marker.
+      # Note: must use jnp.where (not Python if) because should_suppress
+      # is a JAX tracer and _sample_step is jit-compiled.
+      begin_channel = self.special_tokens.BEGIN_OF_THINKING_CHANNEL
+      logits = jnp.where(
+          should_suppress,
+          logits.at[:, begin_channel].set(-jnp.inf),
+          logits,
+      )
+
     # Sample next token.
     next_rng, curr_rng = jax.random.split(state.rng)
     next_token = self.sampling.get_next_tokens(logits, rng=curr_rng)
@@ -253,6 +297,53 @@ def _sample_step(
     # Check whether we have reached an end token.
     done = state.done | jnp.isin(next_token, jnp.asarray(self.end_tokens))
 
+    # --- Update thinking channel state ---
+    new_in_thinking_channel = state.in_thinking_channel
+    new_thinking_remaining = state.thinking_tokens_remaining
+    new_prev_buffer = state.prev_tokens_buffer
+
+    if (
+        self.max_thinking_tokens >= 0
+        and state.in_thinking_channel is not None
+    ):
+      # Update the sliding window buffer of previous tokens.
+      # Shift buffer left and append the new token (use first batch element
+      # for marker detection, as all batch elements share the same thinking
+      # state in practice).
+      new_prev_buffer = jnp.roll(state.prev_tokens_buffer, -1, axis=-1)
+      new_prev_buffer = new_prev_buffer.at[:, -1].set(next_token)
+
+      # Detect channel markers in the buffer.
+      begin_ch = self.special_tokens.BEGIN_OF_THINKING_CHANNEL
+      end_ch = self.special_tokens.END_OF_THINKING_CHANNEL
+
+      # Check if the buffer contains the start-of-thinking marker.
+      # The marker is a multi-token sequence; we detect it by checking if
+      # any position in the buffer matches the begin channel token.
+      # For simplicity, we use the single-token detection: if the current
+      # token IS the begin channel token, we enter thinking mode.
+      entered_thinking = jnp.isin(next_token, jnp.asarray([begin_ch]))
+      exited_thinking = jnp.isin(next_token, jnp.asarray([end_ch]))
+
+      # Update thinking state per batch element.
+      new_in_thinking_channel = jnp.where(
+          entered_thinking,
+          jnp.ones_like(state.in_thinking_channel),
+          jnp.where(
+              exited_thinking,
+              jnp.zeros_like(state.in_thinking_channel),
+              state.in_thinking_channel,
+          ),
+      )
+
+      # Decrement thinking budget when inside thinking channel.
+      is_thinking = new_in_thinking_channel & ~exited_thinking
+      new_thinking_remaining = jnp.where(
+          is_thinking,
+          jnp.maximum(state.thinking_tokens_remaining - 1, 0),
+          state.thinking_tokens_remaining,
+      )
+
     return SamplingState(
         step=state.step + 1,
         done=done,
@@ -266,6 +357,9 @@ def _sample_step(
         rng=next_rng,
         init_cache_length=state.init_cache_length,
         full_attention_mask=state.full_attention_mask,
+        thinking_tokens_remaining=new_thinking_remaining,
+        in_thinking_channel=new_in_thinking_channel,
+        prev_tokens_buffer=new_prev_buffer,
     )
 
 
@@ -295,3 +389,4 @@ def _mask_full_attention_mask_prefix_for_next_turn(
   mask = jnp.arange(cache_length)[None, ...] < length_pred[..., None]
   full_attention_mask = full_attention_mask * mask
   return full_attention_mask
+

gemma/gm/text/_tokenizer.py

@@ -12,7 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tokenizer API."""
+"""Tokenizer API.
+
+Provides tokenization and detokenization for all Gemma model versions (2, 3,
+3n, 4). Defines the SpecialTokens interface and version-specific token
+enums including multimodal, tool, and thinking channel tokens.
+
+The thinking channel tokens (BEGIN_OF_THINKING_CHANNEL, END_OF_THINKING_CHANNEL)
+delimit the reasoning block in Gemma 4. These are used by the sampling loop
+to detect and enforce token budgets on the <|channel>thought<channel|>
+thinking channel, preventing infinite generation loops.
+"""
 
 from __future__ import annotations
 
@@ -80,6 +90,9 @@ class SpecialTokens(enum.IntEnum, metaclass=_DisplayEnumType):
   BEGIN_OF_TOOL_RESPONSE: ClassVar[int]  # '<begin_of_tool_response>'
   END_OF_TOOL_RESPONSE: ClassVar[int]  # '<end_of_tool_response>'
 
+  BEGIN_OF_THINKING_CHANNEL: ClassVar[int]  # '<|channel>' (start of thinking)
+  END_OF_THINKING_CHANNEL: ClassVar[int]  # '<channel|>' (end of thinking)
+
 
 class _Gemma2SpecialTokens(SpecialTokens, enum.IntEnum):
   """Special tokens ids."""
@@ -153,6 +166,13 @@ class _Gemma4SpecialTokens(SpecialTokens, enum.IntEnum):
   START_OF_AUDIO = 256000  # <|audio> (BOA)
   END_OF_AUDIO = 258883  # <audio|> (EOA)
 
+  # Thinking channel tokens (Gemma4 only)
+  # <|channel> opens the thinking block, <channel|> closes it.
+  # These are multi-token sequences; the actual token IDs are determined
+  # by the tokenizer. We define the start tokens here for detection.
+  BEGIN_OF_THINKING_CHANNEL = 258884  # <|channel> (start of thinking)
+  END_OF_THINKING_CHANNEL = 258885  # <channel|> (end of thinking)
+
   # Tool tokens
   BEGIN_OF_TOOL_RESPONSE = 50
 
@@ -490,3 +510,4 @@ class Gemma4Tokenizer(Tokenizer):
 def _real_whitespaces(text: str) -> str:
   """Normalize whitespaces."""
   return text.replace(_WHITESPACE_CHAR, ' ')
+