Ship int8 safetensors with per-row scales for FastEmbed

Stores the released int8 rows as and their per-row scales as , so FastEmbed's static path reproduces the released encoder exactly. Restores to the released version; FastEmbed now handles special tokens and sqrt pooling itself.

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:10a13ff3791bf5bbad0c33e5ed3d8d13ed95f0ab92174b5ffdb36026fc7fd965
-size 62509144
+oid sha256:fa798c355d9bc602f8314d7ca411c27467787cee19ad15e25390c2f9aaec37d0
+size 31376776

scripts/convert-to-safensors.py

@@ -9,7 +9,9 @@ import numpy as np
 from safetensors.numpy import save_file
 
 d = np.load("model.npz")
-W = np.ascontiguousarray(d["rows_fp16"])  # (30522, 1024) float16
-assert W.shape == (30522, 1024) and W.dtype == np.float16
+W = np.ascontiguousarray(d["rows_int8"])  # (30522, 1024) int8, the released table
+scale = d["int8_scale"].astype(np.float32)  # (30522,) per-row dequantization scale
+assert W.shape == (30522, 1024) and W.dtype == np.int8 and scale.shape == (30522,)
 
-save_file({"embeddings": W}, "model.safetensors")
+# FastEmbed's static models multiply each token's row by `weights`, which dequantizes it.
+save_file({"embeddings": W, "weights": scale}, "model.safetensors")

scripts/test-safentesor.py

@@ -2,28 +2,30 @@
 #
 # /// script
 # requires-python = ">=3.12"
-# dependencies = ["statembed-py", "numpy", "tokenizers"]
+# 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
 
-from statembed_py import StaticEmbedding
+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(".", variant="fp16", add_special_tokens=False)
-other_emb = model.encode(["hello world"]).tolist()[0]
-
-model = StaticEmbedding(model_dir=".", normalize=True)
-
-emb = model.embed_text("hello world")
-
-import numpy as np
-
-a, b = np.asarray(emb), np.asarray(other_emb)
-print(f"cosine similarity: {float(a @ b):.6f}")
-print(f"max abs diff:      {float(np.abs(a - b).max()):.8f}")
+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")

zero_encoder.py

@@ -28,7 +28,7 @@ class ZeroQueryEncoder:
              int8 was measured quality-free against it (upper bound 0.00013 nDCG@10).
     """
 
-    def __init__(self, model_dir, variant="int8", add_special_tokens=True):
+    def __init__(self, model_dir, variant="int8"):
         d = Path(model_dir)
         self.config = json.loads((d / "config.json").read_text())
         pre = self.config["preproc"]
@@ -45,10 +45,6 @@ class ZeroQueryEncoder:
         else:
             raise ValueError(f"variant must be 'int8' or 'fp16', got {variant!r}")
         self.variant = variant
-        # static embedding models tend to have this
-        # set to False, whereas most other models
-        # default to True
-        self.add_special_tokens = add_special_tokens
 
         self.tokenizer = Tokenizer.from_file(str(d / "tokenizer.json"))
         n = self.tokenizer.get_vocab_size(with_added_tokens=True)
@@ -79,7 +75,7 @@ class ZeroQueryEncoder:
         if isinstance(texts, str):
             texts = [texts]
         out = np.empty((len(texts), self.dim), dtype=np.float32)
-        for i, enc in enumerate(self.tokenizer.encode_batch(texts, add_special_tokens=self.add_special_tokens)):
+        for i, enc in enumerate(self.tokenizer.encode_batch(texts)):
             out[i] = self._encode_ids(enc.ids)
         return out