Skip to content

Fix custom tokens on Gemma3n/Gemma4 and the Gemma3 tool end token - #819

Open
Nas01010101 wants to merge 2 commits into
google-deepmind:mainfrom
Nas01010101:fix-gemma3n-custom-tokens
Open

Nas01010101 wants to merge 2 commits into
google-deepmind:mainfrom
Nas01010101:fix-gemma3n-custom-tokens

Conversation

@Nas01010101

Copy link
Copy Markdown

Fixes #817. Fixes #818.

Two commits:

  1. Fix custom token lookup for Gemma3n and Gemma4. _add_custom_tokens still checks CUSTOM + i first. If that piece isn't <unused{i}>, it looks the piece up by name. The name lookup builds a dict over the vocab only on that miss, so Gemma2 and Gemma3 do no extra work. This commit also adds CUSTOM = 6 to _Gemma4SpecialTokens, and fixes the "voab" typo in the error message.
  2. Remove Gemma3 tool response end token. This removes BEGIN_OF_TOOL_RESPONSE = 50 from _Gemma3SpecialTokens (id 50 is <unused44>), and Sampler only adds the token to end_tokens when the tokenizer defines it. The same guard also fixes the Gemma2 AttributeError from Issue: Gemma2 sampling crashes with AttributeError when Sampler builds end_tokens (BEGIN_OF_TOOL_RESPONSE) #568, which overlaps with Gemma2 sampling crashes with AttributeError when Sampler builds end_tokens (BEGIN_OF_TOOL_RESPONSE) (#568) #569, fix: guard BEGIN_OF_TOOL_RESPONSE for Gemma2 and fix f-string in _normalize_token #590 and Support non-tool use models with Sampler #707.

Tests

  • _tokenizer_test.py: added test_gemma3_custom_tokens, test_gemma3n_custom_tokens and test_gemma4_custom_tokens.
  • _sampler_test.py: added test_end_tokens_exclude_custom_tokens (Gemma3 with custom_tokens={44: ...}) and test_end_tokens_gemma2.
  • If either fix is reverted, its new tests fail.
  • I ran the CI pytest command locally (same --ignores). On main: 392 passed, 3 failed. On this branch: 397 passed, with the same 3 failures (peft/_lora_test.py and _quantization_test.py, jax._src.prng missing in my local jax).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant