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