File size: 1,068 Bytes
d06c1cc | 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 | #!/usr/bin/env -S uv run --script
#
# /// script
# requires-python = ">=3.12"
# dependencies = ["safetensors", "numpy", "tokenizers"]
# ///
# Checks that model.safetensors, pooled by the model card's rule, gives the released vectors.
import sys
from pathlib import Path
import numpy as np
from safetensors.numpy import load_file
model_dir = Path(__file__).parents[1]
sys.path.insert(0, str(model_dir))
from zero_encoder import ZeroQueryEncoder
model = ZeroQueryEncoder(model_dir, variant="int8")
t = load_file(model_dir / "model.safetensors")
texts = ["hello world", "data data data pipeline", "🤗 emoji query", ""]
for text, released in zip(texts, model.encode(texts)):
# special and unknown tokens included, a token seen c times weighs sqrt(c)
ids, counts = np.unique(model.tokenizer.encode(text).ids, return_counts=True)
v = (np.sqrt(counts) * t["weights"][ids]) @ t["embeddings"][ids].astype(np.float32)
v /= np.linalg.norm(v)
assert np.abs(v - released).max() < 1e-6, text
print("model.safetensors matches the released int8 encoder")
|