#!/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")