bitbudget 0.1.0__py3-none-any.whl

bitbudget/__init__.py

@@ -0,0 +1,28 @@
+"""BitBudget: how much retrieval quality do you keep per byte?
+
+A reproducible benchmark for embedding compression. Given an embedder and a corpus, it
+measures the retrieval quality (nDCG@10, recall@10) retained by each compression method
+against the bytes it stores per vector -- the recall-per-byte frontier.
+
+Add your own method in five lines::
+
+    from bitbudget import method
+    import numpy as np
+
+    @method("my-2bit", bits=2)
+    def my_2bit(demb, qemb):
+        # return (query x doc similarity scores, bytes per stored vector)
+        codes = my_quantise(demb)
+        return qemb @ my_dequantise(codes).T, demb.shape[1] * 2 / 8
+
+Then ``bitbudget run --embedder mxbai --corpus scifact`` scores it alongside the built-ins.
+"""
+from .methods import method, METHODS, list_methods
+from .embedders import embedder, EMBEDDERS, list_embedders
+from .eval import evaluate
+from .indexes import index, INDEXES, list_indexes, bench_indexes
+
+__version__ = "0.1.0"
+__all__ = ["method", "METHODS", "list_methods", "embedder", "EMBEDDERS",
+           "list_embedders", "evaluate", "index", "INDEXES", "list_indexes",
+           "bench_indexes", "__version__"]

bitbudget/_bittrie.c

@@ -0,0 +1,92 @@
+/* Fast batched bit-trie query: coarse-to-fine beam descent over sorted codes, then a
+ * full-precision re-ranking pass, looped over all queries in C. Plain C (no Python.h);
+ * loaded via ctypes from _bittrie_build.py and compiled on demand. Falls back to the
+ * numpy/Python path in bittrie.py if no compiler is available. */
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+#include <math.h>
+
+typedef struct { double cost; long lo, hi; unsigned long long pre; } Node;
+
+static int cmp_node(const void *a, const void *b) {
+    double d = ((const Node *)a)->cost - ((const Node *)b)->cost;
+    return (d > 0) - (d < 0);
+}
+
+static long lower_bound_u64(const uint64_t *a, long lo, long hi, uint64_t key) {
+    while (lo < hi) { long mid = lo + (hi - lo) / 2; if (a[mid] < key) lo = mid + 1; else hi = mid; }
+    return lo;
+}
+
+/* out: nq*k int64 of original doc ids (top-k by inner product), -1 padding if fewer found. */
+/* Each query is independent and writes a disjoint slice of `out`, so the loop parallelises
+ * cleanly. Scratch buffers are allocated per thread inside the parallel region; when compiled
+ * without OpenMP the `omp` pragmas are ignored and this runs as one block, single-threaded. */
+void bt_query_batch(const uint64_t *codes, const int64_t *docids,
+                    const float *Xf, const float *qproj, const float *qf,
+                    int n, int d, int b, int nq, int depth, int beam, int k,
+                    int64_t *out) {
+    if (depth > b) depth = b;
+    if (beam < 1) beam = 1;
+    int cap = 2 * beam + 4;
+
+    #pragma omp parallel
+    {
+        Node *cur = (Node *)malloc(sizeof(Node) * cap);
+        Node *nxt = (Node *)malloc(sizeof(Node) * 2 * cap);
+        double *bscore = (double *)malloc(sizeof(double) * k);
+        long *bid = (long *)malloc(sizeof(long) * k);
+
+        #pragma omp for schedule(static)
+        for (int qi = 0; qi < nq; qi++) {
+            const float *qp = qproj + (size_t)qi * b;
+            const float *qv = qf + (size_t)qi * d;
+            int ncur = 1;
+            cur[0].cost = 0.0; cur[0].lo = 0; cur[0].hi = n; cur[0].pre = 0ULL;
+
+            for (int t = 0; t < depth; t++) {
+                int shift = b - t - 1;
+                int qb = qp[t] > 0.0f ? 1 : 0;
+                double c = fabs((double)qp[t]);
+                int nn = 0;
+                for (int i = 0; i < ncur; i++) {
+                    long lo = cur[i].lo, hi = cur[i].hi;
+                    unsigned long long pre = cur[i].pre;
+                    unsigned long long hi_part = (b - t) >= 64 ? 0ULL : (pre << (b - t));
+                    unsigned long long thresh = hi_part | (1ULL << shift);
+                    long m = lower_bound_u64(codes, lo, hi, (uint64_t)thresh);
+                    if (m > lo) { nxt[nn].cost = cur[i].cost + (qb == 0 ? 0.0 : c);
+                                  nxt[nn].lo = lo; nxt[nn].hi = m; nxt[nn].pre = (pre << 1); nn++; }
+                    if (hi > m) { nxt[nn].cost = cur[i].cost + (qb == 1 ? 0.0 : c);
+                                  nxt[nn].lo = m; nxt[nn].hi = hi; nxt[nn].pre = (pre << 1) | 1ULL; nn++; }
+                }
+                qsort(nxt, nn, sizeof(Node), cmp_node);
+                int keep = nn < beam ? nn : beam;
+                memcpy(cur, nxt, sizeof(Node) * keep);
+                ncur = keep;
+            }
+
+            for (int j = 0; j < k; j++) { bscore[j] = -1e300; bid[j] = -1; }
+            for (int i = 0; i < ncur; i++) {
+                for (long j = cur[i].lo; j < cur[i].hi; j++) {
+                    long id = (long)docids[j];
+                    const float *xr = Xf + (size_t)id * d;
+                    double s = 0.0;
+                    for (int dd = 0; dd < d; dd++) s += (double)qv[dd] * (double)xr[dd];
+                    int mn = 0;
+                    for (int t2 = 1; t2 < k; t2++) if (bscore[t2] < bscore[mn]) mn = t2;
+                    if (s > bscore[mn]) { bscore[mn] = s; bid[mn] = id; }
+                }
+            }
+            for (int a = 0; a < k; a++) {        /* selection-sort the k winners, score descending */
+                int mx = a;
+                for (int b2 = a + 1; b2 < k; b2++) if (bscore[b2] > bscore[mx]) mx = b2;
+                double ts = bscore[a]; bscore[a] = bscore[mx]; bscore[mx] = ts;
+                long ti = bid[a]; bid[a] = bid[mx]; bid[mx] = ti;
+            }
+            for (int j = 0; j < k; j++) out[(size_t)qi * k + j] = bid[j];
+        }
+        free(cur); free(nxt); free(bscore); free(bid);
+    }
+}

bitbudget/_bittrie_build.py

