bloom_atlas.py

"""Slim, shippable variant of the offline existence atlas: a bloom filter over the 16.7M normalized titles.

Why: the exact sqlite atlas is 446MB — too big to ship inside a GGUF/Ollama-tier package. A bloom filter
trades ~33MB for a tunable false-POSITIVE rate and ZERO false-negatives. The no-false-negative property is
exactly right for grounding: every REAL title still matches (reals stay 100% grounded, same as sqlite); only
a small FPR can spuriously "ground" a fabrication (a false-rescue). We size for a low per-query FPR so the
structural zero-false-rescue guarantee degrades only negligibly.

Build:  python bloom_atlas.py build  wiki_titles.db  wiki_titles.bloom   [bits_per_item]
Probe:  python bloom_atlas.py probe  wiki_titles.bloom  ../eval/hl_battery.json
The BloomGrounder mirrors Grounder.grounded() candidate logic exactly, querying the bloom instead of sqlite.
"""
import sys, os, sqlite3, hashlib, struct, math, json, re
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from offline_atlas import norm, STOP, _sig

MAGIC = b"LBLOOM1\n"


def _two_hashes(b):
    d = hashlib.blake2b(b, digest_size=16).digest()
    return struct.unpack("<QQ", d)            # two 64-bit hashes for double-hashing


class Bloom:
    def __init__(self, m_bits, k):
        self.m = int(m_bits); self.k = int(k)
        self.bits = np.zeros((self.m + 7) // 8, dtype=np.uint8)

    def _pos(self, key):
        h1, h2 = _two_hashes(key.encode("utf-8"))
        return [(h1 + i * h2) % self.m for i in range(self.k)]

    def add(self, key):
        for p in self._pos(key):
            self.bits[p >> 3] |= (1 << (p & 7))

    def __contains__(self, key):
        for p in self._pos(key):
            if not (self.bits[p >> 3] >> (p & 7)) & 1:
                return False
        return True

    def save(self, path):
        with open(path, "wb") as f:
            f.write(MAGIC); f.write(struct.pack("<QI", self.m, self.k)); f.write(self.bits.tobytes())

    @classmethod
    def load(cls, path):
        with open(path, "rb") as f:
            assert f.read(len(MAGIC)) == MAGIC, "bad bloom magic"
            m, k = struct.unpack("<QI", f.read(12))
            b = cls(m, k); b.bits = np.frombuffer(f.read(), dtype=np.uint8).copy()
            return b


def build(db_path, out_path, bits_per_item=16):
    con = sqlite3.connect(db_path)
    n = con.execute("SELECT COUNT(*) FROM t").fetchone()[0]
    m = n * int(bits_per_item)
    k = max(1, round((m / n) * math.log(2)))
    fpr = (1 - math.exp(-k * n / m)) ** k
    print(f"n={n} titles, m={m} bits ({m/8/1e6:.1f} MB), k={k}, theoretical FPR={fpr:.4%}", flush=True)
    bl = Bloom(m, k)
    done = 0
    cur = con.execute("SELECT n FROM t")
    while True:
        rows = cur.fetchmany(200000)
        if not rows:
            break
        for (t,) in rows:
            bl.add(t)
        done += len(rows)
        if done % 2000000 == 0:
            print(f"  {done//1000000}M ...", flush=True)
    con.close()
    bl.save(out_path)
    print(f"built {out_path}: {os.path.getsize(out_path)/1e6:.1f} MB", flush=True)


class BloomGrounder:
    """Mirrors offline_atlas.Grounder.grounded() candidate logic, against the bloom filter."""
    def __init__(self, bloom_path):
        self.bl = Bloom.load(bloom_path)

    def _exists(self, nm):
        return nm in self.bl

    def _candidates(self, entity):
        c = [entity]
        if "," in entity:
            c.append(entity.split(",")[0])
        if " by " in entity:
            c.append(entity.split(" by ")[0])
        m = re.search(r"'s\s+(.+)", entity)
        if m:
            c.append(m.group(1))
        c2 = []
        for s in c:
            c2.append(s)
            if "(" in s:
                c2.append(re.sub(r"\([^)]*\)", "", s))
        return [s.strip() for s in c2 if s.strip()]

    def grounded(self, entity):
        for m in self._candidates(entity):
            nm = norm(m)
            if nm and self._exists(nm):
                return {"matched": True, "hit": nm}
        e = norm(entity); toks = e.split()
        esig = [t for t in toks if _sig(t)]
        if not esig:
            return {"matched": False, "hit": ""}
        for L in range(len(toks), 1, -1):
            for i in range(0, len(toks) - L + 1):
                w = toks[i:i + L]
                wsig = [t for t in w if _sig(t)]
                if len(wsig) >= 2 and len(wsig) / len(esig) >= 0.6:
                    wn = " ".join(w)
                    if self._exists(wn):
                        return {"matched": True, "hit": wn}
        return {"matched": False, "hit": ""}


def probe(bloom_path, battery_path):
    g = BloomGrounder(bloom_path)
    items = json.load(open(battery_path))["items"]
    y = np.array([it["label"] for it in items])
    key = "name" if "name" in items[0] else "entity"
    matched = np.array([g.grounded(it[key])["matched"] for it in items])
    fake, real = y == 1, y == 0
    print(f"bloom grounder on {os.path.basename(battery_path)} ({len(items)} items):")
    print(f"  REAL matched {matched[real].mean():.4f} (want 1.0 — bloom has no false-negatives)")
    print(f"  FAKE matched {matched[fake].mean():.4f} (want ~0 — these are bloom false-positives = false-rescues)")
    fr = [items[i][key] for i in range(len(items)) if fake[i] and matched[i]]
    print(f"  false-rescues (fakes matched): {len(fr)}  {fr[:10]}")


if __name__ == "__main__":
    cmd = sys.argv[1]
    if cmd == "build":
        build(sys.argv[2], sys.argv[3], int(sys.argv[4]) if len(sys.argv) > 4 else 16)
    elif cmd == "probe":
        probe(sys.argv[2], sys.argv[3])