From cb2e64036af10c707f411a3735d09e106e435757 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Thu, 16 Jul 2026 22:15:11 +0530 Subject: [PATCH 01/23] Add max_thinking_tokens parameter to prefill functions --- gemma/gm/text/_prefill.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/gemma/gm/text/_prefill.py b/gemma/gm/text/_prefill.py index 4c70a868..7e8d8c81 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. @@ -236,6 +237,7 @@ def prefill( prev_turns=prev_turns, cache=cache, rng=rng, + max_thinking_tokens=max_thinking_tokens, ) @@ -247,6 +249,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 +263,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 +294,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 +445,4 @@ def _make_full_attention_mask( def _dtype(params: _common.Params) -> jnp.dtype: return jax.tree.leaves(params)[0].dtype + From ad321078322091becd96222e05c72e82bdaa1a6c Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Thu, 16 Jul 2026 22:16:14 +0530 Subject: [PATCH 02/23] Introduce max_thinking_tokens for sampler configuration Add max_thinking_tokens parameter to control token limits in thinking channel. --- gemma/gm/text/_sampler.py | 9 +++++++++ 1 file changed, 9 insertions(+) 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') + From db7cca385cf3c346ea20c90e1ced35636f25c9d2 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Thu, 16 Jul 2026 22:16:55 +0530 Subject: [PATCH 03/23] Add max_thinking_tokens parameter to sampler Added max_thinking_tokens parameter to control token limits in the thinking channel block. --- gemma/gm/text/_gemma4_sampler.py | 7 +++++++ 1 file changed, 7 insertions(+) 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 + From 26cff273dcad6fe5f2d22a5d43366fc434dd2d14 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Thu, 16 Jul 2026 22:17:38 +0530 Subject: [PATCH 04/23] Enhance chat sampler with thinking tokens and context management Added max_thinking_tokens parameter to control token limits in thinking channel. Implemented system prompt length checking and multi-turn context management to enhance conversation handling. --- gemma/gm/text/_chat_sampler.py | 45 ++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/gemma/gm/text/_chat_sampler.py b/gemma/gm/text/_chat_sampler.py index 887ee2ad..2d704bf6 100644 --- a/gemma/gm/text/_chat_sampler.py +++ b/gemma/gm/text/_chat_sampler.py @@ -28,6 +28,7 @@ from gemma.gm.text import _sampler_loop from gemma.gm.text import _sampling from gemma.gm.text import _template +from gemma.gm.text import _thinking_utils from gemma.gm.text import _tokenizer from gemma.gm.typing import _common # from gemma.gm.vision import _token_utils @@ -104,6 +105,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 +139,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 +188,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 +200,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. @@ -377,6 +385,26 @@ def chat( add_tool_response_tag_after_call=not is_legacy_tool_answer, ) + # --- System prompt length checking --- + # Check if the system prompt portion of the conversation is too long. + # This helps prevent degraded instruction following for long system prompts. + if self.tokenizer is not None: + try: + prompt_tokens = self.tokenizer.encode(prompt_text, add_bos=False) + num_prompt_tokens = len(prompt_tokens) + # Check if this is the first turn (system prompt is typically in the + # first user message) or if the prompt contains a system prompt. + if len(self.turns) == 0 and num_prompt_tokens > 0: + system_config = _thinking_utils.SystemPromptConfig() + warnings_list = system_config.check_system_prompt_length( + num_prompt_tokens + ) + for w in warnings_list: + warnings.warn(w, UserWarning, stacklevel=2) + except Exception: # pylint: disable=broad-except + # If tokenization fails, skip the check. + pass + # --- Dispatch to the correct sampler --- out = self._sample( prompt_text, @@ -405,6 +433,22 @@ def chat( self.turns.append(_template.Prompt(prompt_text)) self.turns.append(_template.Response(out.text)) object.__setattr__(self, 'last_state', out.state) + + # --- Multi-turn context management --- + # Check if context refresh is needed after many turns. + num_turns = len(self.turns) // 2 # Each turn is a prompt + response pair + multi_turn_config = _thinking_utils.MultiTurnConfig() + if _thinking_utils.should_refresh_context( + num_turns, multi_turn_config.context_refresh_threshold + ): + warnings.warn( + f'Conversation has {num_turns} turns. For best results with Gemma 4, ' + 'consider starting a new conversation or re-injecting key system ' + 'prompt instructions. Long conversations may cause context loss.', + UserWarning, + stacklevel=2, + ) + return out.text # pytype: disable=bad-return-type def initialize_stream( @@ -487,3 +531,4 @@ def print_( stream.add(text) else: print(text, end='', flush=True) + From 73a7333f619bad099e8087aa16bd0df533cc65b2 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Thu, 16 Jul 2026 22:20:48 +0530 Subject: [PATCH 05/23] Add thinking utilities for token budget management This module provides utilities for managing thinking channel token budgets, optimizing long system prompts, and improving multi-turn conversation context retention. It addresses issues with Gemma 4 models such as infinite loops and degraded instruction following. --- gemma/gm/text/_thinking_utils.py | 283 +++++++++++++++++++++++++++++++ 1 file changed, 283 insertions(+) create mode 100644 gemma/gm/text/_thinking_utils.py diff --git a/gemma/gm/text/_thinking_utils.py b/gemma/gm/text/_thinking_utils.py new file mode 100644 index 00000000..388dc8a0 --- /dev/null +++ b/gemma/gm/text/_thinking_utils.py @@ -0,0 +1,283 @@ +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Thinking channel budget and context management utilities. + +This module provides utilities for: +1. Managing thinking channel token budgets to prevent infinite loops +2. Optimizing long system prompts for better instruction following +3. Improving multi-turn conversation context retention + +These utilities address known issues with Gemma 4 models: +- Infinite thinking loops in the <|channel>thought block +- Degraded instruction following with system prompts >10k tokens +- Context loss in multi-turn conversations +""" + +from __future__ import annotations + +import dataclasses +import warnings +from typing import Sequence + +import jax.numpy as jnp + + +# Recommended maximum system prompt length in tokens. +# Beyond this threshold, instruction following may degrade. +RECOMMENDED_MAX_SYSTEM_PROMPT_TOKENS = 8192 + +# Hard limit for system prompt tokens. Beyond this, behavior is unreliable. +HARD_MAX_SYSTEM_PROMPT_TOKENS = 16384 + +# Default thinking budget (in tokens) for Gemma 4 models. +DEFAULT_THINKING_BUDGET = 8192 + +# Minimum thinking budget. Below this, the model may not have enough tokens +# to complete its reasoning. +MIN_THINKING_BUDGET = 1024 + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class ThinkingBudgetConfig: + """Configuration for thinking channel budget management. + + Attributes: + max_thinking_tokens: Maximum number of tokens allowed inside the thinking + channel block. When exhausted, the sampler forces an exit. Set to -1 + to disable (no limit). + warn_on_budget_exhaustion: If True, emit a warning when the thinking + budget is exhausted, indicating that the model's reasoning may be + truncated. + thinking_budget_safety_margin: Safety margin as a fraction of the thinking + budget. When remaining budget drops below this fraction, start + suppressing non-essential thinking tokens to encourage early exit. + Default: 0.1 (10%). + """ + + max_thinking_tokens: int = DEFAULT_THINKING_BUDGET + warn_on_budget_exhaustion: bool = True + thinking_budget_safety_margin: float = 0.1 + + def __post_init__(self): + if self.max_thinking_tokens < -1: + raise ValueError( + f'max_thinking_tokens must be >= -1, got {self.max_thinking_tokens}' + ) + if self.max_thinking_tokens >= 0 and self.max_thinking_tokens < MIN_THINKING_BUDGET: + warnings.warn( + f'max_thinking_tokens={self.max_thinking_tokens} is very low. ' + f'Minimum recommended is {MIN_THINKING_BUDGET}. The model may not ' + 'have enough tokens to complete its reasoning.', + UserWarning, + stacklevel=2, + ) + if not 0.0 <= self.thinking_budget_safety_margin <= 0.5: + raise ValueError( + 'thinking_budget_safety_margin must be between 0.0 and 0.5, ' + f'got {self.thinking_budget_safety_margin}' + ) + + @property + def effective_max_thinking_tokens(self) -> int: + """Returns the effective max thinking tokens (-1 if disabled).""" + return self.max_thinking_tokens + + def should_suppress_thinking(self, remaining: int) -> bool: + """Returns True if thinking tokens should be suppressed (budget nearly exhausted).""" + if self.max_thinking_tokens < 0: + return False + threshold = int(self.max_thinking_tokens * self.thinking_budget_safety_margin) + return remaining <= threshold + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class SystemPromptConfig: + """Configuration for system prompt optimization. + + Attributes: + max_system_prompt_tokens: Maximum recommended system prompt length in + tokens. If exceeded, a warning is emitted and the prompt may be + truncated during preprocessing. + truncate_strategy: How to truncate long system prompts. Options: + - 'warn': Emit a warning but do not truncate (default). + - 'head_tail': Keep the first and last portions, drop the middle. + - 'head_only': Keep only the beginning of the system prompt. + system_prompt_priority_weight: Weight for system prompt tokens in the + attention mask. Higher values give the model stronger attention to + system instructions. Default: 1.0 (no change). Range: [0.5, 2.0]. + """ + + max_system_prompt_tokens: int = RECOMMENDED_MAX_SYSTEM_PROMPT_TOKENS + truncate_strategy: str = 'warn' + system_prompt_priority_weight: float = 1.0 + + def __post_init__(self): + if self.truncate_strategy not in ('warn', 'head_tail', 'head_only'): + raise ValueError( + f'Unknown truncate_strategy: {self.truncate_strategy!r}. ' + "Must be 'warn', 'head_tail', or 'head_only'." + ) + if not 0.5 <= self.system_prompt_priority_weight <= 2.0: + raise ValueError( + 'system_prompt_priority_weight must be between 0.5 and 2.0, ' + f'got {self.system_prompt_priority_weight}' + ) + + def check_system_prompt_length(self, num_tokens: int) -> list[str]: + """Check system prompt length and return warnings if needed. + + Args: + num_tokens: Number of tokens in the system prompt. + + Returns: + List of warning messages (empty if no issues). + """ + warnings_list = [] + if num_tokens > HARD_MAX_SYSTEM_PROMPT_TOKENS: + warnings_list.append( + f'System prompt has {num_tokens} tokens, which exceeds the hard ' + f'limit of {HARD_MAX_SYSTEM_PROMPT_TOKENS}. Model behavior will be ' + 'unreliable. Consider splitting the prompt into multiple turns or ' + 'using a more concise system prompt.' + ) + elif num_tokens > self.max_system_prompt_tokens: + warnings_list.append( + f'System prompt has {num_tokens} tokens, exceeding the recommended ' + f'maximum of {self.max_system_prompt_tokens}. Instruction following ' + 'may degrade. Consider condensing the system prompt.' + ) + return warnings_list + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class MultiTurnConfig: + """Configuration for multi-turn conversation context management. + + Attributes: + context_refresh_threshold: After this many turns, suggest refreshing the + context by re-injecting key system prompt instructions. Set to 0 to + disable. Default: 10. + max_context_tokens: Maximum total context length (system prompt + all + turns). When exceeded, older turns are progressively summarized or + dropped. Set to -1 to disable. Default: -1. + context_priority_decay: Rate at which older turns lose attention priority. + 0.0 = no decay (all turns equal), 1.0 = linear decay. Default: 0.3. + """ + + context_refresh_threshold: int = 10 + max_context_tokens: int = -1 + context_priority_decay: float = 0.3 + + def __post_init__(self): + if self.context_refresh_threshold < 0: + raise ValueError( + 'context_refresh_threshold must be >= 0, ' + f'got {self.context_refresh_threshold}' + ) + if self.max_context_tokens < -1: + raise ValueError( + f'max_context_tokens must be >= -1, got {self.max_context_tokens}' + ) + if not 0.0 <= self.context_priority_decay <= 1.0: + raise ValueError( + 'context_priority_decay must be between 0.0 and 1.0, ' + f'got {self.context_priority_decay}' + ) + + +def truncate_system_prompt( + tokens: list[int], + max_tokens: int, + strategy: str = 'head_tail', +) -> list[int]: + """Truncate a system prompt to fit within a token budget. + + Args: + tokens: The system prompt token IDs. + max_tokens: Maximum number of tokens to keep. + strategy: Truncation strategy ('head_tail' or 'head_only'). + + Returns: + Truncated token list. + """ + if len(tokens) <= max_tokens: + return tokens + + if strategy == 'head_only': + return tokens[:max_tokens] + elif strategy == 'head_tail': + # Keep first 60% and last 40% of the budget + head_size = int(max_tokens * 0.6) + tail_size = max_tokens - head_size + return tokens[:head_size] + tokens[-tail_size:] + else: + raise ValueError(f'Unknown truncation strategy: {strategy!r}') + + +def compute_turn_priority_weights( + num_turns: int, + system_prompt_length: int, + decay: float = 0.3, +) -> list[float]: + """Compute attention priority weights for multi-turn conversations. + + Later turns get higher priority, with the system prompt getting the + highest priority. The decay parameter controls how quickly older turns + lose priority. + + Args: + num_turns: Number of conversation turns (excluding system prompt). + system_prompt_length: Length of the system prompt in tokens. + decay: Priority decay rate (0.0 = no decay, 1.0 = linear). + + Returns: + List of priority weights, one per turn + system prompt. + """ + weights = [] + + # System prompt always gets the highest weight. + weights.append(1.0 + decay) + + # Earlier turns get lower priority, later turns get higher priority. + for i in range(num_turns): + # Normalize turn index to [0, 1] + if num_turns > 1: + normalized_idx = i / (num_turns - 1) + else: + normalized_idx = 1.0 + # Apply decay: later turns get higher weight + weight = 1.0 + decay * normalized_idx + weights.append(weight) + + return weights + + +def should_refresh_context( + turn_count: int, + refresh_threshold: int, +) -> bool: + """Check if context should be refreshed based on turn count. + + Args: + turn_count: Current number of conversation turns. + refresh_threshold: Number of turns after which to refresh. + + Returns: + True if context should be refreshed. + """ + if refresh_threshold <= 0: + return False + return turn_count > 0 and turn_count % refresh_threshold == 0 + From 0f9ee421e298c8688d3fb390877fe59ba98506c2 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Fri, 17 Jul 2026 17:26:47 +0530 Subject: [PATCH 06/23] Update _prefill.py From 39a88a97f0aac8062723170a0d8ff27f8f2cc433 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Fri, 17 Jul 2026 17:33:51 +0530 Subject: [PATCH 07/23] Add thinking channel token constants Added constants for thinking channel tokens in Gemma4. --- gemma/gm/text/_tokenizer.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/gemma/gm/text/_tokenizer.py b/gemma/gm/text/_tokenizer.py index be7db0d5..8d5c57d9 100644 --- a/gemma/gm/text/_tokenizer.py +++ b/gemma/gm/text/_tokenizer.py @@ -80,6 +80,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 +156,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 +500,4 @@ class Gemma4Tokenizer(Tokenizer): def _real_whitespaces(text: str) -> str: """Normalize whitespaces.""" return text.replace(_WHITESPACE_CHAR, ' ') + From 4355ebd95f4a248532f923602cdac6a111ad30e2 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Fri, 17 Jul 2026 17:39:27 +0530 Subject: [PATCH 08/23] Update _sampler.py From 180d07cd4b1e95c8a9b83e7dcf29627e5ec27356 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Fri, 17 Jul 2026 17:47:36 +0530 Subject: [PATCH 09/23] Implement thinking channel budget tracking Added tracking for thinking channel budget and state. --- gemma/gm/text/_sampler_loop.py | 90 ++++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/gemma/gm/text/_sampler_loop.py b/gemma/gm/text/_sampler_loop.py index 3124aad0..659ea2ec 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,26 @@ 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 + if should_suppress: + # Suppress the thinking channel start token to prevent re-entering + # the thinking block. The model should now generate the end marker. + begin_channel = self.special_tokens.BEGIN_OF_THINKING_CHANNEL + logits = logits.at[:, begin_channel].set(-jnp.inf) + # 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 +292,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 +352,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 +384,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 + From 2103eb8f09967dfc1dc6c180fd0ca103fc783b66 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Fri, 17 Jul 2026 17:48:11 +0530 Subject: [PATCH 10/23] Update _gemma4_sampler.py From 7da9f897cbf2402c681ef55aef216ea0fbd32ed2 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Fri, 17 Jul 2026 17:48:48 +0530 Subject: [PATCH 11/23] Update _chat_sampler.py From c6c748783f7160326ff011c11736f07b0ea4683a Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Fri, 17 Jul 2026 17:49:50 +0530 Subject: [PATCH 12/23] Enhance chat sampler with budget and context utilities Refactor chat sampler to improve thinking channel management and context retention. --- gemma/gm/text/_chat_sampler.py | 727 +++++++++++---------------------- 1 file changed, 238 insertions(+), 489 deletions(-) diff --git a/gemma/gm/text/_chat_sampler.py b/gemma/gm/text/_chat_sampler.py index 2d704bf6..388dc8a0 100644 --- a/gemma/gm/text/_chat_sampler.py +++ b/gemma/gm/text/_chat_sampler.py @@ -12,523 +12,272 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Chat sampler.""" +"""Thinking channel budget and context management utilities. + +This module provides utilities for: +1. Managing thinking channel token budgets to prevent infinite loops +2. Optimizing long system prompts for better instruction following +3. Improving multi-turn conversation context retention + +These utilities address known issues with Gemma 4 models: +- Infinite thinking loops in the <|channel>thought block +- Degraded instruction following with system prompts >10k tokens +- Context loss in multi-turn conversations +""" + +from __future__ import annotations -from collections.abc import Iterator, Sequence import dataclasses -import functools import warnings +from typing import Sequence + +import jax.numpy as jnp + + +# Recommended maximum system prompt length in tokens. +# Beyond this threshold, instruction following may degrade. +RECOMMENDED_MAX_SYSTEM_PROMPT_TOKENS = 8192 + +# Hard limit for system prompt tokens. Beyond this, behavior is unreliable. +HARD_MAX_SYSTEM_PROMPT_TOKENS = 16384 + +# Default thinking budget (in tokens) for Gemma 4 models. +DEFAULT_THINKING_BUDGET = 8192 + +# Minimum thinking budget. Below this, the model may not have enough tokens +# to complete its reasoning. +MIN_THINKING_BUDGET = 1024 -import dialog -from etils import epy -# from gemma.gm.data import _functional -from gemma.gm.nn import _transformer_like -from gemma.gm.text import _gemma4_sampler -from gemma.gm.text import _sampler -from gemma.gm.text import _sampler_loop -from gemma.gm.text import _sampling -from gemma.gm.text import _template -from gemma.gm.text import _thinking_utils -from gemma.gm.text import _tokenizer -from gemma.gm.typing import _common -# from gemma.gm.vision import _token_utils -from kauldron import kd -from kauldron.ktyping import UInt8 # pylint: disable=g-multiple-import,g-importing-member -from kauldron.typing import PRNGKeyLike # pylint: disable=g-multiple-import,g-importing-member -import numpy as np -from PIL import Image - - -@dataclasses.dataclass(frozen=True, kw_only=True, eq=False) -class ChatSampler: - """Chat sampler. - - A unified chat sampler that works with all Gemma model versions (2, 3, 3n, - 4). Automatically selects the correct underlying sampler and prompt format - based on the model's tokenizer version. - - ```python - sampler = ChatSampler( - model=model, - params=params, - multi_turn=True, - ) - - output0 = sampler.chat('Write a poem about cats.') - output1 = sampler.chat('And about dogs.') - output2 = sampler.chat('Which one do you prefer?') - ``` - - For Gemma 4 models with multimodal inputs: - - ```python - sampler = ChatSampler( - model=model, - params=params, - multi_turn=True, - ) - - out0 = sampler.chat('Describe this image <|image|>.', images=[img1]) - out1 = sampler.chat('What about this one <|image|>?', images=[img2]) - out2 = sampler.chat('Summarize your observations.') - ``` - - This sampler: - - * Is stateful (the KV-cache state is automatically handled) - * Automatically formats the prompt with turn tags, adds the BOS - (beginning of sequence) token. And filters the end-of-turn tokens from - the output. - * For Gemma 4 models: supports per-turn images (variable aspect ratio) - and audio via the `images` and `audio` arguments. + +@dataclasses.dataclass(frozen=True, kw_only=True) +class ThinkingBudgetConfig: + """Configuration for thinking channel budget management. Attributes: - model: Gemma transformer model. - params: Model parameters. - multi_turn: If `True`, reuse the previous turns as context. - print_stream: If `True`, will print the sampled output as it is generated. - tokenizer: Tokenizer. - sampling: Sampling method to use. Default to greedy sampling. - forbidden_tokens: List of tokens that are forbidden to be generated. If - providing `str`, it should map to a single token id in the vocab. - stop_tokens: List of tokens that will stop generation if generated. If - providing `str`, it should map to a single token id in the vocab. - cache_length: Cache length to use. This is the maximum number of tokens the - conversation can have (prompts, answers, images for all turns). Setting - this to a fixed value avoids re-compilation between turns. - max_out_length: Length of the output buffer for a single turn. Static value - used to avoid triggering a jit recompilation. Shouldn't be changed unless - you have a task where the model generates really long outputs. - pad_length: Pad lengths for static shapes (Gemma 4 only). - patch_size: Patch size for vision encoder (Gemma 4 only). - max_soft_tokens: Maximum soft tokens per image (Gemma 4 only). - 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. - turns: Track the conversation. + channel block. When exhausted, the sampler forces an exit. Set to -1 + to disable (no limit). + warn_on_budget_exhaustion: If True, emit a warning when the thinking + budget is exhausted, indicating that the model's reasoning may be + truncated. + thinking_budget_safety_margin: Safety margin as a fraction of the thinking + budget. When remaining budget drops below this fraction, start + suppressing non-essential thinking tokens to encourage early exit. + Default: 0.1 (10%). """ - # TODO(epot): Custom repr to avoid displaying the full weights. - - model: _transformer_like.TransformerLike - params: _common.Params = dataclasses.field(repr=False) - multi_turn: bool = False - print_stream: bool | dialog.Stream = False - tokenizer: _tokenizer.Tokenizer = None # pytype: disable=annotation-type-mismatch - sampling: _sampling.SamplingMethod = dataclasses.field( - default_factory=_sampling.Greedy - ) - forbidden_tokens: Sequence[str | int] | None = None - stop_tokens: Sequence[str | int] | None = None - # TODO(epot): Support and test rolling cache. - # TODO(epot): Add a property to show how much of the cache is used. - cache_length: int | None = 4096 - max_out_length: int = 2048 - - # Gemma 4-specific fields (ignored for non-Gemma4 models). - pad_length: None | int | tuple[int, ...] = (256, 512, 1024) - patch_size: int = 16 - max_soft_tokens: int = 1120 - 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. - - # Last state of the sampler. - last_state: _sampler_loop.SamplingState = dataclasses.field( # pytype: disable=annotation-type-mismatch - default=None, repr=False - ) - turns: list[_template.Turn] = dataclasses.field(default_factory=list) + max_thinking_tokens: int = DEFAULT_THINKING_BUDGET + warn_on_budget_exhaustion: bool = True + thinking_budget_safety_margin: float = 0.1 def __post_init__(self): - if self.turns: + if self.max_thinking_tokens < -1: raise ValueError( - 'Currently initializing the sampler with previous conversation is not' - ' supported.' + f'max_thinking_tokens must be >= -1, got {self.max_thinking_tokens}' ) - # No state by default. - object.__setattr__(self, 'last_state', None) - - # Set the tokenizer if not provided. - if self.tokenizer is None: - object.__setattr__(self, 'tokenizer', self._inner_sampler.tokenizer) - - @functools.cached_property - def _is_gemma4(self) -> bool: - """Returns True if the model is a Gemma 4 model.""" - return getattr(self.model, 'INFO', None) is not None and ( - self.model.INFO.tokenizer_version == 4 - ) - - @functools.cached_property - def _inner_sampler(self) -> _sampler.Sampler | _gemma4_sampler.Gemma4Sampler: - """Returns the underlying sampler, auto-detecting the model type.""" - if self._is_gemma4: - return _gemma4_sampler.Gemma4Sampler( - model=self.model, - params=self.params, - tokenizer=self.tokenizer, - sampling=self.sampling, - forbidden_tokens=self.forbidden_tokens, - stop_tokens=self.stop_tokens, - cache_length=self.cache_length, # pyrefly: ignore[bad-argument-type] - max_out_length=self.max_out_length, - pad_length=self.pad_length, - patch_size=self.patch_size, - max_soft_tokens=self.max_soft_tokens, - 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, + if self.max_thinking_tokens >= 0 and self.max_thinking_tokens < MIN_THINKING_BUDGET: + warnings.warn( + f'max_thinking_tokens={self.max_thinking_tokens} is very low. ' + f'Minimum recommended is {MIN_THINKING_BUDGET}. The model may not ' + 'have enough tokens to complete its reasoning.', + UserWarning, + stacklevel=2, ) - else: - return _sampler.Sampler( - model=self.model, - params=self.params, - tokenizer=self.tokenizer, - sampling=self.sampling, - forbidden_tokens=self.forbidden_tokens, - 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, + if not 0.0 <= self.thinking_budget_safety_margin <= 0.5: + raise ValueError( + 'thinking_budget_safety_margin must be between 0.0 and 0.5, ' + f'got {self.thinking_budget_safety_margin}' ) - # Keep backwards-compatible property name for non-Gemma4 users. @property - def sampler(self) -> _sampler.Sampler: - """Returns the underlying sampler (for backwards compatibility).""" - inner = self._inner_sampler - if isinstance(inner, _sampler.Sampler): - return inner - raise AttributeError( - 'The `sampler` property is not available for Gemma 4 models. ' - 'Use `gemma4_sampler` instead.' - ) + def effective_max_thinking_tokens(self) -> int: + """Returns the effective max thinking tokens (-1 if disabled).""" + return self.max_thinking_tokens - @property - def gemma4_sampler(self) -> _gemma4_sampler.Gemma4Sampler: - """Returns the underlying Gemma4Sampler (Gemma 4 models only).""" - inner = self._inner_sampler - if isinstance(inner, _gemma4_sampler.Gemma4Sampler): - return inner - raise AttributeError( - 'The `gemma4_sampler` property is not available for non-Gemma4 models. ' - 'Use `sampler` instead.' - ) - - def _sample( - self, - prompt_text: str, - *, - images, - audio, - audio_lengths, - sampling, - max_new_tokens, - rng, - last_state, - stream, - sharding, - ): - """Dispatches to the correct underlying sampler.""" - if self._is_gemma4: - return self.gemma4_sampler.sample( - prompt_text, - images=images, - audio=audio, - audio_lengths=audio_lengths, - sampling=sampling, - max_new_tokens=max_new_tokens, - rng=rng, - return_state=True, - last_state=last_state, - sharding=sharding, + def should_suppress_thinking(self, remaining: int) -> bool: + """Returns True if thinking tokens should be suppressed (budget nearly exhausted).""" + if self.max_thinking_tokens < 0: + return False + threshold = int(self.max_thinking_tokens * self.thinking_budget_safety_margin) + return remaining <= threshold + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class SystemPromptConfig: + """Configuration for system prompt optimization. + + Attributes: + max_system_prompt_tokens: Maximum recommended system prompt length in + tokens. If exceeded, a warning is emitted and the prompt may be + truncated during preprocessing. + truncate_strategy: How to truncate long system prompts. Options: + - 'warn': Emit a warning but do not truncate (default). + - 'head_tail': Keep the first and last portions, drop the middle. + - 'head_only': Keep only the beginning of the system prompt. + system_prompt_priority_weight: Weight for system prompt tokens in the + attention mask. Higher values give the model stronger attention to + system instructions. Default: 1.0 (no change). Range: [0.5, 2.0]. + """ + + max_system_prompt_tokens: int = RECOMMENDED_MAX_SYSTEM_PROMPT_TOKENS + truncate_strategy: str = 'warn' + system_prompt_priority_weight: float = 1.0 + + def __post_init__(self): + if self.truncate_strategy not in ('warn', 'head_tail', 'head_only'): + raise ValueError( + f'Unknown truncate_strategy: {self.truncate_strategy!r}. ' + "Must be 'warn', 'head_tail', or 'head_only'." ) - else: - return self.sampler.sample( # pytype: disable=wrong-arg-types - prompt_text, - images=images, - sampling=sampling, - max_new_tokens=max_new_tokens, - rng=rng, - return_state=True, - last_state=last_state, - stream=bool(stream), + if not 0.5 <= self.system_prompt_priority_weight <= 2.0: + raise ValueError( + 'system_prompt_priority_weight must be between 0.5 and 2.0, ' + f'got {self.system_prompt_priority_weight}' ) - def chat( - self, - prompt: str | dialog.Conversation, - *, - images: list[np.ndarray | Image.Image] | UInt8['N? H W C'] | None = None, # pyrefly: ignore[not-a-type] - audio: list[np.ndarray] | None = None, - audio_lengths: list[int] | None = None, - sampling: _sampling.SamplingMethod | None = None, - rng: PRNGKeyLike | None = None, - max_new_tokens: int | None = None, - multi_turn: bool | None = None, - print_stream: bool | dialog.Stream | None = None, - is_legacy_tool_answer: bool = False, - sharding: kd.sharding.ShardingTree | None = None, # pyrefly: ignore[not-a-type] - ) -> str: - """Samples a string from the model. - - The API always expects new gemma format tokens (``<|image|>``, - ``<|audio|>``, etc.). The ``dialog`` library automatically - converts to the correct format for the underlying model. - - Example: - - ```python - # Text-only (all Gemma versions): - output = sampler.chat('Write a poem about cats.') - - # With images (Gemma 4 or Gemma 3): - output = sampler.chat( - 'Describe this image <|image|>.', - images=[image1], - ) - - # With audio (Gemma 4): - output = sampler.chat( - 'Transcribe this audio <|audio|>.', - audio=[audio_array], - ) - ``` + def check_system_prompt_length(self, num_tokens: int) -> list[str]: + """Check system prompt length and return warnings if needed. Args: - prompt: Prompt to sample from. Can be a single string or a - `dialog.Conversation` object. - images: Images for the prompt. For Gemma 4: list of raw numpy arrays or - PIL Images (variable aspect ratio). For Gemma 2/3: a batched uint8 - array. - audio: List of audio arrays (Gemma 4 only). - audio_lengths: List of audio lengths (Gemma 4 only). - sampling: Sampling method to use. If given, will override the default - sampling method. - rng: Seed to use for the sampling method. If `None`, a random seed is - used. Can be a seed `int` or a `jax.random.PRNGKey` object. - max_new_tokens: If given, will stop sampling after this many tokens. Used - for quicker iterations when debugging. By default, sample until the - end-of-turn token is found, or until the `max_out_length` buffer is - filled. - multi_turn: If `True`, reuse the previous turns as context. Overrides the - `multi_turn` attribute. - print_stream: If `True`, will print the sampled output as it is generated. - Overrides the `print_stream` attribute. - is_legacy_tool_answer: When `True`, indicates that the model has emitted - `` rather than `<|tool_response>`, thus this needs to be corrected. - (this is an internal variable that should never be explicitly set). - sharding: Sharding tree (Gemma 4 only). + num_tokens: Number of tokens in the system prompt. Returns: - The sampled output. + List of warning messages (empty if no issues). """ - if multi_turn is None: - multi_turn = self.multi_turn - stream = self.initialize_stream(print_stream) - - if not multi_turn: - # Non-multi-turn, erase the previous conversations. - object.__setattr__(self, 'last_state', None) - object.__setattr__(self, 'turns', []) - - # --- Unified prompt formatting via `dialog` library --- - if isinstance(prompt, str): - if _has_legacy_gemma3_format(prompt): - if self._is_gemma4: - raise ValueError( - 'Detected deprecated Gemma 3 format tokens (e.g. ' - ') in the prompt, but the model ' - 'is Gemma 4 which uses <|image|>, <|audio|>, etc. ' - 'Please use Gemma 4 format tokens instead.' - ) - else: - warnings.warn( - 'Detected deprecated Gemma 3 format tokens (e.g.' - ' ) in the prompt, but the api expects Gemma 4' - ' format tokens. The legacy format is deprecated and will be' - ' removed in a future release.', - DeprecationWarning, - stacklevel=2, - ) - prompt = dialog.Conversation(dialog.User(prompt)) - else: - if not self._is_gemma4: - prompt = prompt.replace('<|image|>', '') - if prompt.find('<|audio|>') != -1: - raise ValueError( - 'Audio input is not supported for non-Gemma4 models.' - ) - prompt = dialog.Conversation(dialog.User(prompt)) - elif not isinstance(prompt, dialog.Conversation): - raise TypeError(f'Unsupported prompt type: {type(prompt)}') - - last_state = self.last_state - if is_legacy_tool_answer: - # This means the previous model turn ended with an EOS token rather than - # the expected `<|tool_response>` - last_state = _remove_eos_token(last_state, tokenizer=self.tokenizer) - - prompt_text = prompt.as_text( - format=self.tokenizer.FORMAT, - add_tool_response_tag_after_call=not is_legacy_tool_answer, - ) - - # --- System prompt length checking --- - # Check if the system prompt portion of the conversation is too long. - # This helps prevent degraded instruction following for long system prompts. - if self.tokenizer is not None: - try: - prompt_tokens = self.tokenizer.encode(prompt_text, add_bos=False) - num_prompt_tokens = len(prompt_tokens) - # Check if this is the first turn (system prompt is typically in the - # first user message) or if the prompt contains a system prompt. - if len(self.turns) == 0 and num_prompt_tokens > 0: - system_config = _thinking_utils.SystemPromptConfig() - warnings_list = system_config.check_system_prompt_length( - num_prompt_tokens - ) - for w in warnings_list: - warnings.warn(w, UserWarning, stacklevel=2) - except Exception: # pylint: disable=broad-except - # If tokenization fails, skip the check. - pass - - # --- Dispatch to the correct sampler --- - out = self._sample( - prompt_text, - images=images, - audio=audio, - audio_lengths=audio_lengths, - sampling=sampling, - max_new_tokens=max_new_tokens, - rng=rng, - last_state=last_state, - stream=stream, - sharding=sharding, - ) - - # In streaming mode, the output is an iterator, yielding tokens one at a - # time. - if stream: - out = _print_stream(out, stream=stream) - assert isinstance(out, _sampler.SamplerOutput) # For pytype. - assert isinstance(out.text, str) # For pytype. - - # TODO(epot): Remove the end-of-turn token. - - # Save the raw turns text (unformatted). - # Only save the user turn after the sampling has successfully finished. - self.turns.append(_template.Prompt(prompt_text)) - self.turns.append(_template.Response(out.text)) - object.__setattr__(self, 'last_state', out.state) - - # --- Multi-turn context management --- - # Check if context refresh is needed after many turns. - num_turns = len(self.turns) // 2 # Each turn is a prompt + response pair - multi_turn_config = _thinking_utils.MultiTurnConfig() - if _thinking_utils.should_refresh_context( - num_turns, multi_turn_config.context_refresh_threshold - ): - warnings.warn( - f'Conversation has {num_turns} turns. For best results with Gemma 4, ' - 'consider starting a new conversation or re-injecting key system ' - 'prompt instructions. Long conversations may cause context loss.', - UserWarning, - stacklevel=2, + warnings_list = [] + if num_tokens > HARD_MAX_SYSTEM_PROMPT_TOKENS: + warnings_list.append( + f'System prompt has {num_tokens} tokens, which exceeds the hard ' + f'limit of {HARD_MAX_SYSTEM_PROMPT_TOKENS}. Model behavior will be ' + 'unreliable. Consider splitting the prompt into multiple turns or ' + 'using a more concise system prompt.' + ) + elif num_tokens > self.max_system_prompt_tokens: + warnings_list.append( + f'System prompt has {num_tokens} tokens, exceeding the recommended ' + f'maximum of {self.max_system_prompt_tokens}. Instruction following ' + 'may degrade. Consider condensing the system prompt.' ) + return warnings_list - return out.text # pytype: disable=bad-return-type - - def initialize_stream( - self, - stream: dialog.Stream | bool | None, - ) -> dialog.Stream | None: - """Initializes a stream for the sampler.""" - if stream is None: - stream = self.print_stream - - if stream is False: # pylint: disable=g-bool-id-comparison - return None - elif stream is True: # pylint: disable=g-bool-id-comparison - stream = dialog.Stream() - if epy.is_notebook(): - dialog.Model(stream).show() - return stream - elif isinstance(stream, dialog.Stream): - return stream - else: - raise ValueError(f'Unexpected stream type: {type(stream)}') - @property - def conversation(self) -> dialog.Conversation: - """Returns the conversation.""" - return dialog.Conversation(''.join(t.text for t in self.turns)) - - -# Legacy Gemma 3 tokens to detect in user prompts. -_LEGACY_GEMMA3_TOKENS = ('',) - - -def _has_legacy_gemma3_format(prompt: str) -> bool: - """Returns True if the prompt contains legacy Gemma 3 format tokens.""" - return any(token in prompt for token in _LEGACY_GEMMA3_TOKENS) - - -def _remove_eos_token( - state: _sampler_loop.SamplingState, - tokenizer: _tokenizer.Tokenizer, -) -> _sampler_loop.SamplingState: - """Removes the EOS token from the sampling state.""" - cache_info = state.cache_info.set_end_index(state.cache_info.end_index - 1) - return dataclasses.replace( - state, - step=state.step - 1, - # done is True and last_token is EOS => False - # Otherwise, keep the same. - done=state.done ^ (state.last_token == tokenizer.special_tokens.EOS), - last_token_pos=state.last_token_pos - 1, - cache=cache_info.cache, - ) - - -def _print_stream( - out: Iterator[_sampler.SamplerOutput], - *, - stream: dialog.Stream, -) -> _sampler.SamplerOutput: - """Prints the streaming output.""" - text_tokens = [] - - for state in out: - print_(stream, state.text) # pyrefly: ignore[bad-argument-type] - - text_tokens.append(state.text) - if ( - state.text == '' or state.text == '' - ): # Last token is not printed. - continue - out = dataclasses.replace(state, text=''.join(text_tokens)) # pylint: disable=undefined-variable,undefined-loop-variable # pyrefly: ignore[bad-assignment] - return out # pyrefly: ignore[bad-return] - - -def print_( - stream: dialog.Stream, - text: str, -) -> None: - if epy.is_notebook(): - stream.add(text) +@dataclasses.dataclass(frozen=True, kw_only=True) +class MultiTurnConfig: + """Configuration for multi-turn conversation context management. + + Attributes: + context_refresh_threshold: After this many turns, suggest refreshing the + context by re-injecting key system prompt instructions. Set to 0 to + disable. Default: 10. + max_context_tokens: Maximum total context length (system prompt + all + turns). When exceeded, older turns are progressively summarized or + dropped. Set to -1 to disable. Default: -1. + context_priority_decay: Rate at which older turns lose attention priority. + 0.0 = no decay (all turns equal), 1.0 = linear decay. Default: 0.3. + """ + + context_refresh_threshold: int = 10 + max_context_tokens: int = -1 + context_priority_decay: float = 0.3 + + def __post_init__(self): + if self.context_refresh_threshold < 0: + raise ValueError( + 'context_refresh_threshold must be >= 0, ' + f'got {self.context_refresh_threshold}' + ) + if self.max_context_tokens < -1: + raise ValueError( + f'max_context_tokens must be >= -1, got {self.max_context_tokens}' + ) + if not 0.0 <= self.context_priority_decay <= 1.0: + raise ValueError( + 'context_priority_decay must be between 0.0 and 1.0, ' + f'got {self.context_priority_decay}' + ) + + +def truncate_system_prompt( + tokens: list[int], + max_tokens: int, + strategy: str = 'head_tail', +) -> list[int]: + """Truncate a system prompt to fit within a token budget. + + Args: + tokens: The system prompt token IDs. + max_tokens: Maximum number of tokens to keep. + strategy: Truncation strategy ('head_tail' or 'head_only'). + + Returns: + Truncated token list. + """ + if len(tokens) <= max_tokens: + return tokens + + if strategy == 'head_only': + return tokens[:max_tokens] + elif strategy == 'head_tail': + # Keep first 60% and last 40% of the budget + head_size = int(max_tokens * 0.6) + tail_size = max_tokens - head_size + return tokens[:head_size] + tokens[-tail_size:] else: - print(text, end='', flush=True) + raise ValueError(f'Unknown truncation strategy: {strategy!r}') + + +def compute_turn_priority_weights( + num_turns: int, + system_prompt_length: int, + decay: float = 0.3, +) -> list[float]: + """Compute attention priority weights for multi-turn conversations. + + Later turns get higher priority, with the system prompt getting the + highest priority. The decay parameter controls how quickly older turns + lose priority. + + Args: + num_turns: Number of conversation turns (excluding system prompt). + system_prompt_length: Length of the system prompt in tokens. + decay: Priority decay rate (0.0 = no decay, 1.0 = linear). + + Returns: + List of priority weights, one per turn + system prompt. + """ + weights = [] + + # System prompt always gets the highest weight. + weights.append(1.0 + decay) + + # Earlier turns get lower priority, later turns get higher priority. + for i in range(num_turns): + # Normalize turn index to [0, 1] + if num_turns > 1: + normalized_idx = i / (num_turns - 1) + else: + normalized_idx = 1.0 + # Apply decay: later turns get higher weight + weight = 1.0 + decay * normalized_idx + weights.append(weight) + + return weights + + +def should_refresh_context( + turn_count: int, + refresh_threshold: int, +) -> bool: + """Check if context should be refreshed based on turn count. + + Args: + turn_count: Current number of conversation turns. + refresh_threshold: Number of turns after which to refresh. + + Returns: + True if context should be refreshed. + """ + if refresh_threshold <= 0: + return False + return turn_count > 0 and turn_count % refresh_threshold == 0 From 72a72ee6886aab43d21a01d7eccd7a07f72b3b97 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Fri, 17 Jul 2026 20:02:11 +0530 Subject: [PATCH 13/23] Refactor thinking channel suppression logic Refactor suppression logic for thinking channel start token to use jnp.where for JAX compatibility. --- gemma/gm/text/_sampler_loop.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/gemma/gm/text/_sampler_loop.py b/gemma/gm/text/_sampler_loop.py index 659ea2ec..cc1140f6 100644 --- a/gemma/gm/text/_sampler_loop.py +++ b/gemma/gm/text/_sampler_loop.py @@ -274,11 +274,16 @@ def _sample_step( and jnp.any(state.in_thinking_channel) ) should_suppress = budget_exhausted & in_channel - if should_suppress: - # Suppress the thinking channel start token to prevent re-entering - # the thinking block. The model should now generate the end marker. - begin_channel = self.special_tokens.BEGIN_OF_THINKING_CHANNEL - logits = logits.at[:, begin_channel].set(-jnp.inf) + # 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) From 9ed48ce639c467e0516630a132fa321bc4686b03 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Fri, 17 Jul 2026 20:02:53 +0530 Subject: [PATCH 14/23] Document max_thinking_tokens in _prefill.py Added documentation for max_thinking_tokens parameter. --- gemma/gm/text/_prefill.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/gemma/gm/text/_prefill.py b/gemma/gm/text/_prefill.py index 7e8d8c81..8d88a103 100644 --- a/gemma/gm/text/_prefill.py +++ b/gemma/gm/text/_prefill.py @@ -84,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. From 7496fee0a6b6d65fd947b2b3e20b9dc690a73ae3 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Fri, 17 Jul 2026 20:03:38 +0530 Subject: [PATCH 15/23] Refactor ChatSampler for better model compatibility Refactor ChatSampler to improve functionality and structure. --- gemma/gm/text/_chat_sampler.py | 729 ++++++++++++++++++++++----------- 1 file changed, 491 insertions(+), 238 deletions(-) diff --git a/gemma/gm/text/_chat_sampler.py b/gemma/gm/text/_chat_sampler.py index 388dc8a0..f860179e 100644 --- a/gemma/gm/text/_chat_sampler.py +++ b/gemma/gm/text/_chat_sampler.py @@ -12,272 +12,525 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Thinking channel budget and context management utilities. - -This module provides utilities for: -1. Managing thinking channel token budgets to prevent infinite loops -2. Optimizing long system prompts for better instruction following -3. Improving multi-turn conversation context retention - -These utilities address known issues with Gemma 4 models: -- Infinite thinking loops in the <|channel>thought block -- Degraded instruction following with system prompts >10k tokens -- Context loss in multi-turn conversations -""" - -from __future__ import annotations +"""Chat sampler.""" +from collections.abc import Iterator, Sequence import dataclasses +import functools import warnings -from typing import Sequence - -import jax.numpy as jnp - - -# Recommended maximum system prompt length in tokens. -# Beyond this threshold, instruction following may degrade. -RECOMMENDED_MAX_SYSTEM_PROMPT_TOKENS = 8192 - -# Hard limit for system prompt tokens. Beyond this, behavior is unreliable. -HARD_MAX_SYSTEM_PROMPT_TOKENS = 16384 - -# Default thinking budget (in tokens) for Gemma 4 models. -DEFAULT_THINKING_BUDGET = 8192 - -# Minimum thinking budget. Below this, the model may not have enough tokens -# to complete its reasoning. -MIN_THINKING_BUDGET = 1024 - -@dataclasses.dataclass(frozen=True, kw_only=True) -class ThinkingBudgetConfig: - """Configuration for thinking channel budget management. +import dialog +from etils import epy +# from gemma.gm.data import _functional +from gemma.gm.nn import _transformer_like +from gemma.gm.text import _gemma4_sampler +from gemma.gm.text import _sampler +from gemma.gm.text import _sampler_loop +from gemma.gm.text import _sampling +from gemma.gm.text import _template +from gemma.gm.text import _thinking_utils +from gemma.gm.text import _tokenizer +from gemma.gm.typing import _common +# from gemma.gm.vision import _token_utils +from kauldron import kd +from kauldron.ktyping import UInt8 # pylint: disable=g-multiple-import,g-importing-member +from kauldron.typing import PRNGKeyLike # pylint: disable=g-multiple-import,g-importing-member +import numpy as np +from PIL import Image + + +@dataclasses.dataclass(frozen=True, kw_only=True, eq=False) +class ChatSampler: + """Chat sampler. + + A unified chat sampler that works with all Gemma model versions (2, 3, 3n, + 4). Automatically selects the correct underlying sampler and prompt format + based on the model's tokenizer version. + + ```python + sampler = ChatSampler( + model=model, + params=params, + multi_turn=True, + ) + + output0 = sampler.chat('Write a poem about cats.') + output1 = sampler.chat('And about dogs.') + output2 = sampler.chat('Which one do you prefer?') + ``` + + For Gemma 4 models with multimodal inputs: + + ```python + sampler = ChatSampler( + model=model, + params=params, + multi_turn=True, + ) + + out0 = sampler.chat('Describe this image <|image|>.', images=[img1]) + out1 = sampler.chat('What about this one <|image|>?', images=[img2]) + out2 = sampler.chat('Summarize your observations.') + ``` + + This sampler: + + * Is stateful (the KV-cache state is automatically handled) + * Automatically formats the prompt with turn tags, adds the BOS + (beginning of sequence) token. And filters the end-of-turn tokens from + the output. + * For Gemma 4 models: supports per-turn images (variable aspect ratio) + and audio via the `images` and `audio` arguments. Attributes: + model: Gemma transformer model. + params: Model parameters. + multi_turn: If `True`, reuse the previous turns as context. + print_stream: If `True`, will print the sampled output as it is generated. + tokenizer: Tokenizer. + sampling: Sampling method to use. Default to greedy sampling. + forbidden_tokens: List of tokens that are forbidden to be generated. If + providing `str`, it should map to a single token id in the vocab. + stop_tokens: List of tokens that will stop generation if generated. If + providing `str`, it should map to a single token id in the vocab. + cache_length: Cache length to use. This is the maximum number of tokens the + conversation can have (prompts, answers, images for all turns). Setting + this to a fixed value avoids re-compilation between turns. + max_out_length: Length of the output buffer for a single turn. Static value + used to avoid triggering a jit recompilation. Shouldn't be changed unless + you have a task where the model generates really long outputs. + pad_length: Pad lengths for static shapes (Gemma 4 only). + patch_size: Patch size for vision encoder (Gemma 4 only). + max_soft_tokens: Maximum soft tokens per image (Gemma 4 only). + 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. When exhausted, the sampler forces an exit. Set to -1 - to disable (no limit). - warn_on_budget_exhaustion: If True, emit a warning when the thinking - budget is exhausted, indicating that the model's reasoning may be - truncated. - thinking_budget_safety_margin: Safety margin as a fraction of the thinking - budget. When remaining budget drops below this fraction, start - suppressing non-essential thinking tokens to encourage early exit. - Default: 0.1 (10%). + 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. + multi_turn_refresh_threshold: Number of conversation turns after which a + warning is emitted suggesting context refresh. Set to 0 to disable. + Default: 10. + 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. + turns: Track the conversation. """ - max_thinking_tokens: int = DEFAULT_THINKING_BUDGET - warn_on_budget_exhaustion: bool = True - thinking_budget_safety_margin: float = 0.1 + # TODO(epot): Custom repr to avoid displaying the full weights. + + model: _transformer_like.TransformerLike + params: _common.Params = dataclasses.field(repr=False) + multi_turn: bool = False + print_stream: bool | dialog.Stream = False + tokenizer: _tokenizer.Tokenizer = None # pytype: disable=annotation-type-mismatch + sampling: _sampling.SamplingMethod = dataclasses.field( + default_factory=_sampling.Greedy + ) + forbidden_tokens: Sequence[str | int] | None = None + stop_tokens: Sequence[str | int] | None = None + # TODO(epot): Support and test rolling cache. + # TODO(epot): Add a property to show how much of the cache is used. + cache_length: int | None = 4096 + max_out_length: int = 2048 + + # Gemma 4-specific fields (ignored for non-Gemma4 models). + pad_length: None | int | tuple[int, ...] = (256, 512, 1024) + patch_size: int = 16 + max_soft_tokens: int = 1120 + pooling_kernel_size: int = 3 + audio_sample_rate: int = 16000 + audio_seq_length: int = 750 + max_thinking_tokens: int = -1 + multi_turn_refresh_threshold: int = 10 + + # Internal variables, but exposed for power users. + + # Last state of the sampler. + last_state: _sampler_loop.SamplingState = dataclasses.field( # pytype: disable=annotation-type-mismatch + default=None, repr=False + ) + turns: list[_template.Turn] = dataclasses.field(default_factory=list) def __post_init__(self): - if self.max_thinking_tokens < -1: + if self.turns: raise ValueError( - f'max_thinking_tokens must be >= -1, got {self.max_thinking_tokens}' + 'Currently initializing the sampler with previous conversation is not' + ' supported.' ) - if self.max_thinking_tokens >= 0 and self.max_thinking_tokens < MIN_THINKING_BUDGET: - warnings.warn( - f'max_thinking_tokens={self.max_thinking_tokens} is very low. ' - f'Minimum recommended is {MIN_THINKING_BUDGET}. The model may not ' - 'have enough tokens to complete its reasoning.', - UserWarning, - stacklevel=2, + # No state by default. + object.__setattr__(self, 'last_state', None) + + # Set the tokenizer if not provided. + if self.tokenizer is None: + object.__setattr__(self, 'tokenizer', self._inner_sampler.tokenizer) + + @functools.cached_property + def _is_gemma4(self) -> bool: + """Returns True if the model is a Gemma 4 model.""" + return getattr(self.model, 'INFO', None) is not None and ( + self.model.INFO.tokenizer_version == 4 + ) + + @functools.cached_property + def _inner_sampler(self) -> _sampler.Sampler | _gemma4_sampler.Gemma4Sampler: + """Returns the underlying sampler, auto-detecting the model type.""" + if self._is_gemma4: + return _gemma4_sampler.Gemma4Sampler( + model=self.model, + params=self.params, + tokenizer=self.tokenizer, + sampling=self.sampling, + forbidden_tokens=self.forbidden_tokens, + stop_tokens=self.stop_tokens, + cache_length=self.cache_length, # pyrefly: ignore[bad-argument-type] + max_out_length=self.max_out_length, + pad_length=self.pad_length, + patch_size=self.patch_size, + max_soft_tokens=self.max_soft_tokens, + 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, ) - if not 0.0 <= self.thinking_budget_safety_margin <= 0.5: - raise ValueError( - 'thinking_budget_safety_margin must be between 0.0 and 0.5, ' - f'got {self.thinking_budget_safety_margin}' + else: + return _sampler.Sampler( + model=self.model, + params=self.params, + tokenizer=self.tokenizer, + sampling=self.sampling, + forbidden_tokens=self.forbidden_tokens, + 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. @property - def effective_max_thinking_tokens(self) -> int: - """Returns the effective max thinking tokens (-1 if disabled).""" - return self.max_thinking_tokens - - def should_suppress_thinking(self, remaining: int) -> bool: - """Returns True if thinking tokens should be suppressed (budget nearly exhausted).""" - if self.max_thinking_tokens < 0: - return False - threshold = int(self.max_thinking_tokens * self.thinking_budget_safety_margin) - return remaining <= threshold - + def sampler(self) -> _sampler.Sampler: + """Returns the underlying sampler (for backwards compatibility).""" + inner = self._inner_sampler + if isinstance(inner, _sampler.Sampler): + return inner + raise AttributeError( + 'The `sampler` property is not available for Gemma 4 models. ' + 'Use `gemma4_sampler` instead.' + ) -@dataclasses.dataclass(frozen=True, kw_only=True) -class SystemPromptConfig: - """Configuration for system prompt optimization. - - Attributes: - max_system_prompt_tokens: Maximum recommended system prompt length in - tokens. If exceeded, a warning is emitted and the prompt may be - truncated during preprocessing. - truncate_strategy: How to truncate long system prompts. Options: - - 'warn': Emit a warning but do not truncate (default). - - 'head_tail': Keep the first and last portions, drop the middle. - - 'head_only': Keep only the beginning of the system prompt. - system_prompt_priority_weight: Weight for system prompt tokens in the - attention mask. Higher values give the model stronger attention to - system instructions. Default: 1.0 (no change). Range: [0.5, 2.0]. - """ - - max_system_prompt_tokens: int = RECOMMENDED_MAX_SYSTEM_PROMPT_TOKENS - truncate_strategy: str = 'warn' - system_prompt_priority_weight: float = 1.0 - - def __post_init__(self): - if self.truncate_strategy not in ('warn', 'head_tail', 'head_only'): - raise ValueError( - f'Unknown truncate_strategy: {self.truncate_strategy!r}. ' - "Must be 'warn', 'head_tail', or 'head_only'." + @property + def gemma4_sampler(self) -> _gemma4_sampler.Gemma4Sampler: + """Returns the underlying Gemma4Sampler (Gemma 4 models only).""" + inner = self._inner_sampler + if isinstance(inner, _gemma4_sampler.Gemma4Sampler): + return inner + raise AttributeError( + 'The `gemma4_sampler` property is not available for non-Gemma4 models. ' + 'Use `sampler` instead.' + ) + + def _sample( + self, + prompt_text: str, + *, + images, + audio, + audio_lengths, + sampling, + max_new_tokens, + rng, + last_state, + stream, + sharding, + ): + """Dispatches to the correct underlying sampler.""" + if self._is_gemma4: + return self.gemma4_sampler.sample( + prompt_text, + images=images, + audio=audio, + audio_lengths=audio_lengths, + sampling=sampling, + max_new_tokens=max_new_tokens, + rng=rng, + return_state=True, + last_state=last_state, + sharding=sharding, ) - if not 0.5 <= self.system_prompt_priority_weight <= 2.0: - raise ValueError( - 'system_prompt_priority_weight must be between 0.5 and 2.0, ' - f'got {self.system_prompt_priority_weight}' + else: + return self.sampler.sample( # pytype: disable=wrong-arg-types + prompt_text, + images=images, + sampling=sampling, + max_new_tokens=max_new_tokens, + rng=rng, + return_state=True, + last_state=last_state, + stream=bool(stream), ) - def check_system_prompt_length(self, num_tokens: int) -> list[str]: - """Check system prompt length and return warnings if needed. + def chat( + self, + prompt: str | dialog.Conversation, + *, + images: list[np.ndarray | Image.Image] | UInt8['N? H W C'] | None = None, # pyrefly: ignore[not-a-type] + audio: list[np.ndarray] | None = None, + audio_lengths: list[int] | None = None, + sampling: _sampling.SamplingMethod | None = None, + rng: PRNGKeyLike | None = None, + max_new_tokens: int | None = None, + multi_turn: bool | None = None, + print_stream: bool | dialog.Stream | None = None, + is_legacy_tool_answer: bool = False, + sharding: kd.sharding.ShardingTree | None = None, # pyrefly: ignore[not-a-type] + ) -> str: + """Samples a string from the model. + + The API always expects new gemma format tokens (``<|image|>``, + ``<|audio|>``, etc.). The ``dialog`` library automatically + converts to the correct format for the underlying model. + + Example: + + ```python + # Text-only (all Gemma versions): + output = sampler.chat('Write a poem about cats.') + + # With images (Gemma 4 or Gemma 3): + output = sampler.chat( + 'Describe this image <|image|>.', + images=[image1], + ) + + # With audio (Gemma 4): + output = sampler.chat( + 'Transcribe this audio <|audio|>.', + audio=[audio_array], + ) + ``` Args: - num_tokens: Number of tokens in the system prompt. + prompt: Prompt to sample from. Can be a single string or a + `dialog.Conversation` object. + images: Images for the prompt. For Gemma 4: list of raw numpy arrays or + PIL Images (variable aspect ratio). For Gemma 2/3: a batched uint8 + array. + audio: List of audio arrays (Gemma 4 only). + audio_lengths: List of audio lengths (Gemma 4 only). + sampling: Sampling method to use. If given, will override the default + sampling method. + rng: Seed to use for the sampling method. If `None`, a random seed is + used. Can be a seed `int` or a `jax.random.PRNGKey` object. + max_new_tokens: If given, will stop sampling after this many tokens. Used + for quicker iterations when debugging. By default, sample until the + end-of-turn token is found, or until the `max_out_length` buffer is + filled. + multi_turn: If `True`, reuse the previous turns as context. Overrides the + `multi_turn` attribute. + print_stream: If `True`, will print the sampled output as it is generated. + Overrides the `print_stream` attribute. + is_legacy_tool_answer: When `True`, indicates that the model has emitted + `` rather than `<|tool_response>`, thus this needs to be corrected. + (this is an internal variable that should never be explicitly set). + sharding: Sharding tree (Gemma 4 only). Returns: - List of warning messages (empty if no issues). + The sampled output. """ - warnings_list = [] - if num_tokens > HARD_MAX_SYSTEM_PROMPT_TOKENS: - warnings_list.append( - f'System prompt has {num_tokens} tokens, which exceeds the hard ' - f'limit of {HARD_MAX_SYSTEM_PROMPT_TOKENS}. Model behavior will be ' - 'unreliable. Consider splitting the prompt into multiple turns or ' - 'using a more concise system prompt.' - ) - elif num_tokens > self.max_system_prompt_tokens: - warnings_list.append( - f'System prompt has {num_tokens} tokens, exceeding the recommended ' - f'maximum of {self.max_system_prompt_tokens}. Instruction following ' - 'may degrade. Consider condensing the system prompt.' - ) - return warnings_list - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class MultiTurnConfig: - """Configuration for multi-turn conversation context management. - - Attributes: - context_refresh_threshold: After this many turns, suggest refreshing the - context by re-injecting key system prompt instructions. Set to 0 to - disable. Default: 10. - max_context_tokens: Maximum total context length (system prompt + all - turns). When exceeded, older turns are progressively summarized or - dropped. Set to -1 to disable. Default: -1. - context_priority_decay: Rate at which older turns lose attention priority. - 0.0 = no decay (all turns equal), 1.0 = linear decay. Default: 0.3. - """ - - context_refresh_threshold: int = 10 - max_context_tokens: int = -1 - context_priority_decay: float = 0.3 - - def __post_init__(self): - if self.context_refresh_threshold < 0: - raise ValueError( - 'context_refresh_threshold must be >= 0, ' - f'got {self.context_refresh_threshold}' - ) - if self.max_context_tokens < -1: - raise ValueError( - f'max_context_tokens must be >= -1, got {self.max_context_tokens}' - ) - if not 0.0 <= self.context_priority_decay <= 1.0: - raise ValueError( - 'context_priority_decay must be between 0.0 and 1.0, ' - f'got {self.context_priority_decay}' + if multi_turn is None: + multi_turn = self.multi_turn + stream = self.initialize_stream(print_stream) + + if not multi_turn: + # Non-multi-turn, erase the previous conversations. + object.__setattr__(self, 'last_state', None) + object.__setattr__(self, 'turns', []) + + # --- Unified prompt formatting via `dialog` library --- + if isinstance(prompt, str): + if _has_legacy_gemma3_format(prompt): + if self._is_gemma4: + raise ValueError( + 'Detected deprecated Gemma 3 format tokens (e.g. ' + ') in the prompt, but the model ' + 'is Gemma 4 which uses <|image|>, <|audio|>, etc. ' + 'Please use Gemma 4 format tokens instead.' + ) + else: + warnings.warn( + 'Detected deprecated Gemma 3 format tokens (e.g.' + ' ) in the prompt, but the api expects Gemma 4' + ' format tokens. The legacy format is deprecated and will be' + ' removed in a future release.', + DeprecationWarning, + stacklevel=2, + ) + prompt = dialog.Conversation(dialog.User(prompt)) + else: + if not self._is_gemma4: + prompt = prompt.replace('<|image|>', '') + if prompt.find('<|audio|>') != -1: + raise ValueError( + 'Audio input is not supported for non-Gemma4 models.' + ) + prompt = dialog.Conversation(dialog.User(prompt)) + elif not isinstance(prompt, dialog.Conversation): + raise TypeError(f'Unsupported prompt type: {type(prompt)}') + + last_state = self.last_state + if is_legacy_tool_answer: + # This means the previous model turn ended with an EOS token rather than + # the expected `<|tool_response>` + last_state = _remove_eos_token(last_state, tokenizer=self.tokenizer) + + prompt_text = prompt.as_text( + format=self.tokenizer.FORMAT, + add_tool_response_tag_after_call=not is_legacy_tool_answer, + ) + + # --- System prompt length checking --- + # Check if the system prompt portion of the conversation is too long. + # This helps prevent degraded instruction following for long system prompts. + if self.tokenizer is not None: + try: + prompt_tokens = self.tokenizer.encode(prompt_text, add_bos=False) + num_prompt_tokens = len(prompt_tokens) + # Check if this is the first turn (system prompt is typically in the + # first user message) or if the prompt contains a system prompt. + if len(self.turns) == 0 and num_prompt_tokens > 0: + system_config = _thinking_utils.SystemPromptConfig() + warnings_list = system_config.check_system_prompt_length( + num_prompt_tokens + ) + for w in warnings_list: + warnings.warn(w, UserWarning, stacklevel=2) + except (ValueError, TypeError, KeyError): # Tokenization failures. + pass + + # --- Dispatch to the correct sampler --- + out = self._sample( + prompt_text, + images=images, + audio=audio, + audio_lengths=audio_lengths, + sampling=sampling, + max_new_tokens=max_new_tokens, + rng=rng, + last_state=last_state, + stream=stream, + sharding=sharding, + ) + + # In streaming mode, the output is an iterator, yielding tokens one at a + # time. + if stream: + out = _print_stream(out, stream=stream) + assert isinstance(out, _sampler.SamplerOutput) # For pytype. + assert isinstance(out.text, str) # For pytype. + + # TODO(epot): Remove the end-of-turn token. + + # Save the raw turns text (unformatted). + # Only save the user turn after the sampling has successfully finished. + self.turns.append(_template.Prompt(prompt_text)) + self.turns.append(_template.Response(out.text)) + object.__setattr__(self, 'last_state', out.state) + + # --- Multi-turn context management --- + # Check if context refresh is needed after many turns. + num_turns = len(self.turns) // 2 # Each turn is a prompt + response pair + if _thinking_utils.should_refresh_context( + num_turns, self.multi_turn_refresh_threshold + ): + warnings.warn( + f'Conversation has {num_turns} turns. For best results with Gemma 4, ' + 'consider starting a new conversation or re-injecting key system ' + 'prompt instructions. Long conversations may cause context loss.', + UserWarning, + stacklevel=2, ) - -def truncate_system_prompt( - tokens: list[int], - max_tokens: int, - strategy: str = 'head_tail', -) -> list[int]: - """Truncate a system prompt to fit within a token budget. - - Args: - tokens: The system prompt token IDs. - max_tokens: Maximum number of tokens to keep. - strategy: Truncation strategy ('head_tail' or 'head_only'). - - Returns: - Truncated token list. - """ - if len(tokens) <= max_tokens: - return tokens - - if strategy == 'head_only': - return tokens[:max_tokens] - elif strategy == 'head_tail': - # Keep first 60% and last 40% of the budget - head_size = int(max_tokens * 0.6) - tail_size = max_tokens - head_size - return tokens[:head_size] + tokens[-tail_size:] - else: - raise ValueError(f'Unknown truncation strategy: {strategy!r}') - - -def compute_turn_priority_weights( - num_turns: int, - system_prompt_length: int, - decay: float = 0.3, -) -> list[float]: - """Compute attention priority weights for multi-turn conversations. - - Later turns get higher priority, with the system prompt getting the - highest priority. The decay parameter controls how quickly older turns - lose priority. - - Args: - num_turns: Number of conversation turns (excluding system prompt). - system_prompt_length: Length of the system prompt in tokens. - decay: Priority decay rate (0.0 = no decay, 1.0 = linear). - - Returns: - List of priority weights, one per turn + system prompt. - """ - weights = [] - - # System prompt always gets the highest weight. - weights.append(1.0 + decay) - - # Earlier turns get lower priority, later turns get higher priority. - for i in range(num_turns): - # Normalize turn index to [0, 1] - if num_turns > 1: - normalized_idx = i / (num_turns - 1) + return out.text # pytype: disable=bad-return-type + + def initialize_stream( + self, + stream: dialog.Stream | bool | None, + ) -> dialog.Stream | None: + """Initializes a stream for the sampler.""" + if stream is None: + stream = self.print_stream + + if stream is False: # pylint: disable=g-bool-id-comparison + return None + elif stream is True: # pylint: disable=g-bool-id-comparison + stream = dialog.Stream() + if epy.is_notebook(): + dialog.Model(stream).show() + return stream + elif isinstance(stream, dialog.Stream): + return stream else: - normalized_idx = 1.0 - # Apply decay: later turns get higher weight - weight = 1.0 + decay * normalized_idx - weights.append(weight) - - return weights - - -def should_refresh_context( - turn_count: int, - refresh_threshold: int, -) -> bool: - """Check if context should be refreshed based on turn count. + raise ValueError(f'Unexpected stream type: {type(stream)}') - Args: - turn_count: Current number of conversation turns. - refresh_threshold: Number of turns after which to refresh. - - Returns: - True if context should be refreshed. - """ - if refresh_threshold <= 0: - return False - return turn_count > 0 and turn_count % refresh_threshold == 0 + @property + def conversation(self) -> dialog.Conversation: + """Returns the conversation.""" + return dialog.Conversation(''.join(t.text for t in self.turns)) + + +# Legacy Gemma 3 tokens to detect in user prompts. +_LEGACY_GEMMA3_TOKENS = ('',) + + +def _has_legacy_gemma3_format(prompt: str) -> bool: + """Returns True if the prompt contains legacy Gemma 3 format tokens.""" + return any(token in prompt for token in _LEGACY_GEMMA3_TOKENS) + + +def _remove_eos_token( + state: _sampler_loop.SamplingState, + tokenizer: _tokenizer.Tokenizer, +) -> _sampler_loop.SamplingState: + """Removes the EOS token from the sampling state.""" + cache_info = state.cache_info.set_end_index(state.cache_info.end_index - 1) + return dataclasses.replace( + state, + step=state.step - 1, + # done is True and last_token is EOS => False + # Otherwise, keep the same. + done=state.done ^ (state.last_token == tokenizer.special_tokens.EOS), + last_token_pos=state.last_token_pos - 1, + cache=cache_info.cache, + ) + + +def _print_stream( + out: Iterator[_sampler.SamplerOutput], + *, + stream: dialog.Stream, +) -> _sampler.SamplerOutput: + """Prints the streaming output.""" + text_tokens = [] + + for state in out: + print_(stream, state.text) # pyrefly: ignore[bad-argument-type] + + text_tokens.append(state.text) + if ( + state.text == '' or state.text == '' + ): # Last token is not printed. + continue + out = dataclasses.replace(state, text=''.join(text_tokens)) # pylint: disable=undefined-variable,undefined-loop-variable # pyrefly: ignore[bad-assignment] + return out # pyrefly: ignore[bad-return] + + +def print_( + stream: dialog.Stream, + text: str, +) -> None: + if epy.is_notebook(): + stream.add(text) + else: + print(text, end='', flush=True) From 62fcaca6352597379355f4fcdcc784508b232187 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Fri, 17 Jul 2026 20:04:06 +0530 Subject: [PATCH 16/23] Remove unused functions from _thinking_utils.py Removed unused functions for truncating system prompts and computing turn priority weights. --- gemma/gm/text/_thinking_utils.py | 70 -------------------------------- 1 file changed, 70 deletions(-) diff --git a/gemma/gm/text/_thinking_utils.py b/gemma/gm/text/_thinking_utils.py index 388dc8a0..48938d8d 100644 --- a/gemma/gm/text/_thinking_utils.py +++ b/gemma/gm/text/_thinking_utils.py @@ -29,9 +29,6 @@ import dataclasses import warnings -from typing import Sequence - -import jax.numpy as jnp # Recommended maximum system prompt length in tokens. @@ -197,73 +194,6 @@ def __post_init__(self): ) -def truncate_system_prompt( - tokens: list[int], - max_tokens: int, - strategy: str = 'head_tail', -) -> list[int]: - """Truncate a system prompt to fit within a token budget. - - Args: - tokens: The system prompt token IDs. - max_tokens: Maximum number of tokens to keep. - strategy: Truncation strategy ('head_tail' or 'head_only'). - - Returns: - Truncated token list. - """ - if len(tokens) <= max_tokens: - return tokens - - if strategy == 'head_only': - return tokens[:max_tokens] - elif strategy == 'head_tail': - # Keep first 60% and last 40% of the budget - head_size = int(max_tokens * 0.6) - tail_size = max_tokens - head_size - return tokens[:head_size] + tokens[-tail_size:] - else: - raise ValueError(f'Unknown truncation strategy: {strategy!r}') - - -def compute_turn_priority_weights( - num_turns: int, - system_prompt_length: int, - decay: float = 0.3, -) -> list[float]: - """Compute attention priority weights for multi-turn conversations. - - Later turns get higher priority, with the system prompt getting the - highest priority. The decay parameter controls how quickly older turns - lose priority. - - Args: - num_turns: Number of conversation turns (excluding system prompt). - system_prompt_length: Length of the system prompt in tokens. - decay: Priority decay rate (0.0 = no decay, 1.0 = linear). - - Returns: - List of priority weights, one per turn + system prompt. - """ - weights = [] - - # System prompt always gets the highest weight. - weights.append(1.0 + decay) - - # Earlier turns get lower priority, later turns get higher priority. - for i in range(num_turns): - # Normalize turn index to [0, 1] - if num_turns > 1: - normalized_idx = i / (num_turns - 1) - else: - normalized_idx = 1.0 - # Apply decay: later turns get higher weight - weight = 1.0 + decay * normalized_idx - weights.append(weight) - - return weights - - def should_refresh_context( turn_count: int, refresh_threshold: int, From dfa2d221aefdf2eaa6a23ed531ff6e07bb11de04 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Thu, 30 Jul 2026 18:57:09 +0530 Subject: [PATCH 17/23] Enhance docstring for Tokenizer API Expanded the docstring to provide detailed information about the Tokenizer API, including its functionality and specific tokens used in Gemma models. --- gemma/gm/text/_tokenizer.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/gemma/gm/text/_tokenizer.py b/gemma/gm/text/_tokenizer.py index 8d5c57d9..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 From 842c924446d02ed4b5299a7913b6752b47b34682 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Thu, 30 Jul 2026 18:58:37 +0530 Subject: [PATCH 18/23] Update _sampler_loop.py From 16dee63490d2be7e66cd8ad8807438a23b63daaa Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Thu, 30 Jul 2026 18:59:15 +0530 Subject: [PATCH 19/23] Update _prefill.py From 36bbcc058909dcd9fe5c9510dcd59460cbd55186 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Thu, 30 Jul 2026 19:00:06 +0530 Subject: [PATCH 20/23] Update _sampler.py From 2b6ab6d4524f9b2fb3da3acd90ef7a2b161e3cb3 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Thu, 30 Jul 2026 19:00:33 +0530 Subject: [PATCH 21/23] Update _gemma4_sampler.py From 315bb23bbd8eb83c0f2e941313f86a2313736490 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Thu, 30 Jul 2026 19:01:06 +0530 Subject: [PATCH 22/23] Update _chat_sampler.py --- gemma/gm/text/_chat_sampler.py | 38 ---------------------------------- 1 file changed, 38 deletions(-) diff --git a/gemma/gm/text/_chat_sampler.py b/gemma/gm/text/_chat_sampler.py index f860179e..d6b75b88 100644 --- a/gemma/gm/text/_chat_sampler.py +++ b/gemma/gm/text/_chat_sampler.py @@ -28,7 +28,6 @@ from gemma.gm.text import _sampler_loop from gemma.gm.text import _sampling from gemma.gm.text import _template -from gemma.gm.text import _thinking_utils from gemma.gm.text import _tokenizer from gemma.gm.typing import _common # from gemma.gm.vision import _token_utils @@ -109,9 +108,6 @@ class ChatSampler: 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. - multi_turn_refresh_threshold: Number of conversation turns after which a - warning is emitted suggesting context refresh. Set to 0 to disable. - Default: 10. 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. @@ -143,7 +139,6 @@ class ChatSampler: audio_sample_rate: int = 16000 audio_seq_length: int = 750 max_thinking_tokens: int = -1 - multi_turn_refresh_threshold: int = 10 # Internal variables, but exposed for power users. @@ -389,25 +384,6 @@ def chat( add_tool_response_tag_after_call=not is_legacy_tool_answer, ) - # --- System prompt length checking --- - # Check if the system prompt portion of the conversation is too long. - # This helps prevent degraded instruction following for long system prompts. - if self.tokenizer is not None: - try: - prompt_tokens = self.tokenizer.encode(prompt_text, add_bos=False) - num_prompt_tokens = len(prompt_tokens) - # Check if this is the first turn (system prompt is typically in the - # first user message) or if the prompt contains a system prompt. - if len(self.turns) == 0 and num_prompt_tokens > 0: - system_config = _thinking_utils.SystemPromptConfig() - warnings_list = system_config.check_system_prompt_length( - num_prompt_tokens - ) - for w in warnings_list: - warnings.warn(w, UserWarning, stacklevel=2) - except (ValueError, TypeError, KeyError): # Tokenization failures. - pass - # --- Dispatch to the correct sampler --- out = self._sample( prompt_text, @@ -437,20 +413,6 @@ def chat( self.turns.append(_template.Response(out.text)) object.__setattr__(self, 'last_state', out.state) - # --- Multi-turn context management --- - # Check if context refresh is needed after many turns. - num_turns = len(self.turns) // 2 # Each turn is a prompt + response pair - if _thinking_utils.should_refresh_context( - num_turns, self.multi_turn_refresh_threshold - ): - warnings.warn( - f'Conversation has {num_turns} turns. For best results with Gemma 4, ' - 'consider starting a new conversation or re-injecting key system ' - 'prompt instructions. Long conversations may cause context loss.', - UserWarning, - stacklevel=2, - ) - return out.text # pytype: disable=bad-return-type def initialize_stream( From d036cadd1cd31632873b465d715a902ca04080a9 Mon Sep 17 00:00:00 2001 From: Aaditya Bhadane Date: Thu, 30 Jul 2026 19:06:23 +0530 Subject: [PATCH 23/23] Delete gemma/gm/text/_thinking_utils.py --- gemma/gm/text/_thinking_utils.py | 213 ------------------------------- 1 file changed, 213 deletions(-) delete mode 100644 gemma/gm/text/_thinking_utils.py diff --git a/gemma/gm/text/_thinking_utils.py b/gemma/gm/text/_thinking_utils.py deleted file mode 100644 index 48938d8d..00000000 --- a/gemma/gm/text/_thinking_utils.py +++ /dev/null @@ -1,213 +0,0 @@ -# Copyright 2026 DeepMind Technologies Limited. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Thinking channel budget and context management utilities. - -This module provides utilities for: -1. Managing thinking channel token budgets to prevent infinite loops -2. Optimizing long system prompts for better instruction following -3. Improving multi-turn conversation context retention - -These utilities address known issues with Gemma 4 models: -- Infinite thinking loops in the <|channel>thought block -- Degraded instruction following with system prompts >10k tokens -- Context loss in multi-turn conversations -""" - -from __future__ import annotations - -import dataclasses -import warnings - - -# Recommended maximum system prompt length in tokens. -# Beyond this threshold, instruction following may degrade. -RECOMMENDED_MAX_SYSTEM_PROMPT_TOKENS = 8192 - -# Hard limit for system prompt tokens. Beyond this, behavior is unreliable. -HARD_MAX_SYSTEM_PROMPT_TOKENS = 16384 - -# Default thinking budget (in tokens) for Gemma 4 models. -DEFAULT_THINKING_BUDGET = 8192 - -# Minimum thinking budget. Below this, the model may not have enough tokens -# to complete its reasoning. -MIN_THINKING_BUDGET = 1024 - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class ThinkingBudgetConfig: - """Configuration for thinking channel budget management. - - Attributes: - max_thinking_tokens: Maximum number of tokens allowed inside the thinking - channel block. When exhausted, the sampler forces an exit. Set to -1 - to disable (no limit). - warn_on_budget_exhaustion: If True, emit a warning when the thinking - budget is exhausted, indicating that the model's reasoning may be - truncated. - thinking_budget_safety_margin: Safety margin as a fraction of the thinking - budget. When remaining budget drops below this fraction, start - suppressing non-essential thinking tokens to encourage early exit. - Default: 0.1 (10%). - """ - - max_thinking_tokens: int = DEFAULT_THINKING_BUDGET - warn_on_budget_exhaustion: bool = True - thinking_budget_safety_margin: float = 0.1 - - def __post_init__(self): - if self.max_thinking_tokens < -1: - raise ValueError( - f'max_thinking_tokens must be >= -1, got {self.max_thinking_tokens}' - ) - if self.max_thinking_tokens >= 0 and self.max_thinking_tokens < MIN_THINKING_BUDGET: - warnings.warn( - f'max_thinking_tokens={self.max_thinking_tokens} is very low. ' - f'Minimum recommended is {MIN_THINKING_BUDGET}. The model may not ' - 'have enough tokens to complete its reasoning.', - UserWarning, - stacklevel=2, - ) - if not 0.0 <= self.thinking_budget_safety_margin <= 0.5: - raise ValueError( - 'thinking_budget_safety_margin must be between 0.0 and 0.5, ' - f'got {self.thinking_budget_safety_margin}' - ) - - @property - def effective_max_thinking_tokens(self) -> int: - """Returns the effective max thinking tokens (-1 if disabled).""" - return self.max_thinking_tokens - - def should_suppress_thinking(self, remaining: int) -> bool: - """Returns True if thinking tokens should be suppressed (budget nearly exhausted).""" - if self.max_thinking_tokens < 0: - return False - threshold = int(self.max_thinking_tokens * self.thinking_budget_safety_margin) - return remaining <= threshold - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class SystemPromptConfig: - """Configuration for system prompt optimization. - - Attributes: - max_system_prompt_tokens: Maximum recommended system prompt length in - tokens. If exceeded, a warning is emitted and the prompt may be - truncated during preprocessing. - truncate_strategy: How to truncate long system prompts. Options: - - 'warn': Emit a warning but do not truncate (default). - - 'head_tail': Keep the first and last portions, drop the middle. - - 'head_only': Keep only the beginning of the system prompt. - system_prompt_priority_weight: Weight for system prompt tokens in the - attention mask. Higher values give the model stronger attention to - system instructions. Default: 1.0 (no change). Range: [0.5, 2.0]. - """ - - max_system_prompt_tokens: int = RECOMMENDED_MAX_SYSTEM_PROMPT_TOKENS - truncate_strategy: str = 'warn' - system_prompt_priority_weight: float = 1.0 - - def __post_init__(self): - if self.truncate_strategy not in ('warn', 'head_tail', 'head_only'): - raise ValueError( - f'Unknown truncate_strategy: {self.truncate_strategy!r}. ' - "Must be 'warn', 'head_tail', or 'head_only'." - ) - if not 0.5 <= self.system_prompt_priority_weight <= 2.0: - raise ValueError( - 'system_prompt_priority_weight must be between 0.5 and 2.0, ' - f'got {self.system_prompt_priority_weight}' - ) - - def check_system_prompt_length(self, num_tokens: int) -> list[str]: - """Check system prompt length and return warnings if needed. - - Args: - num_tokens: Number of tokens in the system prompt. - - Returns: - List of warning messages (empty if no issues). - """ - warnings_list = [] - if num_tokens > HARD_MAX_SYSTEM_PROMPT_TOKENS: - warnings_list.append( - f'System prompt has {num_tokens} tokens, which exceeds the hard ' - f'limit of {HARD_MAX_SYSTEM_PROMPT_TOKENS}. Model behavior will be ' - 'unreliable. Consider splitting the prompt into multiple turns or ' - 'using a more concise system prompt.' - ) - elif num_tokens > self.max_system_prompt_tokens: - warnings_list.append( - f'System prompt has {num_tokens} tokens, exceeding the recommended ' - f'maximum of {self.max_system_prompt_tokens}. Instruction following ' - 'may degrade. Consider condensing the system prompt.' - ) - return warnings_list - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class MultiTurnConfig: - """Configuration for multi-turn conversation context management. - - Attributes: - context_refresh_threshold: After this many turns, suggest refreshing the - context by re-injecting key system prompt instructions. Set to 0 to - disable. Default: 10. - max_context_tokens: Maximum total context length (system prompt + all - turns). When exceeded, older turns are progressively summarized or - dropped. Set to -1 to disable. Default: -1. - context_priority_decay: Rate at which older turns lose attention priority. - 0.0 = no decay (all turns equal), 1.0 = linear decay. Default: 0.3. - """ - - context_refresh_threshold: int = 10 - max_context_tokens: int = -1 - context_priority_decay: float = 0.3 - - def __post_init__(self): - if self.context_refresh_threshold < 0: - raise ValueError( - 'context_refresh_threshold must be >= 0, ' - f'got {self.context_refresh_threshold}' - ) - if self.max_context_tokens < -1: - raise ValueError( - f'max_context_tokens must be >= -1, got {self.max_context_tokens}' - ) - if not 0.0 <= self.context_priority_decay <= 1.0: - raise ValueError( - 'context_priority_decay must be between 0.0 and 1.0, ' - f'got {self.context_priority_decay}' - ) - - -def should_refresh_context( - turn_count: int, - refresh_threshold: int, -) -> bool: - """Check if context should be refreshed based on turn count. - - Args: - turn_count: Current number of conversation turns. - refresh_threshold: Number of turns after which to refresh. - - Returns: - True if context should be refreshed. - """ - if refresh_threshold <= 0: - return False - return turn_count > 0 and turn_count % refresh_threshold == 0 -