@@ -0,0 +1,94 @@
+"""Compile the bit-trie C kernel on demand and load it via ctypes.
+
+The package ships pure-Python on PyPI (the .c is data, not a built extension), so install never
+needs a compiler. The first time the fast path is requested we compile _bittrie.c into a cached
+shared library with the system C compiler; if no compiler is available we return None and the
+caller falls back to the numpy/Python path. The compiled library is cached by source hash, so the
+build happens once per machine.
+"""
+import ctypes
+import hashlib
+import os
+import subprocess
+
+_LIB = None
+_TRIED = False
+
+
+def _cache_dir():
+    d = os.environ.get("BITBUDGET_CACHE") or os.path.join(os.path.expanduser("~"), ".cache", "bitbudget")
+    os.makedirs(d, exist_ok=True)
+    return d
+
+
+def _brew_libomp():
+    """Return the Homebrew libomp prefix on macOS, or None."""
+    try:
+        p = subprocess.run(["brew", "--prefix", "libomp"], capture_output=True, text=True)
+        prefix = p.stdout.strip()
+        if p.returncode == 0 and prefix and os.path.isdir(prefix):
+            return prefix
+    except Exception:
+        pass
+    return None
+
+
+def _variants():
+    """(cflags, lflags) compile attempts, OpenMP first, single-thread C last."""
+    v = []
+    omp = _brew_libomp()
+    if omp:                                                    # macOS clang + Homebrew libomp (try first)
+        v.append((["-Xpreprocessor", "-fopenmp", f"-I{omp}/include"],
+                  [f"-L{omp}/lib", "-lomp", f"-Wl,-rpath,{omp}/lib"]))
+    v.append((["-fopenmp"], ["-fopenmp"]))                     # gcc / OpenMP-capable clang (Linux)
+    v.append(([], []))                                         # no OpenMP: single-threaded C
+    return v
+
+
+def _set_sig(lib):
+    f = lib.bt_query_batch
+    f.restype = None
+    P = ctypes.POINTER
+    f.argtypes = [P(ctypes.c_uint64), P(ctypes.c_int64), P(ctypes.c_float),
+                  P(ctypes.c_float), P(ctypes.c_float),
+                  ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int,
+                  ctypes.c_int, ctypes.c_int, ctypes.c_int, P(ctypes.c_int64)]
+    return lib
+
+
+def get_lib():
+    """Return the loaded ctypes library exposing bt_query_batch, or None if unavailable.
+
+    Tries to build a multithreaded (OpenMP) kernel first, falling back to single-threaded C,
+    then to None (the caller uses the numpy/Python path). The compiled library is cached by
+    source hash; set BITBUDGET_NO_OMP=1 to force the single-threaded build.
+    """
+    global _LIB, _TRIED
+    if _TRIED:
+        return _LIB
+    _TRIED = True
+    src = os.path.join(os.path.dirname(__file__), "_bittrie.c")
+    if not os.path.exists(src):
+        return None
+    h = hashlib.sha1(open(src, "rb").read()).hexdigest()[:12]
+    so = os.path.join(_cache_dir(), f"_bittrie_{h}.so")
+    if os.path.exists(so):                                     # cached: load directly
+        try:
+            _LIB = _set_sig(ctypes.CDLL(so)); return _LIB
+        except Exception:
+            try: os.remove(so)
+            except Exception: pass
+    cc = os.environ.get("CC") or "cc"
+    variants = [([], [])] if os.environ.get("BITBUDGET_NO_OMP") else _variants()
+    for cflags, lflags in variants:
+        cmd = [cc, "-O3", "-shared", "-fPIC", *cflags, "-o", so, src, *lflags, "-lm"]
+        try:
+            subprocess.run(cmd, check=True, capture_output=True)
+        except Exception:
+            continue
+        try:                                                  # must actually load (libomp present at runtime)
+            _LIB = _set_sig(ctypes.CDLL(so)); return _LIB
+        except Exception:
+            try: os.remove(so)
+            except Exception: pass
+    return None

bitbudget/bittrie.py

@@ -0,0 +1,108 @@
+"""Bit-trie index: a van Emde Boas / PATRICIA-style radix trie over compact codes.
+
+The sorted array of packed sign-codes *is* the trie: a node is a contiguous range of codes
+sharing a bit prefix. Search is coarse-to-fine beam descent (the most discriminative bit is the
+shallowest), then a re-ranking pass. Unsupervised (random rotation + sign), numpy-only, one knob
+(``beam``). It stores compact codes rather than full vectors, so it sits on the organisation axis
+at a fraction of a graph's footprint. Research entry to the BitBudget index benchmark; see the
+projection-quantisation-organisation survey for the motivation.
+
+This reference implementation packs into a single uint64, so ``n_bits <= 64``.
+"""
+import numpy as np
+
+
+class BitTrieIndex:
+    def __init__(self, n_bits=64, seed=0):
+        assert n_bits <= 64, "this reference packs into one uint64; use a word-array for >64 bits"
+        self.b = n_bits
+        self.seed = seed
+
+    def fit(self, X):
+        X = np.ascontiguousarray(X, dtype=np.float32)
+        n, d = X.shape
+        rng = np.random.RandomState(self.seed)
+        R = rng.randn(d, self.b).astype(np.float32)               # unsupervised random rotation (RaBitQ)
+        R, _ = np.linalg.qr(R) if self.b <= d else (R, None)
+        P = X @ R
+        order = np.argsort(-P.var(axis=0))                        # coarse-to-fine: MSB = highest variance
+        self.R = R[:, order]
+        P = P[:, order]
+        bits = (P > 0).astype(np.uint64)
+        weights = (np.uint64(1) << np.arange(self.b - 1, -1, -1, dtype=np.uint64))
+        codes = bits @ weights
+        order_idx = np.argsort(codes, kind="stable")
+        self.codes = codes[order_idx]                             # sorted: the implicit trie
+        self.docids = order_idx.astype(np.int64)
+        self.X = X                                                # kept only for optional float re-rank (cold)
+        self.n, self.d = n, d
+        return self
+
+    def _descend(self, qbits, qconf, beam, depth):
+        """Confidence-ordered beam over the sorted codes. Stop at ``depth`` bits so each surviving
+        prefix is still a fat bucket (the coarse-to-fine / anytime property)."""
+        b = self.b
+        state = [(0.0, 0, self.n, 0, 0)]                          # (cost, lo, hi, prefix_int, depth)
+        for t in range(depth):
+            shift = b - t - 1
+            nxt = []
+            for cost, lo, hi, pre, _ in state:
+                thresh = (pre << (b - t)) | (1 << shift)
+                m = lo + int(np.searchsorted(self.codes[lo:hi], np.uint64(thresh), "left"))
+                qb = int(qbits[t]); c = float(qconf[t])
+                left = (lo, m, (pre << 1))
+                right = (m, hi, (pre << 1) | 1)
+                for child_bit, (clo, chi, cpre) in ((0, left), (1, right)):
+                    if chi <= clo:
+                        continue
+                    add = 0.0 if child_bit == qb else c           # sibling costs its margin
+                    nxt.append((cost + add, clo, chi, cpre, t + 1))
+            nxt.sort(key=lambda e: e[0])
+            state = nxt[:beam]
+        return state
+
+    def query(self, q, topk=10, beam=64, depth=28, rerank="float"):
+        q = np.asarray(q, dtype=np.float32)
+        p = q @ self.R
+        qbits = (p > 0).astype(int)
+        qconf = np.abs(p)
+        ranges = self._descend(qbits, qconf, beam, depth)
+        cand = np.concatenate([self.docids[lo:hi] for _, lo, hi, _, _ in ranges]) if ranges else np.empty(0, np.int64)
+        if cand.size == 0:
+            return cand
+        if rerank == "float":                                     # DiskANN-style: cold float for re-rank only
+            s = self.X[cand] @ q
+        else:                                                     # asymmetric: float query vs sign code, no float
+            shifts = np.arange(self.b - 1, -1, -1, dtype=np.uint64)
+            sign = np.where(((self.codes_full(cand)[:, None] >> shifts) & np.uint64(1)) > 0, 1.0, -1.0)
+            s = sign @ (q @ self.R)
+        return cand[np.argsort(-s)[:topk]]
+
+    def query_batch(self, Q, topk=10, beam=64, depth=28):
+        """Top-k for every query in Q. Uses the compiled C kernel when available (one to two
+        orders of magnitude faster), else falls back to the pure-Python per-query path."""
+        Q = np.ascontiguousarray(Q, dtype=np.float32)
+        depth = min(depth, self.b)
+        from ._bittrie_build import get_lib
+        lib = get_lib()
+        if lib is None:
+            return np.array([self.query(q, topk, beam, depth, rerank="float") for q in Q])
+        import ctypes
+        qproj = np.ascontiguousarray(Q @ self.R, dtype=np.float32)
+        Xf = np.ascontiguousarray(self.X, dtype=np.float32)
+        out = np.full((len(Q), topk), -1, dtype=np.int64)
+        cast = lambda a, t: a.ctypes.data_as(ctypes.POINTER(t))
+        lib.bt_query_batch(
+            cast(self.codes, ctypes.c_uint64), cast(self.docids, ctypes.c_int64),
+            cast(Xf, ctypes.c_float), cast(qproj, ctypes.c_float), cast(Q, ctypes.c_float),
+            self.n, self.d, self.b, len(Q), int(depth), int(beam), int(topk),
+            cast(out, ctypes.c_int64))
+        return out
+
+    def codes_full(self, cand):
+        inv = np.empty(self.n, np.int64); inv[self.docids] = np.arange(self.n)
+        return self.codes[inv[cand]]
+
+    def index_bytes(self):
+        """Bytes that must be in RAM to route (codes + leaf ids); the trie is implicit in the sort."""
+        return self.codes.nbytes + self.docids.nbytes

