google-deepmind / google-deepmind/gemma
Issue: Gemma2 sampling crashes with AttributeError when Sampler builds end_tokens (BEGIN_OF_TOOL_RESPONSE)
- Dominant language
- Python
- Stars
- 5.7k
- Forks
- 1k
- Avg merge
- 10h 33m
- Merged PRs (30d)
- 2
Description
## 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:
```text
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`
Contributor guide
Assessment
This issue has not been assessed yet.