Upload build_index.py with huggingface_hub

build_index.py

@@ -0,0 +1,73 @@
+"""
+One-time script: encode 500 MS-MARCO docs with dfrokido/bge-large-e8-snap,
+build float and rfsnap normalized index tensors, save to hf_space/data/index.pt.
+
+Run from repo root:
+    python hf_space/build_index.py
+"""
+from __future__ import annotations
+import json
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from sentence_transformers import SentenceTransformer
+
+sys.path.insert(0, str(Path(__file__).parent))
+from e8_utils import nestquant_snap
+
+CORPUS_PATH = Path("runs/msmarco_local_100k/corpus.jsonl")
+MODEL_ID    = "dfrokido/bge-large-e8-snap"
+N_DOCS      = 500
+OUT_PATH    = Path("hf_space/data/index.pt")
+BATCH_SIZE  = 64
+
+
+def main():
+    OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
+
+    print(f"Loading {N_DOCS} docs from {CORPUS_PATH}...")
+    docs = []
+    with open(CORPUS_PATH) as f:
+        for i, line in enumerate(f):
+            if i >= N_DOCS:
+                break
+            docs.append(json.loads(line))
+    doc_ids   = [d["doc_id"] for d in docs]
+    doc_texts = [d["text"]   for d in docs]
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Encoding with {MODEL_ID} on {device}...")
+    model = SentenceTransformer(MODEL_ID, device=device)
+    emb = model.encode(doc_texts, batch_size=BATCH_SIZE,
+                       convert_to_tensor=True, show_progress_bar=True)
+    emb = F.normalize(emb.float().cpu(), p=2, dim=1)
+
+    snap_emb  = nestquant_snap(emb)
+    snap_norm = F.normalize(snap_emb, p=2, dim=1)
+    float_norm = emb  # already normalized
+
+    sizes = {
+        "float32_mb": emb.numel() * 4 / 1e6,
+        "rfsnap_mb":  emb.shape[0] * (emb.shape[1] // 8) * 3 / 1e6,
+    }
+
+    payload = {
+        "doc_ids":    doc_ids,
+        "doc_texts":  doc_texts,
+        "float_norm": float_norm,
+        "snap_norm":  snap_norm,
+        "d_model":    emb.shape[1],
+        "n_docs":     len(docs),
+        "sizes":      sizes,
+    }
+    torch.save(payload, OUT_PATH)
+    print(f"Saved {OUT_PATH}")
+    print(f"  float32 index: {sizes['float32_mb']:.2f} MB")
+    print(f"  rfsnap  index: {sizes['rfsnap_mb']:.2f} MB")
+    print(f"  compression:   {sizes['float32_mb']/sizes['rfsnap_mb']:.1f}x")
+
+
+if __name__ == "__main__":
+    main()