Issue: Gemma2 sampling crashes with AttributeError when Sampler builds end_tokens (BEGIN_OF_TOOL_RESPONSE)

[markknoffler]

Description

gm.text.Sampler unconditionally includes tokenizer.special_tokens.BEGIN_OF_TOOL_RESPONSE when building end_tokens for SamplerLoop. The Gemma2 tokenizer’s special-tokens enum (_Gemma2SpecialTokens) does not define BEGIN_OF_TOOL_RESPONSE (that token was introduced in Gemma3 for tool/function calling). As a result, using the Sampler with a Gemma2 model (or any model that uses tokenizer version 2) raises:

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

Location

• File: gemma/gm/text/_sampler.py
• Approximate lines: 344–357 (inside Sampler.sample(), when constructing _sampler_loop.SamplerLoop)
• Relevant tokenizer definitions: gemma/gm/text/_tokenizer.py
  • _Gemma2SpecialTokens (lines 79–99): has EOS, END_OF_TURN, etc.; no BEGIN_OF_TOOL_RESPONSE
  • _Gemma3SpecialTokens (lines 102–128): defines BEGIN_OF_TOOL_RESPONSE = 50

Steps to reproduce

1. Create a Sampler with a model that uses tokenizer version 2 (e.g. Gemma2), or explicitly pass a Gemma2Tokenizer.
2. Call sampler.sample(prompt) (or any code path that builds the SamplerLoop and thus end_tokens).
3. The crash occurs when the code builds end_tokens and accesses self.tokenizer.special_tokens.BEGIN_OF_TOOL_RESPONSE.

Minimal repro: use a dummy model with INFO.tokenizer_version = 2, instantiate Sampler(model=..., params=...), then call sampler.sample("test"). The failure happens before any real model forward pass.

Expected behavior

Sampling with a Gemma2 model (tokenizer version 2) should not crash. The tool-response token should only be included in end_tokens when the tokenizer actually defines it (Gemma3+).

Actual behavior

AttributeError when building end_tokens for Gemma2, because _Gemma2SpecialTokens has no attribute BEGIN_OF_TOOL_RESPONSE.

Environment

• Repository: google-deepmind/gemma (or fork)
• Branch: main
• Python: 3.12+
• Relevant code: gemma/gm/text/_sampler.py, gemma/gm/text/_tokenizer.py

[Ayushhgit]

Opened PR #570 to address this by conditionally including BEGIN_OF_TOOL_RESPONSE only when defined on the tokenizer’s special tokens class. Happy to adjust if a tokenizer_version-based check is preferred.

[pratik376]

replace this:
end_tokens = [
self.tokenizer.special_tokens.EOS,
self.tokenizer.special_tokens.END_OF_TURN,
self.tokenizer.special_tokens.BEGIN_OF_TOOL_RESPONSE,
]
with this:

special = self.tokenizer.special_tokens

end_tokens = [
special.EOS,
special.END_OF_TURN,
]

if hasattr(special, "BEGIN_OF_TOOL_RESPONSE"):
end_tokens.append(special.BEGIN_OF_TOOL_RESPONSE)

[markknoffler]

replace this: end_tokens = [ self.tokenizer.special_tokens.EOS, self.tokenizer.special_tokens.END_OF_TURN, self.tokenizer.special_tokens.BEGIN_OF_TOOL_RESPONSE, ] with this:

special = self.tokenizer.special_tokens

end_tokens = [ special.EOS, special.END_OF_TURN, ]

if hasattr(special, "BEGIN_OF_TOOL_RESPONSE"): end_tokens.append(special.BEGIN_OF_TOOL_RESPONSE)

Opened PR #570 to address this by conditionally including BEGIN_OF_TOOL_RESPONSE only when defined on the tokenizer’s special tokens class. Happy to adjust if a tokenizer_version-based check is preferred.

Hey @Ayushhgit and @pratik376 — appreciate the enthusiasm and the solutions!

Just a heads-up: I actually opened this issue and already submitted a fix in PR #569 a few hours ago. The approach there similarly guards against the missing BEGIN_OF_TOOL_RESPONSE attribute on Gemma2's tokenizer.

@Ayushhgit, your PR #570 looks like it covers the same ground. @pratik376, the inline snippet you shared is essentially the same pattern as what's already in #569.

[Ayushhgit]

@markknoffler Thanks for the heads-up! I didn’t see that PR #569 had already addressed this.
I’ll go ahead and close this one to avoid duplication. Appreciate the clarification.