Gemma2 sampling crashes with AttributeError when Sampler builds end_tokens (BEGIN_OF_TOOL_RESPONSE) (#568)

[markknoffler]

Fix for Issue #568

Root cause

gm.text.Sampler.sample() builds the SamplerLoop.end_tokens tuple and unconditionally includes:

• self.tokenizer.special_tokens.BEGIN_OF_TOOL_RESPONSE

Gemma2 uses _Gemma2SpecialTokens which does not define BEGIN_OF_TOOL_RESPONSE (tool/function calling tokens are Gemma3+). That unconditional attribute access raises:

AttributeError: type object '_Gemma2SpecialTokens' has no attribute 'BEGIN_OF_TOOL_RESPONSE'

Fix summary

Build end_tokens conditionally:

• Always include EOS and END_OF_TURN
• Include BEGIN_OF_TOOL_RESPONSE only if it exists on the tokenizer’s special_tokens enum (Gemma3+)
• Append any user-provided stop tokens as usual

This keeps tool behavior unchanged for Gemma3+ while preventing Gemma2 from crashing.

Patch sketch

In gemma/gm/text/_sampler.py, replace the hardcoded end_tokens=(...) tuple with logic like:

base_end_tokens = (
    self.tokenizer.special_tokens.EOS,
    self.tokenizer.special_tokens.END_OF_TURN,
)

tool_token = getattr(
    self.tokenizer.special_tokens, "BEGIN_OF_TOOL_RESPONSE", None
)
if tool_token is not None:
  base_end_tokens = (*base_end_tokens, tool_token)

end_tokens = (*base_end_tokens, *self._normalized_stop_tokens)

and pass end_tokens=end_tokens when constructing _sampler_loop.SamplerLoop(...).

Why this is correct

• Gemma2: getattr(..., None) returns None, so the tool token is skipped → no crash.
• Gemma3 / Gemma3n: BEGIN_OF_TOOL_RESPONSE exists, so the token is included → tool sampling behavior preserved.