google-deepmind / google-deepmind/gemma
[Examples] Avoid hardcoded label token IDs in classification.py in examples
- Dominant language
- Python
- Stars
- 5.7k
- Forks
- 1k
- Avg merge
- 10h 33m
- Merged PRs (30d)
- 2
Description
**I noticed that in "examples/classification.py" maps labels to hardcoded token IDs:**
0 -> 1294 ("No")
1 -> 3553 ("Yes")
While this works for the current tokenizer, it makes the example brittle if:
- the tokenizer vocabulary changes,
- a different Gemma variant is used, or
- the example is adapted to another model.
Additionally, the file contains minor typos:
- "grammaticaly" -> "grammatically"
- "respectivelly" -> "respectively"
**Proposed Solution**
1. Compute the token IDs for "Yes" and "No" dynamically using the tokenizer instance.
2. Add a small safety check to ensure that these labels map to a single token.
3. Fix the spelling errors in the prompt template and comments.
This would make the example more robust while keeping its behavior unchanged.
I would be happy to implement this change.
Contributor guide
Assessment
This issue has not been assessed yet.