chore: conversion and test script, make special tokens in ZeroEncoder configurable

scripts/convert-to-safensors.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env -S uv run --script
+#
+# /// script
+# requires-python = ">=3.12"
+# dependencies = ["safetensors", "numpy"]
+# ///
+
+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
+
+save_file({"embedding.weight": W}, "model.safetensors")

scripts/test-safentesor.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env -S uv run --script
+#
+# /// script
+# requires-python = ">=3.12"
+# dependencies = ["statembed-py", "numpy", "tokenizers"]
+# ///
+
+import sys
+from pathlib import Path
+
+from statembed_py import StaticEmbedding
+
+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}")

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"):
+    def __init__(self, model_dir, variant="int8", add_special_tokens=True):
         d = Path(model_dir)
         self.config = json.loads((d / "config.json").read_text())
         pre = self.config["preproc"]
@@ -45,6 +45,10 @@ 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)
@@ -75,7 +79,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)):
+        for i, enc in enumerate(self.tokenizer.encode_batch(texts, add_special_tokens=self.add_special_tokens)):
             out[i] = self._encode_ids(enc.ids)
         return out