diff --git a/gemma/gm/text/_chat_sampler.py b/gemma/gm/text/_chat_sampler.py index 887ee2ad..d6b75b88 100644 --- a/gemma/gm/text/_chat_sampler.py +++ b/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) + diff --git a/gemma/gm/text/_gemma4_sampler.py b/gemma/gm/text/_gemma4_sampler.py index 153e6bf3..5578cd0e 100644 --- a/gemma/gm/text/_gemma4_sampler.py +++ b/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 + diff --git a/gemma/gm/text/_prefill.py b/gemma/gm/text/_prefill.py index 4c70a868..8d88a103 100644 --- a/gemma/gm/text/_prefill.py +++ b/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 + diff --git a/gemma/gm/text/_sampler.py b/gemma/gm/text/_sampler.py index 567b598e..2a3a6cdb 100644 --- a/gemma/gm/text/_sampler.py +++ b/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') + diff --git a/gemma/gm/text/_sampler_loop.py b/gemma/gm/text/_sampler_loop.py index 3124aad0..cc1140f6 100644 --- a/gemma/gm/text/_sampler_loop.py +++ b/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 + diff --git a/gemma/gm/text/_tokenizer.py b/gemma/gm/text/_tokenizer.py index be7db0d5..6d3610a0 100644 --- a/gemma/gm/text/_tokenizer.py +++ b/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 +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] # '' END_OF_TOOL_RESPONSE: ClassVar[int] # '' + BEGIN_OF_THINKING_CHANNEL: ClassVar[int] # '<|channel>' (start of thinking) + END_OF_THINKING_CHANNEL: ClassVar[int] # '' (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 # (EOA) + # Thinking channel tokens (Gemma4 only) + # <|channel> opens the thinking block, 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 # (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, ' ') +