Text-only ONNX (fp32 + int8) of google/embeddinggemma-2 for Vespa

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+onnx/int8/model.onnx_data filter=lfs diff=lfs merge=lfs -text
+onnx/model.onnx_data filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,61 @@
+---
+license: apache-2.0
+pipeline_tag: feature-extraction
+library_name: onnx
+base_model:
+- google/embeddinggemma-2
+- onnx-community/embeddinggemma-2-ONNX
+tags:
+- embedding
+- sentence-similarity
+- vespa
+---
+
+# EmbeddingGemma 2 — text-only ONNX (for Vespa)
+
+Text-only ONNX export of [google/embeddinggemma-2](https://huggingface.co/google/embeddinggemma-2),
+derived from [onnx-community/embeddinggemma-2-ONNX](https://huggingface.co/onnx-community/embeddinggemma-2-ONNX)
+for use with Vespa's [`hugging-face-embedder`](https://docs.vespa.ai/en/rag/embedding.html#huggingface-embedder).
+
+The upstream graph requires `image_features`, `video_features` and `audio_features` inputs that are
+concatenated after the text tokens. These are replaced with empty `[0, 512]` constants, so the
+graph takes only `input_ids` and `attention_mask`. Weights are unchanged. See `scripts/make_text_only.py`.
+
+| File | Precision | Size |
+| --- | --- | --- |
+| `onnx/model.onnx` (+ `model.onnx_data`) | fp32 | 1.08 GB |
+| `onnx/int8/model.onnx` (+ `model.onnx_data`) | int8 (dynamic quantization) | 314 MB |
+
+Outputs: `last_hidden_state` `[batch, seq, 768]` (use mean pooling + L2 normalize) and `sentence_embedding` `[batch, 768]`.
+
+Do **not** use fp16. Per the model card, activations overflow float16.
+
+## Verification
+
+`scripts/verify.py` compares Vespa-style inference (tokenizer.json with special tokens, mean pooling over
+`last_hidden_state`, L2 normalization) against `sentence-transformers` 6.1 / transformers 5.19 in fp32,
+on 11 texts including multilingual, code, and a 2,718-token document (exercises sliding-window attention).
+Token IDs are identical to the reference.
+
+| Variant | Cosine vs. reference (min, 768d / 512d / 256d / 128d) | Same ranking |
+| --- | --- | --- |
+| fp32 | 1.000000 / 1.000000 / 1.000000 / 1.000000 (max abs diff 6e-7) | yes |
+| int8 | 0.999896 / 0.999901 / 0.999910 / 0.999940 | yes |
+
+## Vespa usage
+
+```xml
+<component id="embeddinggemma2" type="hugging-face-embedder">
+    <transformer-model url="https://huggingface.co/vespa-engine/embeddinggemma-2-ONNX/resolve/main/onnx/model.onnx"/>
+    <tokenizer-model url="https://huggingface.co/vespa-engine/embeddinggemma-2-ONNX/resolve/main/tokenizer.json"/>
+    <max-tokens>8192</max-tokens>
+    <pooling-strategy>mean</pooling-strategy>
+    <normalize>true</normalize>
+    <prepend>
+        <query>task: search result | query: </query>
+        <document>title: none | text: </document>
+    </prepend>
+</component>
+```
+
+Supports Matryoshka truncation to 512, 256 and 128 dimensions (re-normalize after truncating).

onnx/int8/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d6ee6b22d9c48d013d9c931d52f7261d73b29cdb4b255264aaa27200d5e8fb1b
+size 488970

onnx/int8/model.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:278a7ff1248c3618e4bd11a607fc54f7bdc7778854230f3956d3f86bd9db4f3b
+size 313724928

onnx/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cbd21b1b6ea4f47ef4e8789d138c27d396d14a369cc7b7eb4855f360d0fac563
+size 421349

onnx/model.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9fd452bfc6916e5a92f080f8dada35598ca4512b21c65897d4c5974d5f9630a2
+size 1084170240

scripts/make_text_only.py

@@ -0,0 +1,45 @@
+"""Make onnx-community/embeddinggemma-2-ONNX text-only for Vespa's hugging-face-embedder.
+
+The upstream graph requires image_features/video_features/audio_features inputs
+(each [num_tokens, 512]) that are concatenated after the text embeddings.
+For text-only inference these are empty, so we replace them with constant
+[0, 512] initializers. The weights are untouched: each variant is written as
+<out>/<variant>/model.onnx referencing model.onnx_data (hard-linked to the
+upstream data file), matching how modelhub/Vespa store external data.
+"""
+
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import onnx
+from onnx import numpy_helper
+
+SRC = Path(sys.argv[1])  # onnx-community snapshot dir
+OUT = Path(sys.argv[2])
+VARIANTS = {"fp32": "model", "int8": "model_quantized", "q4": "model_q4"}
+MEDIA_INPUTS = ["image_features", "video_features", "audio_features"]
+
+for variant, name in VARIANTS.items():
+    dst = OUT / variant
+    dst.mkdir(parents=True, exist_ok=True)
+    m = onnx.load(SRC / "onnx" / f"{name}.onnx", load_external_data=False)
+    for inp in [i for i in m.graph.input if i.name in MEDIA_INPUTS]:
+        dtype = onnx.helper.tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type)
+        width = inp.type.tensor_type.shape.dim[1].dim_value
+        m.graph.input.remove(inp)
+        m.graph.initializer.append(
+            numpy_helper.from_array(np.zeros((0, width), dtype=dtype), inp.name)
+        )
+    for t in m.graph.initializer:
+        for e in t.external_data:
+            if e.key == "location":
+                assert e.value == f"{name}.onnx_data", e.value
+                e.value = "model.onnx_data"
+    onnx.save_model(m, dst / "model.onnx")
+    data = dst / "model.onnx_data"
+    data.unlink(missing_ok=True)
+    os.link(os.path.realpath(SRC / "onnx" / f"{name}.onnx_data"), data)
+    # onnx.checker rejects hard-linked data files; we validate by loading in ORT instead.
+    print(variant, [i.name for i in m.graph.input], [o.name for o in m.graph.output])

scripts/verify.py

@@ -0,0 +1,81 @@
+"""Verify text-only ONNX variants against the sentence-transformers fp32 reference.
+
+ONNX side mimics Vespa's hugging-face-embedder: tokenizer.json via `tokenizers`
+(add_special_tokens=True), batch size 1, mean pooling over last_hidden_state with
+the attention mask, then L2 normalize (and truncate+renormalize for MRL dims).
+"""
+
+import sys
+
+import numpy as np
+import onnxruntime as ort
+import torch
+from sentence_transformers import SentenceTransformer
+from tokenizers import Tokenizer
+
+OUT = sys.argv[1]
+TOKENIZER = sys.argv[2]
+Q = "task: search result | query: "
+D = "title: none | text: "
+
+long_doc = " ".join(
+    f"Paragraph {i}: The aurora borealis appears when charged solar particles collide with "
+    f"gases in Earth's upper atmosphere, producing light in green, red and violet hues."
+    for i in range(80)
+)
+texts = [
+    Q + "What causes the northern lights?",
+    D + "The northern lights are caused by charged particles from the sun.",
+    Q + "Which planet is known as the Red Planet?",
+    D + "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
+    D + "Venus is often called Earth's twin because of its similar size and proximity.",
+    Q + "Hvordan lager man brunost?",
+    D + "Brunost lages ved å koke myse til sukkeret karamelliseres. 日本語のテキストも含む。",
+    "task: code retrieval | query: reverse a linked list in python",
+    D + "def reverse(head):\n    prev = None\n    while head:\n        head.next, prev, head = prev, head, head.next\n    return prev",
+    "a",
+    D + long_doc,
+]
+
+st = SentenceTransformer(
+    "google/embeddinggemma-2",
+    model_kwargs={"torch_dtype": torch.float32},
+    config_kwargs={"vision_config": None, "audio_config": None},
+)
+with torch.no_grad():
+    ref = st.encode(texts, normalize_embeddings=True, batch_size=1, convert_to_numpy=True)
+
+tok = Tokenizer.from_file(TOKENIZER)
+st_ids = [st.tokenize([t])["input_ids"][0].tolist() for t in texts]
+vespa_ids = [tok.encode(t, add_special_tokens=True).ids for t in texts]
+print("token ids identical:", all(a == b for a, b in zip(st_ids, vespa_ids)))
+print("long doc tokens:", len(vespa_ids[-1]), "| sample ids:", vespa_ids[0][:4], "...", vespa_ids[0][-2:])
+
+
+def norm(x):
+    return x / np.linalg.norm(x, axis=-1, keepdims=True)
+
+
+for variant in ["fp32", "int8", "q4"]:
+    sess = ort.InferenceSession(f"{OUT}/{variant}/model.onnx", providers=["CPUExecutionProvider"])
+    embs, sent = [], []
+    for ids in vespa_ids:
+        ii = np.array([ids], dtype=np.int64)
+        am = np.ones_like(ii)
+        lhs, se = sess.run(["last_hidden_state", "sentence_embedding"], {"input_ids": ii, "attention_mask": am})
+        embs.append((lhs[0] * am[0, :, None]).sum(0) / am.sum())
+        sent.append(se[0])
+    embs = norm(np.stack(embs))
+    sent = norm(np.stack(sent))
+    print(f"\n== {variant}")
+    print("  any NaN:", bool(np.isnan(embs).any()))
+    print("  meanpool vs sentence_embedding output, min cos:", float((embs * sent).sum(-1).min()))
+    for dim in [768, 512, 256, 128]:
+        a, b = norm(embs[:, :dim]), norm(ref[:, :dim])
+        cos = (a * b).sum(-1)
+        print(f"  dim {dim}: cos vs ST min={cos.min():.6f} mean={cos.mean():.6f} | max abs diff={np.abs(a - b).max():.2e}")
+    # Retrieval sanity: q0->d1, q2->d3 should rank first among docs; compare score matrices
+    qi, di = [0, 2, 5, 7], [1, 3, 4, 6, 8]
+    s_onnx, s_ref = embs[qi] @ embs[di].T, ref[qi] @ ref[di].T
+    print("  ranking identical to ST:", bool((s_onnx.argsort(1) == s_ref.argsort(1)).all()),
+          "| max score diff:", f"{np.abs(s_onnx - s_ref).max():.2e}")

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d777ef5bdc1aa36227abdfb77c3e49e7b9c892d16e1b6bda41c393504828be4
+size 32170510

tokenizer_config.json

@@ -0,0 +1,57 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": true,
+  "local_files_only": false,
+  "mask_token": "<mask>",
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>",
+    "video_token": "<|video|>"
+  },
+  "pad_token": "<pad>",
+  "padding_side": "right",
+  "processor_class": "EmbeddingGemma2Processor",
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>",
+  "video_token": "<|video|>"
+}