honzajara commited on
Commit
1aa71cb
Β·
verified Β·
1 Parent(s): 30e1d21

Add embeddinggemma-300m

Browse files

Adds **google/embeddinggemma-300m** as a built-in model, usable as `ts/embeddinggemma-300m`.

> **Blocked on typesense/typesense#3038.** This model needs the `gemma` tokenizer type added there; until that ships, `model_type: "gemma"` is rejected by `validate_and_init_local_model`. Opening this now so both halves can be reviewed together.

EmbeddingGemma is a 300M-parameter multilingual model trained on 100+ languages, at 768 dims β€” a third of the size of the `multilingual-e5-large` / `bge-m3` class and the smallest capable multilingual option in the list.

### Files

| File | Notes |
|---|---|
| `embeddinggemma-300m/model.onnx` | exported from `google/embeddinggemma-300m` |
| `embeddinggemma-300m/model.onnx_data` | external weights, 1.23 GB |
| `embeddinggemma-300m/tokenizer.model` | the Gemma SentencePiece vocab, unmodified |
| `embeddinggemma-300m/config.json` | `model_type: gemma`, with the model's own prompts as prefixes |

`.gitattributes` gains an LFS rule for `embeddinggemma-300m/model.onnx_data`, matching the existing `multilingual-e5-large/model.onnx_data` line.

### About the export

EmbeddingGemma's sentence embedding is not the raw token output. Its pipeline is mean pooling β†’ `Dense(768β†’3072)` β†’ `Dense(3072β†’768)` β†’ L2 normalise, and the Dense layers come *after* pooling, so they cannot be left to `TextEmbedder::mean_pooling`.

The graph therefore ends with the whole pipeline baked in, then broadcasts the finished 768-d vector across the sequence axis. Mean pooling over the attention mask returns that vector unchanged, so Typesense gets the real sentence embedding without any change to the pooling code. This is the same approach used for `bge-m3` in #17.

One compatibility note for whoever regenerates this: `torch.onnx.export` produces IR version 10, and the pinned onnxruntime (`rel-1.14.1`) accepts at most IR 8. Setting `ir_version = 8` before saving is enough β€” opset 18 itself is fine.

### Prefixes

EmbeddingGemma ships asymmetric prompts. Mapped onto Typesense's prefixes, with the trailing space included since prefixes are concatenated raw:

- `indexing_prefix`: `"title: none | text: "` (the model's `document` prompt)
- `query_prefix`: `"task: search result | query: "` (the model's `query` prompt)

### Testing

The ONNX was loaded into `typesense/typesense:31.0.rc14`: the session builds, `num_dim` resolves to 768, documents index, and nothing is logged as an error. That build has no `gemma` tokenizer yet, so the full path was verified by feeding the ids produced by the compiled `GemmaTokenizer` from #3038 through this ONNX and applying `TextEmbedder::mean_pooling`:

- **cosine 1.000000** against `sentence-transformers` on 15 inputs β€” Latin, Chinese, Cyrillic, Arabic and Japanese text, emoji, mixed case, whitespace, numerics, and a 3603-token input exercising truncation.
- **cosine 1.000000** for a padded batch of 8, exercising `batch_encode`.
- **cosine 1.000000** on the query path, comparing `query_prefix + q` against `prompt_name="query"`.

The document path lands at 0.991–0.997 rather than 1.000000. That is not the export: `src/index.cpp:8348` appends `" "` after each field value, and removing that single trailing space gives exactly 1.000000. It applies to every model in this repo equally, so I have left it alone.

### One note for maintainers

`data_md5` is required for any model with external weights β€” without it `download_public_model` re-downloads the weights on every start β€” but the README's config table does not list it. Same note as on #17.

πŸ€– Generated with [Claude Code](https://claude.com/claude-code)

https://claude.ai/code/session_01A9ydwqcWZWFG2bLUnDLt8E

.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zst filter=lfs diff=lfs merge=lfs -text
34
  *tfevents* filter=lfs diff=lfs merge=lfs -text
35
  multilingual-e5-large/model.onnx_data filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zst filter=lfs diff=lfs merge=lfs -text
34
  *tfevents* filter=lfs diff=lfs merge=lfs -text
35
  multilingual-e5-large/model.onnx_data filter=lfs diff=lfs merge=lfs -text
36
+ embeddinggemma-300m/model.onnx_data filter=lfs diff=lfs merge=lfs -text
embeddinggemma-300m/config.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "model_md5": "f6e4aadaf1c8578e8df09bb597cd1fbb",
3
+ "vocab_file_name": "tokenizer.model",
4
+ "vocab_md5": "00d2276cbec4474f6cf3df98fbc18cbb",
5
+ "data_md5": "7cfd839859507c1b12041a6a0634ad19",
6
+ "model_type": "gemma",
7
+ "indexing_prefix": "title: none | text: ",
8
+ "query_prefix": "task: search result | query: "
9
+ }
embeddinggemma-300m/model.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:220f4b52a42d756876cd10587031842cde2bacec674851fb0f795e0a1c53cd95
3
+ size 7424960
embeddinggemma-300m/model.onnx_data ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b46e4dd4313141059b9eb6b5129e6752bfade7bd6bfde8f1460249d274a5b50d
3
+ size 1230326784
embeddinggemma-300m/tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1299c11d7cf632ef3b4e11937501358ada021bbdf7c47638d13c0ee982f2e79c
3
+ size 4689074