File size: 3,609 Bytes
7dd7a5b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | """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}")
|