ziv-ai/embedder-fast-onnx

ziv init --model fast
ziv start

pip install onnxruntime tokenizers huggingface_hub hf_transfer numpy

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="ziv-ai/embedder-fast-onnx",
    repo_type="model",
    local_dir=".ziv/models/embedder-fast-onnx",
    local_dir_use_symlinks=False,
)

import os
import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

model_dir = ".ziv/models/embedder-fast-onnx"

tokenizer = Tokenizer.from_file(os.path.join(model_dir, "tokenizer.json"))
tokenizer.enable_truncation(max_length=256)
tokenizer.enable_padding(pad_token="[PAD]", length=None)

session = ort.InferenceSession(
    os.path.join(model_dir, "model.onnx"),
    providers=["CPUExecutionProvider"]
)

texts = ["How does authentication work?", "Explain the login flow"]
encoded = tokenizer.encode_batch(texts)

input_ids      = np.array([e.ids            for e in encoded], dtype=np.int64)
attention_mask = np.array([e.attention_mask for e in encoded], dtype=np.int64)
token_type_ids = np.array([e.type_ids       for e in encoded], dtype=np.int64)

outputs = session.run(None, {
    "input_ids":      input_ids,
    "attention_mask": attention_mask,
    "token_type_ids": token_type_ids,
})

# Mean pooling
token_embeddings = outputs[0]
mask = attention_mask[..., None].astype(np.float32)
embeddings = (token_embeddings * mask).sum(1) / mask.sum(1).clip(min=1e-9)

# L2 normalize
norms = np.linalg.norm(embeddings, axis=1, keepdims=True).clip(min=1e-12)
embeddings = embeddings / norms

print(embeddings.shape)           # (2, 384)
print(embeddings @ embeddings.T)  # cosine similarity matrix (2x2)

@article{wang2020minilm,
  title={MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers},
  author={Wang, Wenhui and Wei, Furu and Dong, Li and Bao, Hangbo and Yang, Nan and Zhou, Ming},
  journal={arXiv preprint arXiv:2002.10957},
  year={2020}
}