Skip to content

fix(embedding): remap unsupported EmbeddingGemma token IDs - #2136

Merged
igorls merged 1 commit into
MemPalace:developfrom
fatkobra:fix/2114-embeddinggemma-token-bounds
Aug 3, 2026
Merged

fix(embedding): remap unsupported EmbeddingGemma token IDs#2136
igorls merged 1 commit into
MemPalace:developfrom
fatkobra:fix/2114-embeddinggemma-token-bounds

Conversation

@fatkobra

@fatkobra fatkobra commented Aug 2, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

  • derive the EmbeddingGemma text-model vocabulary boundary from the tokenizer's base vocabulary
  • detect tokenizer IDs that the text-only ONNX embedding table cannot accept
  • remap unsupported IDs to the tokenizer's valid <unk> token before invoking ONNX Runtime
  • leave valid token batches unchanged
  • emit a one-time warning containing the unsupported IDs
  • add an offline regression for <image_soft_token>

Root cause

The packaged EmbeddingGemma tokenizer defines <image_soft_token> as added token ID 262144.

The text-only ONNX export's embedding matrix has 262,144 rows, so valid indices end at 262143.

EmbeddinggemmaONNX.__call__() previously passed the tokenizer output directly to session.run(). When source text contained the literal <image_soft_token>, ONNX Runtime attempted to gather row 262144 and aborted the complete embedding batch.

Implementation

The new helper obtains the text-model boundary with:

tokenizer.get_vocab_size(with_added_tokens=False)

How to test

python3 -m ruff format mempalace/embedding.py tests/test_embeddinggemma.py
python3 -m ruff format --check .
python3 -m ruff check .
python3 -m pytest tests/test_embeddinggemma.py \
  -k "out_of_range_added_token_is_remapped_to_unknown" \
  -q \
  -rs
python3 -m pytest tests/test_embeddinggemma.py -q -rs
python3 -m pytest \
  tests/test_embeddinggemma.py \
  tests/test_embedding.py \
  tests/test_embedding_wrapper.py \
  -q \
  -rs
python -m pytest tests/ -v

Checklist

  • Tests pass (python -m pytest tests/ -v)
  • No hardcoded paths
  • Linter passes (ruff check .)

@igorls
igorls merged commit efa1820 into MemPalace:develop Aug 3, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants