google-deepmind / google-deepmind/gemma

[Examples] Avoid hardcoded label token IDs in classification.py in examples

Open
#496 3 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.