feat(index): add IVF+PQ compression flag for the FAISS build

embed/src/pixelrag_embed/index.py

@@ -148,15 +148,21 @@ def build_ivf(
     train_sample: int = 500_000,
     metric: str = "ip",
     gpu_id: int = -1,
+    pq_m: int = 0,
+    pq_nbits: int = 8,
 ):
-    """Build FAISS IVFFlat index.
+    """Build a FAISS IVF index — IVFFlat by default, IVFPQ when pq_m > 0.
 
     Args:
         nlist: number of IVF clusters (default 4096, good for ~30M vectors)
         nprobe: default search nprobe stored in the index
         train_sample: number of vectors to sample for K-means training
         metric: 'ip' (inner product / cosine for L2-normalized vectors) or 'l2'
         gpu_id: GPU to use for training (-1 = CPU only)
+        pq_m: PQ sub-quantizers. 0 keeps the uncompressed IVFFlat index; >0 builds
+            an IVFPQ index that stores pq_m bytes/vector (at nbits=8) instead of
+            dim*4 — ~128x smaller for dim=2048, pq_m=64. Must divide dim evenly.
+        pq_nbits: bits per PQ sub-quantizer (4 or 8), only used when pq_m > 0.
     """
     import faiss
 
@@ -195,7 +201,16 @@ def build_ivf(
     train_data = embeddings[train_indices]
 
     quantizer = faiss.IndexFlatIP(dim) if metric == "ip" else faiss.IndexFlatL2(dim)
-    index = faiss.IndexIVFFlat(quantizer, dim, nlist, metric_type)
+    if pq_m > 0:
+        if dim % pq_m != 0:
+            raise ValueError(
+                f"--pq-m ({pq_m}) must divide the embedding dim ({dim}) evenly"
+            )
+        # Pass metric_type explicitly: faiss.IndexIVFPQ defaults to METRIC_L2,
+        # which is wrong for the IP-normalized embeddings this index uses.
+        index = faiss.IndexIVFPQ(quantizer, dim, nlist, pq_m, pq_nbits, metric_type)
+    else:
+        index = faiss.IndexIVFFlat(quantizer, dim, nlist, metric_type)
 
     if gpu_id >= 0:
         # GPU-accelerated training: move CPU index to GPU, train, move back
@@ -250,6 +265,9 @@ def build_ivf(
         "nlist": nlist,
         "nprobe": nprobe,
         "metric": metric,
+        "index_type": "ivfpq" if pq_m > 0 else "ivfflat",
+        "pq_m": pq_m,
+        "pq_nbits": pq_nbits if pq_m > 0 else None,
         "index_file": index_path,
         "metadata_file": metadata_path,
     }
@@ -419,6 +437,20 @@ def main():
     p_build.add_argument(
         "--nlist", type=int, default=4096, help="Number of IVF clusters (default: 4096)"
     )
+    p_build.add_argument(
+        "--pq-m",
+        type=int,
+        default=0,
+        help="PQ sub-quantizers for IVFPQ compression (0 = uncompressed IVFFlat). "
+        "Must divide the embedding dim evenly; e.g. 64 for dim 2048 → 64 B/vector.",
+    )
+    p_build.add_argument(
+        "--pq-nbits",
+        type=int,
+        default=8,
+        choices=[4, 8],
+        help="Bits per PQ sub-quantizer (default: 8). Only used when --pq-m > 0.",
+    )
     p_build.add_argument(
         "--nprobe",
         type=int,
@@ -515,6 +547,8 @@ def main():
                 train_sample=args.train_sample,
                 metric=args.metric,
                 gpu_id=args.gpu_id,
+                pq_m=args.pq_m,
+                pq_nbits=args.pq_nbits,
             )
     elif args.command == "test":
         test_search(args.index_dir, nprobe=args.nprobe, k=args.k)

index/src/pixelrag_index/pipelines.py

@@ -451,6 +451,10 @@ def _repl(m: re.Match) -> str:
             nlist,
         )
         cmd += ["--nlist", str(nlist)]
+        if index_cfg.get("pq_m"):
+            cmd += ["--pq-m", str(index_cfg["pq_m"])]
+            if index_cfg.get("pq_nbits"):
+                cmd += ["--pq-nbits", str(index_cfg["pq_nbits"])]
     subprocess.run(cmd, check=True)
 
     logger.info("Index built at %s", output)

tests/test_ivfpq_index.py

@@ -0,0 +1,71 @@
+"""IVFPQ compression path for the FAISS index build (issue #153).
+
+`build_ivf(pq_m>0)` must produce a trained, searchable IndexIVFPQ that keeps the
+inner-product metric — `faiss.IndexIVFPQ` defaults to L2, which would silently
+wreck ranking on these L2-normalized embeddings.
+"""
+
+import json
+import sys
+
+import numpy as np
+import pytest
+
+faiss = pytest.importorskip("faiss")
+sys.path.insert(0, "embed/src")
+from pixelrag_embed.index import build_ivf  # noqa: E402
+
+
+def _write_shard(emb_dir, n=1024, dim=32):
+    rng = np.random.default_rng(0)
+    emb = rng.standard_normal((n, dim)).astype(np.float32)
+    emb /= np.linalg.norm(emb, axis=1, keepdims=True)  # L2-normalized (cosine/IP)
+    np.savez(
+        emb_dir / "shard_000.npz",
+        embeddings=emb,
+        article_ids=np.arange(n, dtype=np.int64),
+        tile_indices=np.zeros(n, dtype=np.int32),
+        chunk_indices=np.arange(n, dtype=np.int32),
+        y_offsets=np.zeros(n, dtype=np.int32),
+        tile_heights=np.full(n, 100, dtype=np.int32),
+    )
+    return emb
+
+
+def test_ivfpq_index_is_built_trained_and_searchable(tmp_path):
+    emb_dir = tmp_path / "emb"
+    emb_dir.mkdir()
+    emb = _write_shard(emb_dir, dim=32)
+    out = tmp_path / "out"
+
+    build_ivf(str(emb_dir), str(out), nlist=8, nprobe=8, pq_m=8, pq_nbits=4)
+
+    index = faiss.read_index(str(out / "index.faiss"))
+    assert isinstance(index, faiss.IndexIVFPQ)
+    # The bug in the issue's own sketch: without the metric arg this would be L2.
+    assert index.metric_type == faiss.METRIC_INNER_PRODUCT
+    assert index.ntotal == emb.shape[0]
+    _, ids = index.search(emb[:1], 5)
+    assert ids.shape == (1, 5) and (ids[0] >= 0).all()
+
+    summary = json.loads((out / "summary.json").read_text())
+    assert summary["index_type"] == "ivfpq"
+    assert summary["pq_m"] == 8 and summary["pq_nbits"] == 4
+
+
+def test_pq_m_must_divide_dim(tmp_path):
+    emb_dir = tmp_path / "emb"
+    emb_dir.mkdir()
+    _write_shard(emb_dir, dim=32)
+    with pytest.raises(ValueError, match="divide"):
+        build_ivf(str(emb_dir), str(tmp_path / "out"), nlist=8, pq_m=7)
+
+
+def test_default_stays_ivfflat(tmp_path):
+    emb_dir = tmp_path / "emb"
+    emb_dir.mkdir()
+    _write_shard(emb_dir, dim=32)
+    out = tmp_path / "out"
+    build_ivf(str(emb_dir), str(out), nlist=8)
+    index = faiss.read_index(str(out / "index.faiss"))
+    assert isinstance(index, faiss.IndexIVFFlat)