bitbudget/cli.py

@@ -0,0 +1,231 @@
+"""BitBudget command line.
+
+    bitbudget methods                              # list registered compression methods
+    bitbudget embedders                            # list registered embedders
+    bitbudget embed   --embedder mxbai --corpus scifact      # torch step -> cached embeddings
+    bitbudget eval    --embedder mxbai --corpus scifact      # numpy step -> results card
+    bitbudget run     --embedder mxbai --corpus scifact nfcorpus arguana fiqa   # embed + eval
+    bitbudget leaderboard results/*.json           # aggregate cards into a markdown table
+    bitbudget indexes                              # list registered indexes (organisation axis)
+    bitbudget bench-index --synthetic 100000 128   # recall vs QPS vs bytes (flat/hnsw/ivfpq/bittrie)
+
+`run` embeds (torch) and evaluates (numpy) in one process, which is safe because the core
+methods use no faiss. The faiss-backed indexes in `bench-index` import faiss only when run, so
+run it in its own process (no torch) to avoid the OpenMP clash.
+"""
+import argparse
+import glob
+import json
+import os
+
+import numpy as np
+
+from . import datasets
+from .eval import evaluate, aggregate
+from .methods import list_methods, METHODS
+from .indexes import bench_indexes, list_indexes, INDEXES
+
+
+def _emb_path(cache, embedder, corpus):
+    return os.path.join(os.path.expanduser(cache), f"emb_{embedder}_{corpus}.npz")
+
+
+def cmd_embed(a):
+    from .embedders import EMBEDDERS                      # lazy: imports torch
+    emb = EMBEDDERS[a.embedder]
+    for corpus in a.corpus:
+        corp, queries, qrels = datasets.load(corpus, a.cache, a.split)
+        cids, qids = list(corp), list(queries)
+        print(f"[embed] {a.embedder} x {corpus}: {len(cids)} docs, {len(qids)} queries")
+        demb = emb([corp[c] for c in cids], is_query=False, device=a.device, batch_size=a.batch_size)
+        qemb = emb([queries[q] for q in qids], is_query=True, device=a.device, batch_size=a.batch_size)
+        os.makedirs(os.path.expanduser(a.cache), exist_ok=True)
+        cidset = set(cids)
+        np.savez(_emb_path(a.cache, a.embedder, corpus), demb=demb, qemb=qemb,
+                 cids=np.array(cids), qids=np.array(qids),
+                 qrels=json.dumps({q: {d: s for d, s in qrels[q].items() if d in cidset} for q in qids}))
+        print(f"[embed] cached -> {_emb_path(a.cache, a.embedder, corpus)}  {demb.shape}")
+
+
+def _load_emb(a, corpus):
+    z = np.load(_emb_path(a.cache, a.embedder, corpus), allow_pickle=True)
+    return z["demb"], z["qemb"], z["cids"], z["qids"], json.loads(str(z["qrels"]))
+
+
+def cmd_eval(a, _embed_first=False):
+    methods = a.methods or list_methods()
+    per_corpus = {}
+    for corpus in a.corpus:
+        if _embed_first:
+            cmd_embed(argparse.Namespace(**{**vars(a), "corpus": [corpus]}))
+        demb, qemb, cids, qids, qrels = _load_emb(a, corpus)
+        rows = evaluate(demb, qemb, cids, qids, qrels, methods)
+        per_corpus[corpus] = rows
+        print(f"\n=== {a.embedder} x {corpus} ({demb.shape[1]}-d) ===")
+        print(f"  {'method':16s}{'bytes':>8s}{'nDCG@10':>9s}{'recall@10':>11s}{'% float':>9s}")
+        for r in rows:
+            print(f"  {r['method']:16s}{r['bytes']:8.0f}{r['ndcg']:9.3f}{r['recall']:11.3f}{r['pct_float']:8.0f}%")
+    os.makedirs(a.out, exist_ok=True)
+    card = dict(embedder=a.embedder, corpora=a.corpus,
+                dim=int(_load_emb(a, a.corpus[0])[0].shape[1]),
+                per_corpus=per_corpus, aggregate=aggregate(per_corpus))
+    path = os.path.join(a.out, f"card_{a.embedder}.json")
+    json.dump(card, open(path, "w"), indent=2)
+    print(f"\n=== aggregate over {len(a.corpus)} corpora (mean +/- std) ===")
+    for r in card["aggregate"]:
+        print(f"  {r['method']:16s}{r['bytes']:8.0f}B  nDCG {r['ndcg']:.3f}  "
+              f"{r['pct_float']:.0f}+/-{r['pct_float_std']:.0f}% of float")
+    print(f"\nwrote results card -> {path}")
+    return card
+
+
+def cmd_methods(a):
+    print("registered methods (name | axis | bits):")
+    for n in list_methods():
+        fn = METHODS[n]
+        print(f"  {n:16s} {getattr(fn,'axis','?'):12s} {fn.bits}")
+
+
+def cmd_indexes(a):
+    print("registered indexes (name | axis):")
+    for n in list_indexes():
+        print(f"  {n:14s} {getattr(INDEXES[n], 'axis', 'organisation')}")
+
+
+def cmd_bench_index(a):
+    """Build each index over document vectors and report recall@k, QPS and bytes/vec.
+
+    faiss (its OpenMP) and the bit-trie (libomp for the C kernel) cannot share one process on
+    macOS (two OpenMP runtimes). When both groups are requested we run each in its own subprocess
+    and merge; the data is deterministic from the arguments, so the split is transparent.
+    """
+    sel = a.indexes or list_indexes()
+    faiss_grp = [n for n in sel if getattr(INDEXES[n], "axis", "") == "organisation"]
+    other_grp = [n for n in sel if n not in faiss_grp]
+    if faiss_grp and other_grp and not getattr(a, "no_split", False):
+        import subprocess
+        import sys
+        import tempfile
+        import shutil
+        merged, meta = [], {}
+        for grp in (faiss_grp, other_grp):
+            td = tempfile.mkdtemp()
+            cmd = [sys.executable, "-m", "bitbudget.cli", "bench-index",
+                   *_bench_data_args(a), "--k", str(a.k), "--out", td, "--no-split",
+                   "--indexes", *grp]
+            subprocess.run(cmd, check=True, env=os.environ.copy())
+            c = json.load(open(os.path.join(td, "index_card.json")))
+            merged += c["rows"]; meta = c
+            shutil.rmtree(td, ignore_errors=True)
+        merged.sort(key=lambda r: r["bytes"])
+        os.makedirs(a.out, exist_ok=True)
+        card = dict(source=meta["source"], k=a.k, dim=meta["dim"], n=meta["n"], rows=merged)
+        json.dump(card, open(os.path.join(a.out, "index_card.json"), "w"), indent=2)
+        print("\n=== board (faiss + bit-trie run in separate processes) ===")
+        for r in merged:
+            print(f"  {r['method']:14s}{r['bytes']:8.0f}B  recall@{a.k}={r['recall']:.3f}  {r['qps']:9.0f} qps")
+        print(f"wrote index card -> {os.path.join(a.out, 'index_card.json')}")
+        return card
+
+    xb, xq, label = _bench_load(a)
+    print(f"[bench-index] {label}: base={tuple(xb.shape)}, queries={len(xq)}, k={a.k}")
+    rows = bench_indexes(xb, xq, k=a.k, which=sel)
+    os.makedirs(a.out, exist_ok=True)
+    card = dict(source=label, k=a.k, dim=int(xb.shape[1]), n=int(len(xb)), rows=rows)
+    path = os.path.join(a.out, "index_card.json")
+    json.dump(card, open(path, "w"), indent=2)
+    print(f"\nwrote index card -> {path}")
+    return card
+
+
+def _bench_data_args(a):
+    if a.synthetic:
+        return ["--synthetic", str(a.synthetic[0]), str(a.synthetic[1])]
+    if a.npz:
+        return ["--npz", a.npz]
+    return ["--embedder", a.embedder, "--corpus", a.corpus, "--cache", a.cache]
+
+
+def _bench_load(a):
+    if a.synthetic:
+        N, D = a.synthetic
+        rng = np.random.RandomState(0)
+        centers = rng.randn(max(2, N // 500), D).astype(np.float32) * 4.0
+        xb = centers[rng.randint(0, len(centers), N)] + rng.randn(N, D).astype(np.float32)
+        nq = min(10000, max(1, N // 10))                  # enough queries to time throughput stably
+        xq = xb[rng.choice(N, nq, replace=False)] + 0.1 * rng.randn(nq, D).astype(np.float32)
+        return xb, xq, f"synthetic N={N} D={D}"
+    if a.npz:
+        z = np.load(os.path.expanduser(a.npz), allow_pickle=True)
+        xb = z["base"] if "base" in z.files else z["demb"]
+        xq = z["query"] if "query" in z.files else z["qemb"]
+        return xb, xq, a.npz
+    if a.embedder and a.corpus:
+        z = np.load(_emb_path(a.cache, a.embedder, a.corpus), allow_pickle=True)
+        return z["demb"], z["qemb"], f"{a.embedder} x {a.corpus}"
+    raise SystemExit("bench-index needs one of: --synthetic N D | --npz PATH | --embedder X --corpus Y")
+
+
+def cmd_embedders(a):
+    from .embedders import list_embedders, EMBEDDERS
+    print("registered embedders (name | dim | matryoshka):")
+    for n in list_embedders():
+        fn = EMBEDDERS[n]
+        print(f"  {n:10s} dim={fn.dim}  matryoshka={fn.matryoshka}  ({fn.model})")
+
+
+def cmd_leaderboard(a):
+    cards = [json.load(open(p)) for pat in a.cards for p in glob.glob(pat)]
+    print("# BitBudget leaderboard\n")
+    for c in sorted(cards, key=lambda c: -c["dim"]):
+        print(f"## {c['embedder']} ({c['dim']}-d), mean over {len(c['corpora'])} corpora\n")
+        print("| Method | Axis | Bytes/vec | nDCG@10 | % of float |")
+        print("|---|---|---|---|---|")
+        for r in c["aggregate"]:
+            print(f"| {r['method']} | {r['axis']} | {r['bytes']:.0f} | "
+                  f"{r['ndcg']:.3f} | {r['pct_float']:.0f} ± {r['pct_float_std']:.0f} |")
+        print()
+
+
+def main(argv=None):
+    p = argparse.ArgumentParser(prog="bitbudget", description="recall-per-byte benchmark for embedding compression")
+    sub = p.add_subparsers(dest="cmd", required=True)
+    for name in ("embed", "eval", "run"):
+        s = sub.add_parser(name)
+        s.add_argument("--embedder", required=True)
+        s.add_argument("--corpus", nargs="+", default=datasets.CORPORA)
+        s.add_argument("--methods", nargs="+", default=None)
+        s.add_argument("--cache", default="~/.cache/bitbudget")
+        s.add_argument("--out", default="results")
+        s.add_argument("--device", default="cpu", help="cpu | mps (embedding step)")
+        s.add_argument("--batch-size", type=int, default=64)
+        s.add_argument("--split", default="test")
+    sub.add_parser("methods")
+    sub.add_parser("embedders")
+    sub.add_parser("indexes")
+    lb = sub.add_parser("leaderboard"); lb.add_argument("cards", nargs="+")
+    bi = sub.add_parser("bench-index")
+    bi.add_argument("--synthetic", nargs=2, type=int, metavar=("N", "D"),
+                    help="benchmark on N clustered Gaussian vectors of dimension D")
+    bi.add_argument("--npz", help="load base/query (keys base,query or demb,qemb) from an .npz")
+    bi.add_argument("--embedder", help="use a cached embedding (with --corpus): demb as base")
+    bi.add_argument("--corpus")
+    bi.add_argument("--cache", default="~/.cache/bitbudget")
+    bi.add_argument("--k", type=int, default=10)
+    bi.add_argument("--indexes", nargs="+", default=None, help="subset of indexes to run")
+    bi.add_argument("--out", default="results")
+    bi.add_argument("--no-split", dest="no_split", action="store_true",
+                    help="run all indexes in one process (do not isolate faiss from the bit-trie)")
+    a = p.parse_args(argv)
+    if a.cmd == "embed":   cmd_embed(a)
+    elif a.cmd == "eval":  cmd_eval(a)
+    elif a.cmd == "run":   cmd_eval(a, _embed_first=True)
+    elif a.cmd == "methods":     cmd_methods(a)
+    elif a.cmd == "embedders":   cmd_embedders(a)
+    elif a.cmd == "indexes":     cmd_indexes(a)
+    elif a.cmd == "leaderboard": cmd_leaderboard(a)
+    elif a.cmd == "bench-index": cmd_bench_index(a)
+
+
+if __name__ == "__main__":
+    main()

bitbudget/datasets.py

@@ -0,0 +1,49 @@
+"""BEIR corpus loading. Datasets auto-download to the cache dir on first use."""
+import json
+import os
+import shutil
+import urllib.request
+import zipfile
+
+BEIR_URL = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{}.zip"
+# the standard small/medium BEIR corpora the leaderboard is defined over
+CORPORA = ["scifact", "nfcorpus", "arguana", "fiqa"]
+
+
+def _download(name, cache):
+    d = os.path.join(cache, name)
+    if not os.path.isdir(d):
+        os.makedirs(cache, exist_ok=True)
+        z = os.path.join(cache, name + ".zip")
+        if not os.path.exists(z):
+            req = urllib.request.Request(BEIR_URL.format(name), headers={"User-Agent": "Mozilla/5.0"})
+            with urllib.request.urlopen(req) as r, open(z + ".part", "wb") as f:
+                shutil.copyfileobj(r, f)
+            os.replace(z + ".part", z)
+        with zipfile.ZipFile(z) as zf:
+            zf.extractall(cache)
+    return d
+
+
+def load(name, cache="~/.cache/bitbudget", split="test"):
+    cache = os.path.expanduser(cache)
+    d = _download(name, cache)
+    corpus = {}
+    with open(os.path.join(d, "corpus.jsonl")) as f:
+        for line in f:
+            o = json.loads(line)
+            corpus[o["_id"]] = (o.get("title", "") + " " + o.get("text", "")).strip()
+    queries = {}
+    with open(os.path.join(d, "queries.jsonl")) as f:
+        for line in f:
+            o = json.loads(line)
+            queries[o["_id"]] = o["text"]
+    qrels = {}
+    with open(os.path.join(d, "qrels", split + ".tsv")) as f:
+        next(f)
+        for line in f:
+            qid, did, score = line.strip().split("\t")
+            if int(score) > 0:
+                qrels.setdefault(qid, {})[did] = int(score)
+    queries = {q: queries[q] for q in qrels if q in queries}
+    return corpus, queries, qrels

bitbudget/embedders.py

@@ -0,0 +1,46 @@
+"""Embedders. A named embedder turns a list of texts into a float32 matrix. These import
+sentence-transformers (torch) and so must run in a process that does NOT also import faiss
+(OpenMP clash on macOS); the CLI keeps embedding and faiss-backed evaluation in separate
+processes. Register your own with ``@embedder``.
+"""
+EMBEDDERS = {}
+
+
+def embedder(name, model=None, dim=None, query_prompt="", matryoshka=False):
+    """Register an embedder. `query_prompt` is prepended to queries only (asymmetric models)."""
+    def deco(fn):
+        fn.bitbudget_name = name
+        fn.model, fn.dim, fn.query_prompt, fn.matryoshka = model, dim, query_prompt, matryoshka
+        EMBEDDERS[name] = fn
+        return fn
+    return deco
+
+
+def list_embedders():
+    return sorted(EMBEDDERS)
+
+
+def _st(model_name, device, max_seq_length):
+    from sentence_transformers import SentenceTransformer
+    m = SentenceTransformer(model_name, device=device, trust_remote_code=True)
+    if max_seq_length:
+        m.max_seq_length = max_seq_length
+    return m
+
+
+def _make_st_embedder(name, model_name, dim, query_prompt, matryoshka, max_seq_length=256):
+    @embedder(name, model=model_name, dim=dim, query_prompt=query_prompt, matryoshka=matryoshka)
+    def fn(texts, is_query=False, device="cpu", batch_size=64):
+        import numpy as np
+        m = _st(model_name, device, max_seq_length)
+        if is_query and query_prompt:
+            texts = [query_prompt + t for t in texts]
+        return m.encode(texts, batch_size=batch_size, normalize_embeddings=True,
+                        show_progress_bar=True).astype(np.float32)
+    return fn
+
+
+# built-in embedders used in the paper
+_make_st_embedder("minilm", "sentence-transformers/all-MiniLM-L6-v2", 384, "", False)
+_make_st_embedder("mxbai", "mixedbread-ai/mxbai-embed-large-v1", 1024,
+                  "Represent this sentence for searching relevant passages: ", True)

bitbudget/eval.py

@@ -0,0 +1,44 @@
+"""The BitBudget protocol: given document/query embeddings and judgements, score every
+registered compression method on retrieval quality against the bytes it stores per vector."""
+import numpy as np
+
+from .methods import METHODS
+from .metrics import ndcg_at_k, recall_at_k
+
+
+def evaluate(demb, qemb, cids, qids, qrels, methods=None, recall_k=10, ndcg_k=10):
+    """Run `methods` (default: all registered) and return a list of result dicts:
+    {method, axis, bytes, ndcg, recall, pct_float}. nDCG uses the graded qrels; recall uses
+    the exact float top-k as ground truth; pct_float is nDCG relative to the float32 baseline.
+    """
+    methods = methods or list(METHODS)
+    exact = qemb @ demb.T                              # exact float scores = recall ground truth
+    base = ndcg_at_k(exact, qids, cids, qrels, ndcg_k)
+    rows = []
+    for name in methods:
+        fn = METHODS[name]
+        scores, bpv = fn(demb, qemb)
+        n = ndcg_at_k(scores, qids, cids, qrels, ndcg_k)
+        r = recall_at_k(scores, exact, recall_k)
+        rows.append(dict(method=name, axis=getattr(fn, "axis", "?"), bytes=float(bpv),
+                         ndcg=round(n, 4), recall=round(r, 4),
+                         pct_float=round(100 * n / base, 1) if base else float("nan")))
+    rows.sort(key=lambda x: x["bytes"])
+    return rows
+
+
+def aggregate(per_corpus):
+    """Mean +/- std across corpora for each method, given a {corpus: rows} mapping."""
+    by = {}
+    for corpus, rows in per_corpus.items():
+        for r in rows:
+            by.setdefault(r["method"], []).append(r)
+    out = []
+    for name, rs in by.items():
+        nd = np.array([r["ndcg"] for r in rs]);  pc = np.array([r["pct_float"] for r in rs])
+        out.append(dict(method=name, axis=rs[0]["axis"], bytes=rs[0]["bytes"],
+                        ndcg=round(float(nd.mean()), 3), ndcg_std=round(float(nd.std()), 3),
+                        pct_float=round(float(pc.mean()), 1), pct_float_std=round(float(pc.std()), 1),
+                        n_corpora=len(rs)))
+    out.sort(key=lambda x: x["bytes"])
+    return out

bitbudget/indexes.py

@@ -0,0 +1,136 @@
+"""The organisation axis: build an index over document vectors and measure recall@k, query
+throughput (QPS) and the bytes it stores per vector.
+
+This is the axis the compression leaderboard does not cover. A graph index such as HNSW does not
+shrink the footprint, it *adds* bytes (the vectors plus a graph) and buys throughput; a compact
+code such as the bit-trie shrinks it. Reporting recall, QPS and bytes together makes the
+trade-off explicit, in the spirit of ann-benchmarks.
+
+Register an index with ``@index``. The faiss-backed indexes (flat/hnsw/ivfpq) import faiss lazily
+and are skipped with a message if faiss is not installed; the bit-trie is numpy-only. Vectors are
+L2-normalised on ingest so inner-product and L2 rankings agree. Run in its own process (no torch)
+to avoid the faiss/OpenMP clash (see README).
+"""
+import time
+import numpy as np
+
+INDEXES = {}
+
+
+def index(name, axis="organisation"):
+    """Register an index. An index is ``fn(xb, xq, k) -> (topk_indices, bytes_per_vec, qps)``."""
+    def deco(fn):
+        fn.bitbudget_name = name
+        fn.axis = axis
+        INDEXES[name] = fn
+        return fn
+    return deco
+
+
+def list_indexes():
+    return sorted(INDEXES)
+
+
+def _time(search, repeats=3):
+    best, I = float("inf"), None
+    for _ in range(repeats):
+        t = time.time(); I = search(); best = min(best, time.time() - t)
+    return I, best
+
+
+# ----------------------------------------------------------------- faiss-backed (organisation)
+@index("flat")
+def flat(xb, xq, k):
+    import faiss
+    D = xb.shape[1]
+    ix = faiss.IndexFlatIP(D); ix.add(xb)
+    I, t = _time(lambda: ix.search(xq, k)[1])
+    return I, 4 * D, len(xq) / t                       # full float vectors, exact
+
+
+@index("hnsw")
+def hnsw(xb, xq, k, M=24, efc=200, efs=128):
+    import faiss
+    D = xb.shape[1]
+    ix = faiss.IndexHNSWFlat(D, M, faiss.METRIC_INNER_PRODUCT)
+    ix.hnsw.efConstruction = efc; ix.add(xb); ix.hnsw.efSearch = efs
+    I, t = _time(lambda: ix.search(xq, k)[1])
+    return I, 4 * D + M * 2 * 4, len(xq) / t           # vectors + graph edges (both levels approx)
+
+
+@index("ivfpq")
+def ivfpq(xb, xq, k, nprobe=64):
+    import faiss
+    n, D = xb.shape
+    m = D // 2
+    while D % m:
+        m -= 1
+    nlist = max(1, min(8192, n // 100))
+    ix = faiss.IndexIVFPQ(faiss.IndexFlatIP(D), D, nlist, m, 8, faiss.METRIC_INNER_PRODUCT)
+    samp = xb[np.random.RandomState(0).choice(n, min(200000, n), replace=False)]
+    ix.train(samp); ix.add(xb); ix.nprobe = min(nprobe, nlist)
+    I, t = _time(lambda: ix.search(xq, k)[1])
+    return I, float(m), len(xq) / t                    # PQ code bytes
+
+
+# ----------------------------------------------------------------- numpy bit-trie (research entry)
+@index("bittrie", axis="quantisation")
+def bittrie(xb, xq, k, nbits=64, beam=128, depth=24):
+    from .bittrie import BitTrieIndex
+    ix = BitTrieIndex(n_bits=nbits, seed=0).fit(xb)
+    ix.query_batch(xq[:2], topk=k, beam=beam, depth=depth)   # warm up (compile C kernel once)
+    I, t = _time(lambda: ix.query_batch(xq, topk=k, beam=beam, depth=depth), repeats=2)
+    return I, nbits / 8.0, len(xq) / t                 # compact routing code (float kept cold for re-rank)
+
+
+# ----------------------------------------------------------------- harness
+def _exact(xb, xq, k, chunk=200000):
+    """Exact top-k by inner product, chunked over documents to bound memory. The chunk shrinks
+    as the query count grows so the (nq x chunk) score block stays small regardless of nq."""
+    nq = len(xq)
+    chunk = max(1024, min(chunk, 20_000_000 // max(1, nq)))   # keep nq*chunk ~ 20M elements
+    best_s = np.full((nq, k), -1e30, np.float32)
+    best_i = np.zeros((nq, k), np.int64)
+    for s in range(0, len(xb), chunk):
+        e = min(s + chunk, len(xb))
+        sc = xq @ xb[s:e].T
+        cs = np.concatenate([best_s, sc], 1)
+        ci = np.concatenate([best_i, np.broadcast_to(np.arange(s, e), (nq, e - s))], 1)
+        part = np.argpartition(-cs, k - 1, axis=1)[:, :k]
+        best_s = np.take_along_axis(cs, part, 1)
+        best_i = np.take_along_axis(ci, part, 1)
+    order = np.argsort(-best_s, 1)
+    return np.take_along_axis(best_i, order, 1)
+
+
+def bench_indexes(xb, xq, k=10, gt=None, which=None, verbose=True):
+    """Build each index and report recall@k, QPS and bytes/vec. Returns rows sorted by footprint.
+    Vectors are L2-normalised so the inner-product ground truth is metric-consistent across indexes.
+    """
+    xb = np.ascontiguousarray(xb, dtype=np.float32)
+    xq = np.ascontiguousarray(xq, dtype=np.float32)
+    xb /= (np.linalg.norm(xb, axis=1, keepdims=True) + 1e-9)
+    xq /= (np.linalg.norm(xq, axis=1, keepdims=True) + 1e-9)
+    if gt is None:
+        gt = _exact(xb, xq, k)
+    rows = []
+    # run faiss-backed indexes before the numpy bit-trie: the bit-trie loads its own OpenMP
+    # runtime (libomp) for the C kernel, which clashes with faiss's OpenMP if faiss runs after it
+    # in the same process (the macOS two-runtimes problem). Ordering it last avoids the clash.
+    names = sorted(which or list_indexes(), key=lambda n: (n == "bittrie", n))
+    for name in names:
+        fn = INDEXES[name]
+        try:
+            I, bpv, qps = fn(xb, xq, k)
+        except ImportError:
+            if verbose:
+                print(f"  [skip {name}: install bitbudget[faiss]]")
+            continue
+        rec = float(np.mean([len(set(I[i][:k].tolist()) & set(gt[i][:k].tolist())) / k
+                             for i in range(len(xq))]))
+        rows.append(dict(method=name, axis=getattr(fn, "axis", "organisation"),
+                         bytes=float(bpv), recall=round(rec, 4), qps=round(float(qps), 1)))
+        if verbose:
+            print(f"  {name:14s}{bpv:8.0f}B  recall@{k}={rec:.3f}  {qps:9.0f} qps")
+    rows.sort(key=lambda r: r["bytes"])
+    return rows

bitbudget/methods.py

@@ -0,0 +1,114 @@
+"""Compression methods. A method maps float document/query embeddings to a
+query-by-document similarity matrix under some compression, and reports the bytes it stores
+per document vector. Register your own with the ``@method`` decorator.
+
+    @method("my-method", bits=2)
+    def my_method(demb, qemb):
+        ...
+        return scores, bytes_per_vec        # scores: (n_queries, n_docs); bytes: float
+
+Everything here is numpy-only so that evaluation runs in the same process as torch embedding
+without the faiss/OpenMP clash (see README). faiss-backed methods live in methods_faiss.py and
+run in a separate evaluation process.
+"""
+import numpy as np
+
+METHODS = {}
+
+
+def method(name, bits=None, axis="quantisation"):
+    """Register a compression method. `bits` and `axis` are metadata for the leaderboard."""
+    def deco(fn):
+        fn.bitbudget_name = name
+        fn.bits = bits
+        fn.axis = axis
+        METHODS[name] = fn
+        return fn
+    return deco
+
+
+def list_methods():
+    return sorted(METHODS)
+
+
+# ---------------------------------------------------------------- baselines / quantisation
+@method("float32", bits=32, axis="none")
+def float32(demb, qemb):
+    return qemb @ demb.T, 4 * demb.shape[1]
+
+
+@method("int8", bits=8)
+def int8(demb, qemb):
+    scale = np.abs(demb).max(0, keepdims=True) / 127.0 + 1e-9
+    dq = np.round(demb / scale) * scale
+    return qemb @ dq.T, demb.shape[1]
+
+
+@method("binary", bits=1)
+def binary(demb, qemb):
+    return np.sign(qemb) @ np.sign(demb).T, demb.shape[1] / 8.0
+
+
+@method("binary+rerank", bits=1)
+def binary_rerank(demb, qemb, depth=100):
+    bq, bd = np.sign(qemb), np.sign(demb)
+    cand = np.argsort(-(bq @ bd.T), axis=1)[:, :depth]
+    scores = np.full((qemb.shape[0], demb.shape[0]), -1e9, dtype=np.float32)
+    for i in range(qemb.shape[0]):
+        scores[i, cand[i]] = qemb[i] @ demb[cand[i]].T
+    return scores, demb.shape[1] / 8.0          # compact code stored; full vectors fetched on re-rank
+
+
+@method("rabitq", bits=1)
+def rabitq(demb, qemb, seed=0):
+    D = demb.shape[1]
+    R = np.linalg.qr(np.random.RandomState(seed).randn(D, D))[0].astype(np.float32)
+    return (qemb @ R) @ np.sign(demb @ R).T, D / 8.0
+
+
+# ---------------------------------------------------------------- product quantisation (numpy)
+def _kmeans(X, k, iters=10, seed=0):
+    rng = np.random.RandomState(seed)
+    C = X[rng.choice(len(X), size=min(k, len(X)), replace=False)].copy()
+    for _ in range(iters):
+        d = ((X ** 2).sum(1, keepdims=True) - 2 * X @ C.T + (C ** 2).sum(1))
+        a = d.argmin(1)
+        for j in range(len(C)):
+            m = a == j
+            if m.any():
+                C[j] = X[m].mean(0)
+    return C, a
+
+
+@method("pq", bits=8, axis="quantisation")
+def pq(demb, qemb, m=None, ksub=256):
+    """Product quantisation: m sub-quantisers, 256 centroids each (m bytes/vector). Asymmetric
+    scoring against reconstructed documents."""
+    D = demb.shape[1]
+    m = m or max(1, D // 8)
+    assert D % m == 0, f"PQ needs D ({D}) divisible by m ({m})"
+    sub = D // m
+    recon = np.empty_like(demb)
+    for s in range(m):
+        sl = slice(s * sub, (s + 1) * sub)
+        C, a = _kmeans(demb[:, sl], ksub, seed=s)
+        recon[:, sl] = C[a]
+    return qemb @ recon.T, float(m)
+
+
+# ---------------------------------------------------------------- projection axis (dimensions)
+@method("matryoshka", bits=32, axis="projection")
+def matryoshka(demb, qemb, dim=None):
+    """Prefix-truncation. Fair only for Matryoshka-trained embedders; on others it is naive."""
+    dim = dim or demb.shape[1] // 4
+    return qemb[:, :dim] @ demb[:, :dim].T, 4 * dim
+
+
+@method("pca", bits=32, axis="projection")
+def pca(demb, qemb, dim=None):
+    dim = dim or demb.shape[1] // 4
+    mean = demb.mean(0)
+    dc = demb - mean
+    _, V = np.linalg.eigh(dc.T @ dc)
+    W = V[:, ::-1][:, :dim]
+    return ((qemb - mean) @ W) @ (dc @ W).T, 4 * dim

bitbudget/metrics.py

@@ -0,0 +1,37 @@
+"""Retrieval metrics. nDCG@k uses graded BEIR judgements; recall@k uses the exact
+floating-point neighbours as ground truth (the ANN-Benchmarks convention)."""
+import numpy as np
+
+
+def _topk(scores_row, k):
+    k = min(k, scores_row.shape[0])
+    idx = np.argpartition(-scores_row, k - 1)[:k]
+    return idx[np.argsort(-scores_row[idx])]
+
+
+def ndcg_at_k(scores, qids, cids, qrels, k=10):
+    """Mean nDCG@k over queries with graded relevance judgements `qrels`."""
+    disc = 1.0 / np.log2(np.arange(2, k + 2))
+    out = []
+    cids = [str(c) for c in cids]
+    for i, q in enumerate(qids):
+        rel = qrels.get(str(q), {})
+        if not rel:
+            continue
+        top = _topk(scores[i], k)
+        gains = np.array([rel.get(cids[j], 0) for j in top], dtype=float)
+        dcg = np.sum((2 ** gains - 1) * disc[:len(gains)])
+        ideal = np.sort(np.array(list(rel.values()), dtype=float))[::-1][:k]
+        idcg = np.sum((2 ** ideal - 1) * disc[:len(ideal)])
+        out.append(dcg / idcg if idcg > 0 else 0.0)
+    return float(np.mean(out)) if out else float("nan")
+
+
+def recall_at_k(scores, exact_scores, k=10):
+    """k-recall@k against the exact floating-point top-k (ground truth from `exact_scores`)."""
+    out = []
+    for i in range(scores.shape[0]):
+        gt = set(_topk(exact_scores[i], k).tolist())
+        got = set(_topk(scores[i], k).tolist())
+        out.append(len(gt & got) / float(k))
+    return float(np.mean(out))

bitbudget-0.1.0.dist-info/METADATA

@@ -0,0 +1,148 @@
+Metadata-Version: 2.4
+Name: bitbudget
+Version: 0.1.0
+Summary: How much retrieval quality do you keep per byte? A reproducible benchmark for embedding compression.
+Author: Sean Moran
+License: MIT
+Project-URL: Paper, https://arxiv.org/abs/2510.04127
+Project-URL: Leaderboard, https://github.com/sjmoran/bitbudget/blob/main/LEADERBOARD.md
+Keywords: retrieval,embeddings,quantisation,hashing,compression,ANN,RAG
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.21
+Provides-Extra: embed
+Requires-Dist: sentence-transformers>=2.2; extra == "embed"
+Provides-Extra: faiss
+Requires-Dist: faiss-cpu>=1.7.4; extra == "faiss"
+Provides-Extra: all
+Requires-Dist: sentence-transformers>=2.2; extra == "all"
+Requires-Dist: faiss-cpu>=1.7.4; extra == "all"
+Dynamic: license-file
+
+# BitBudget
+
+**How much retrieval quality do you keep per byte?**
+
+BitBudget is a small, reproducible benchmark for **embedding compression**. Give it an
+embedder and a corpus and it reports the retrieval quality (nDCG@10, recall@10) that each
+compression method retains against the **bytes it stores per vector** — the recall‑per‑byte
+frontier that every RAG and vector‑database deployment actually lives on.
+
+It is the companion benchmark to the survey *“Projection and Quantisation: A Unifying View of
+Learning to Hash, from Random Projections to the RAG Era”* and exists to answer one question
+that today is mostly answered by vendor blog posts: **when you binarise / int8 / RaBitQ /
+product‑quantise / Matryoshka‑truncate your embeddings, what do you actually lose?**
+
+## The headline finding
+
+> **Bits beat dimensions.** Spending a fixed byte budget on *more coarsely quantised*
+> coordinates beats spending it on *fewer full‑precision* coordinates, at every budget and
+> for every embedder we have tried. One‑bit codes with a cheap re‑ranking pass are **32×
+> smaller than float at no measurable loss**.
+
+```
+mxbai‑embed‑large (1024‑d), mean over 4 BEIR corpora
+  binary+rerank      128 B   nDCG 0.509   100% of float   ← 32× smaller, lossless
+  pq                 128 B   nDCG 0.488    96%
+  rabitq             128 B   nDCG 0.487    96%
+  matryoshka        1024 B   nDCG 0.439    86%             ← 4× smaller, projection axis
+  float32           4096 B   nDCG 0.508   100%
+```
+
+See **[LEADERBOARD.md](LEADERBOARD.md)** for the full table.
+
+## Install
+
+```bash
+pip install bitbudget            # evaluation only (numpy)
+pip install "bitbudget[all]"     # + sentence-transformers (embedding) + faiss
+```
+
+## Quickstart
+
+```bash
+bitbudget methods                                   # list compression methods
+bitbudget run --embedder mxbai --corpus scifact     # embed + evaluate, print a results card
+bitbudget leaderboard results/card_*.json           # render a markdown leaderboard
+
+bitbudget indexes                                   # list indexes (organisation axis)
+bitbudget bench-index --synthetic 100000 128        # recall vs QPS vs bytes: flat/hnsw/ivfpq/bittrie
+```
+
+`run` embeds (torch) and evaluates (numpy) in one process. The corpora auto‑download.
+
+### The organisation axis (`bench-index`)
+
+The compression leaderboard answers *quality per byte*; `bench-index` answers the orthogonal
+*recall per query-second*. It builds an index over the document vectors and reports recall@k,
+throughput (QPS) and bytes per vector, so HNSW and IVF‑PQ (which buy throughput and *add* bytes)
+can be compared against compact‑code indexes on one frontier. Run it on synthetic data, on a
+cached embedding (`--embedder mxbai --corpus scifact`), or on your own vectors (`--npz`). The
+faiss‑backed indexes need `pip install bitbudget[faiss]`; the numpy `bittrie` runs without it.
+
+The `bittrie` index ships a small C kernel (`_bittrie.c`) for the query hot‑path, compiled on
+first use and cached (no compiler needed to *install* — the wheel stays pure‑Python, and it falls
+back to numpy if no compiler is present). It builds **multithreaded** when OpenMP is available
+(GCC/clang on Linux, Homebrew `libomp` on macOS) and single‑threaded otherwise; results are
+bit‑identical to the numpy path, and recall/footprint are algorithmic and unchanged either way.
+
+Because faiss carries its own OpenMP runtime, it cannot share a process with the bit‑trie's
+`libomp` on macOS. `bench-index` therefore runs the faiss indexes and the bit‑trie in **separate
+subprocesses** and merges the results, so a single `bitbudget bench-index ...` works everywhere
+(pass `--no-split` to force one process, e.g. on Linux where both share one OpenMP runtime).
+
+> **macOS note.** torch and faiss each bundle their own OpenMP runtime and crash if imported
+> in the same process. The core methods are numpy‑only, so `run` is safe; if you add a
+> faiss‑backed method, run `bitbudget embed` (torch) and `bitbudget eval` (numpy/faiss)
+> as separate processes.
+
+## The protocol (frozen, so results are comparable)
+
+- **Corpora:** the BEIR subsets `scifact`, `nfcorpus`, `arguana`, `fiqa` (small enough to run
+  on a laptop, diverse enough to be honest). Numbers are the mean over corpora; `±` is the
+  standard deviation across them.
+- **Metrics:** `nDCG@10` against the graded BEIR judgements, and `recall@10` against the exact
+  floating‑point neighbours. `% of float` is nDCG relative to the uncompressed embedding.
+- **Memory:** bytes stored per document vector (`4D` float, `D` int8, `D/8` binary, `M` for an
+  `M`‑byte product code, `4·dim` for a truncated/PCA‑reduced vector).
+- **Embedders:** `minilm` (384‑d) and `mxbai` (1024‑d, Matryoshka) ship built in.
+
+## Add your method in five lines
+
+This is the point of the benchmark: drop in your compressor and it is scored against every
+built‑in on the same protocol.
+
+```python
+from bitbudget import method
+import numpy as np
+
+@method("my-2bit", bits=2)
+def my_2bit(demb, qemb):
+    codes = my_quantise(demb)                       # your compression
+    scores = qemb @ my_reconstruct(codes).T         # (queries x docs) similarity
+    return scores, demb.shape[1] * 2 / 8            # scores, bytes per stored vector
+```
+
+```bash
+bitbudget run --embedder mxbai --corpus scifact --methods my-2bit binary+rerank float32
+```
+
+Then open a pull request adding your row to [LEADERBOARD.md](LEADERBOARD.md). See
+[CONTRIBUTING.md](CONTRIBUTING.md).
+
+## Cite
+
+If BitBudget helps your work, please cite the survey:
+
+```bibtex
+@article{moran2025projection,
+  title   = {Projection and Quantisation: A Unifying View of Learning to Hash,
+             from Random Projections to the RAG Era},
+  author  = {Moran, Sean},
+  journal = {arXiv preprint arXiv:2510.04127},
+  year    = {2025}
+}
+```
+
+MIT licensed.

bitbudget-0.1.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+bitbudget/__init__.py,sha256=83ZS2fMiaZojifkwbfMYYTXL8SwIHG9GPcL8TykB-sE,1157
+bitbudget/_bittrie.c,sha256=jbFyD_NwvkupJtC9DfgvYz0XuHTZQe4B-hdXI4gj0Hc,4496
+bitbudget/_bittrie_build.py,sha256=3HCcmnfWbVzQD4BVPM9z8aDT8r6N-reEShfRtxgbf6c,3672
+bitbudget/bittrie.py,sha256=3EQa1Qu6VAf8nxm8_sGvD60Iwbzz4ncKM8Ar3xPfIAs,5516
+bitbudget/cli.py,sha256=6dubvJF7yas9ytXm-vyUEhlvzIPFMVuNbGdUjM-slsE,10992
+bitbudget/datasets.py,sha256=0lQJbD3ALfIQF7wzCT-gIkMJ8tD1pm3z3EM4IpB_mOQ,1799
+bitbudget/embedders.py,sha256=_F6pHVRhQaK724317Fc1uIjDmeTnPKaV5nuAzdoZ9Bk,1909
+bitbudget/eval.py,sha256=O3nUaD6plbeOgrgD57EOrkNA54zgfPOrI2w-aapjzR0,2090
+bitbudget/indexes.py,sha256=yQ2O3NzcPoaONdpLagZrNz_AzT8NQcNQN70R4FXnBlM,5830
+bitbudget/methods.py,sha256=D4B66oKXUmq5xezc0w9WHQ9nThOY9dhagbuvWWzxiWk,3909
+bitbudget/metrics.py,sha256=7K2D3lgtpUDsACP6smPjytw5lbckWjE6HRpdzKdzCs8,1450
+bitbudget-0.1.0.dist-info/licenses/LICENSE,sha256=nUSDmf1Ud4HyOItkD93o4_Vs-hHzoheLsIW_BI9zfEo,1067
+bitbudget-0.1.0.dist-info/METADATA,sha256=wIstOwmYomqfT8Chz4mzZefQiwOGNvgLmct3OX3qXiU,6717
+bitbudget-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+bitbudget-0.1.0.dist-info/entry_points.txt,sha256=f-1IR4Yvl1Zygj95DC9w0LgaSPOIjsFQg6Oz4v8tPeo,49
+bitbudget-0.1.0.dist-info/top_level.txt,sha256=jBmwsZBMtA-eonXvguqSaVB05M8M6W1-UAl3OY6VFls,10
+bitbudget-0.1.0.dist-info/RECORD,,

bitbudget-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

bitbudget-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+bitbudget = bitbudget.cli:main

bitbudget-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Sean Moran
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

bitbudget-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+bitbudget