| """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}") |
| |
| 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}") |
|
|
