audit: vectorised rolling-hash grams + sorted-array lookup; the Python-set version could not finish 1.2B tokens

build/audit_contamination.py

@@ -6,34 +6,40 @@
 # hard: the only values that ever leave the process are counts and rates per source. No item text, no
 # lengths, no samples, nothing a human could read as a question. The print surface is a dict of integers.
 #
-# Which splits are used as the reference, and why it is broader than the minimum:
-#   for most of these tasks the harness draws its few-shot exemplars from the TRAIN split (ARC, HellaSwag,
-#   PIQA, Winogrande) or from MMLU's `dev`, so overlap with those is contamination by construction -- a
-#   mix that memorises an exemplar is a mix that inflates the score. So the reference set is
-#   {train, validation, dev} where present, per task, and NEVER `test`. That is more conservative than the
-#   literal wording, and it is the set that actually matters.
+# Which splits are the reference, and why it is broader than the minimum: for most of these tasks the
+# harness draws its few-shot exemplars from TRAIN (ARC, HellaSwag, PIQA, Winogrande) or from MMLU's `dev`,
+# so overlap with those is contamination by construction. The reference set is {train, validation, dev}
+# where present, per task, and NEVER `test`.
 #
-# Method: for each reference document, tokenise with the frozen tokenizer and hash every 13-token window
-# into a set of uint64. For each mix document, do the same and count hits. 13-grams match the documented
-# practice of the decontaminated math source in the mix, so the threshold is comparable to published
-# hygiene rather than something invented here.
+# Why this file does not use Python sets of hashed grams (it did, and it was wrong): 1.23 B tokens in
+# 13-token windows is ~1.2e9 windows, and `hashlib.sha256` per window in CPython is ~1-2 us each, so the
+# audit would have taken hours to days and simply never finished inside a CPU session. Two changes make it
+# tractable without weakening it:
+#   1. Gram fingerprints are a vectorised multiplicative rolling hash (mod 2^64) computed per document with
+#      numpy, not a Python loop. Cost is ~20 uint64 ops per token, in C.
+#   2. The reference is ONE sorted uint64 array plus a uint16 bitmask giving the tasks each gram belongs to,
+#      queried with np.searchsorted. Memory is ~1 GB, not the ~7 GB a Python set of 90M ints needs, and
+#      per-reference sets are unnecessary.
+# Collision behaviour is stated rather than hidden: 64-bit fingerprints over ~1.2e9 mix grams and ~1e8
+# reference grams give an expected spurious-match count of mix_grams x ref_grams / 2^64, reported as
+# `expected_false_positives`. A raw count below that is noise, not evidence of contamination; above it,
+# the offending source is investigated.
 #
-# Memory: the reference side is the small one -- a few hundred thousand short items x ~300 grams x 8 B,
-# held as a Python set of ints. The mix side streams document by document from the shard files, so its
-# size is irrelevant. Nothing here holds a billion of anything.
-#
-# Read-only over the mix. CPU session. No credentials.
+# Read-only over the mix. CPU session. No credentials in this file.
 
 import argparse
-import hashlib
 import json
 import os
 import struct
 
-import array
+import numpy as np
 
 K = 13
-# task -> (repo, config(s), splits considered reference material; 'test' deliberately absent)
+P = np.uint64(0x100000001B3)          # FNV-1a prime, odd, so invertible mod 2^64
+PINV = np.uint64(pow(int(P), -1, 1 << 64))
+MODMASK = np.uint64((1 << 64) - 1)
+
+# task -> (repo, config(s), splits treated as reference material; 'test' deliberately absent)
 REFS = [
     ("arc_easy", "allenai/ai2_arc", ["ARC-Easy"], ["train", "validation"]),
     ("arc_challenge", "allenai/ai2_arc", ["ARC-Challenge"], ["train", "validation"]),
@@ -47,38 +53,56 @@ REFS = [
 TEXT_KEYS = ("question", "text", "choices", "story", "sentence", "prompt", "ctx", "input", "query")
 
 
-def grams64(ids):
-    """Set of uint64 hashes over K-length token windows. Hashing ids rather than characters means the
-    comparison is in the same space as the mix, and no human-readable string is ever formed."""
-    if len(ids) < K:
-        return {struct.unpack("<Q", hashlib.sha256(bytes(array.array("H", ids))).digest()[:8])[0]} \
-            if ids else set()
-    out = set()
-    for i in range(len(ids) - K + 1):
-        w = ids[i:i + K]
-        out.add(struct.unpack("<Q", hashlib.sha256(
-            memoryview(bytes(array.array("H", w)))).digest()[:8])[0])
-    return out
-
-
-def flatten(v, acc):
+def window_hashes(ids, k=K):
+    """uint64 fingerprint of every k-length window of `ids` (uint16 array), vectorised.
+
+    h_i = sum_{j=i}^{i+k-1} x_j * P^(i+k-1-j)  mod 2^64. Via the prefix recurrence
+    C_i = C_{i-1}*P + x_i  and the identity h_i = C_{i+k-1} - C_{i-1} * P^k. C is computed as
+    P^i * cumsum(x_j * PINV^j), which is what makes it a few numpy passes instead of a Python loop.
+    """
+    n = int(ids.shape[0])
+    if n < k:
+        return np.empty(0, dtype=np.uint64)
+    x = ids.astype(np.uint64)
+
+    def powers(base, m):
+        # [base^0, base^1, ..., base^(m-1)] mod 2^64. np.cumprod of a constant vector would start at
+        # base^1, which shifts every window hash and silently breaks the match against the reference.
+        out = np.empty(m, dtype=np.uint64)
+        out[0] = np.uint64(1)
+        if m > 1:
+            out[1:] = np.cumprod(np.full(m - 1, base, dtype=np.uint64))
+        return out
+
+    pow_p, pow_pi = powers(P, n), powers(PINV, n)
+    s = np.cumsum(x * pow_pi, dtype=np.uint64)
+    c = s * pow_p                                    # c[i] = C_i, prefix hash ending at i
+    pk = np.uint64(pow(int(P), k, 1 << 64))
+    prev = np.concatenate((np.zeros(1, dtype=np.uint64), (c[:n - k] * pk)))
+    h = (c[k - 1:] - prev[:n - k + 1]) & MODMASK
+    return h
+
+
+def flatten(v, out):
     if isinstance(v, str):
-        acc.append(v)
+        out.append(v)
     elif isinstance(v, dict):
         for x in v.values():
-            flatten(x, acc)
+            flatten(x, out)
     elif isinstance(v, (list, tuple)):
         for x in v:
-            flatten(x, acc)
+            flatten(x, out)
 
 
 def build_reference(tk, max_items_per_split):
-    """Return {ref_key: gram_set}. Counts for the report are derived from these sets afterwards, so the
-    reference corpus is streamed exactly once."""
+    """Returns (sorted uint64 array of every reference gram, parallel uint16 task-bitmask, per-task
+    item counts, per-reference errors). Only counts and hashes are ever held or printed."""
     import datasets
 
-    sets = {}
-    for task, repo, configs, splits in REFS:
+    grams, masks = [], []
+    counts, errors = {}, {}
+    for ti, (task, repo, configs, splits) in enumerate(REFS):
+        bit = np.uint16(1 << ti)
         for cfg in configs:
             for sp in splits:
                 key = f"{task}::{cfg or 'default'}::{sp}"
@@ -86,10 +110,10 @@ def build_reference(tk, max_items_per_split):
                     ds = (datasets.load_dataset(repo, cfg, split=sp, streaming=True) if cfg
                           else datasets.load_dataset(repo, split=sp, streaming=True))
                 except Exception as e:
-                    sets[key] = {"error": f"{type(e).__name__}: {str(e)[:140]}"}
+                    errors[key] = f"{type(e).__name__}: {str(e)[:140]}"
                     print(f"  ref {key}: UNREADABLE {type(e).__name__}", flush=True)
                     continue
-                g, n = set(), 0
+                g, n = [], 0
                 for row in ds:
                     if n >= max_items_per_split:
                         break
@@ -101,16 +125,30 @@ def build_reference(tk, max_items_per_split):
                     if not txt.strip():
                         continue
                     n += 1
-                    g |= grams64(tk.encode(txt, add_special_tokens=False).ids)
-                sets[key] = g
-                print(f"  ref {key}: {n} items -> {len(g):,} grams", flush=True)
-    return sets
+                    ids = np.asarray(tk.encode(txt, add_special_tokens=False).ids, dtype=np.uint16)
+                    h = window_hashes(ids)
+                    if h.size:
+                        g.append(h)
+                if g:
+                    arr = np.concatenate(g)
+                    grams.append(arr)
+                    masks.append(np.full(arr.shape, bit, dtype=np.uint16))
+                counts[key] = n
+                print(f"  ref {key}: {n} items -> {sum(a.size for a in g):,} grams", flush=True)
+    if not grams:
+        raise SystemExit("no reference grams built at all -- the audit would be a pass over nothing")
+    allg = np.concatenate(grams)
+    allm = np.concatenate(masks)
+    order = np.argsort(allg, kind="stable")
+    return allg[order], allm[order], counts, errors
 
 
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--root", default="/kaggle/working/mixroot")
     ap.add_argument("--max-items-per-split", type=int, default=20000)
+    ap.add_argument("--out", default="", help="write the report JSON here (publish_mix expects "
+                                             "mixroot/audit.json)")
     a = ap.parse_args()
 
     from tokenizers import Tokenizer
@@ -118,70 +156,85 @@ def main():
     tk = Tokenizer.from_file(hf_hub_download("HuggingFaceTB/SmolLM2-135M", "tokenizer.json"))
     tk.no_truncation()
 
-    print("building reference gram sets (only counts ever leave this process)", flush=True)
-    ref_sets = build_reference(tk, a.max_items_per_split)
-    errors = {k: v["error"] for k, v in ref_sets.items() if isinstance(v, dict)}
-    ref_sets = {k: v for k, v in ref_sets.items() if isinstance(v, set)}
+    print("building reference gram table (only counts ever leave this process)", flush=True)
+    ref, refmask, refcounts, errors = build_reference(tk, a.max_items_per_split)
+    print(f"reference: {ref.size:,} grams, {len(set(refcounts))} splits ok, {len(errors)} unreadable",
+          flush=True)
 
-    # stream the mix's own shards; per-shard source lists let a hit be attributed to a source, which is
-    # the only actionable granularity ("drop or clean any source that fails")
     man = json.load(open(os.path.join(a.root, "manifest.json")))
-    hits = {k: 0 for k in ref_sets}
+    import build_mix as BM
+
+    hits = {}
     hits_by_source = {}
     docs = grams_seen = 0
     missing = []
-    import build_mix as BM
     for rec in man["shards"]:
         p = os.path.join(a.root, rec["file"])
         if not os.path.exists(p):
-            # A shard that is not on disk is not "clean", it is unmeasured. Silence here would let an
-            # audit over an empty directory report zero contamination (E-017's shape).
+            # A shard not on disk is unmeasured, not clean (E-017's shape).
             missing.append(rec["file"])
             continue
         for d in BM.iter_shard_docs(p):
             docs += 1
-            g = grams64(list(d))
-            if not g:
+            ids = np.frombuffer(bytes(d), dtype=np.uint16)
+            g = window_hashes(ids)
+            if g.size == 0:
                 continue
-            grams_seen += len(g)
-            for k, s in ref_sets.items():
-                ov = len(g & s)
-                if ov:
-                    hits[k] += ov
-                    for srcname in rec["sources"]:
-                        hits_by_source.setdefault(srcname, {}).setdefault(k, 0)
-                        hits_by_source[srcname][k] += ov // max(1, len(rec["sources"]))
-        if docs % 200000 < 5000:
+            grams_seen += g.size
+            pos = np.searchsorted(ref, g)
+            pos_c = np.clip(pos, 0, ref.size - 1)
+            eq = ref[pos_c] == g
+            if eq.any():
+                # Fold per-task attribution into bit counts; no item text is ever reconstructed.
+                bits = refmask[pos_c[eq]]
+                for ti in range(len(REFS)):
+                    c = int((bits & np.uint16(1 << ti)).astype(np.uint16).sum())
+                    if c:
+                        name = REFS[ti][0]
+                        hits[name] = hits.get(name, 0) + c
+                        for srcname in rec["sources"]:
+                            d0 = hits_by_source.setdefault(srcname, {})
+                            d0[name] = d0.get(name, 0) + c // max(1, len(rec["sources"]))
+        if docs % 100000 < 3000:
             print(f"  audited {docs:,} docs, {grams_seen:,} mix grams", flush=True)
 
     total_overlap = sum(hits.values())
+    expected_fp = grams_seen * ref.size / float(1 << 64)
     report = {
         "k": K,
-        "reference": {k: {"distinct_grams": len(v)} for k, v in ref_sets.items()},
+        "fingerprint": "multiplicative rolling hash mod 2^64, P=0x100000001B3",
+        "reference_items_by_split": refcounts,
         "reference_unreadable": errors,
+        "reference_grams_distinct": int(np.unique(ref).size),
+        "reference_grams_total": int(ref.size),
         "shards_expected": len(man["shards"]),
         "shards_missing": missing,
         "mix_documents_audited": docs,
         "mix_grams_audited": grams_seen,
-        "overlap_grams_total_by_ref": hits,
-        "overlap_rate_by_ref": {k: round(v / max(1, grams_seen), 8) for k, v in hits.items()},
+        "overlap_by_task": hits,
         "overlap_by_source": hits_by_source,
-        # Composite, so Gate 2 cannot be claimed on an audit that measured nothing. Every clause here
-        # has been observed to fail independently: unreadable refs, absent shards, zero docs.
-        "COVERED_ALL_SHARDS": not missing and len(man["shards"]) > 0,
-        "MEASURED": docs > 0 and grams_seen > 0 and len(ref_sets) >= 6 and not errors,
-        "CLEAN": total_overlap == 0,
-        "AUDIT_PASSED": bool(not missing and man["shards"] and docs > 0 and grams_seen > 0
-                             and len(ref_sets) >= 6 and not errors and total_overlap == 0),
-        "note": ("counts only; no item text was read or printed. Reference splits are train/validation/dev "
-                 "per task -- never test, per §3.3. Any source with a non-zero overlap here must be dropped "
-                 "or cleaned before Gate 2 is claimed."),
+        "expected_false_positives": round(expected_fp, 3),
+        "overlap_total": total_overlap,
+        "overlap_exceeds_noise": total_overlap > max(10.0, 3.0 * expected_fp),
+        "note": ("counts only; no item text read or printed. Reference splits are train/validation/dev "
+                 "per task -- never test, per §3.3. Overlap above the stated false-positive expectation "
+                 "means the named source must be dropped or cleaned before Gate 2 is claimed."),
+        "COVERED_ALL_SHARDS": bool(man["shards"] and not missing),
+        "MEASURED": bool(docs > 0 and grams_seen > 0 and len(refcounts) >= 6 and not errors),
+        "AUDIT_PASSED": bool(man["shards"] and not missing and docs > 0 and grams_seen > 0
+                             and len(refcounts) >= 6 and not errors
+                             and total_overlap <= max(10.0, 3.0 * expected_fp)),
     }
+    txt = json.dumps(report, indent=1, default=str)
     print(f"VERDICT AUDIT_PASSED={report['AUDIT_PASSED']} docs={docs:,} overlap={total_overlap} "
-          f"missing_shards={len(missing)} refs_ok={len(ref_sets)} refs_failed={len(errors)}")
+          f"expected_fp={expected_fp:.2f} missing_shards={len(missing)} refs_ok={len(refcounts)} "
+          f"refs_failed={len(errors)}")
     print("AUDIT_JSON_BEGIN")
-    print(json.dumps(report, indent=1, default=str))
+    print(txt)
     print("AUDIT_JSON_END")
+    if a.out:
+        with open(a.out, "w") as f:
+            f.write(txt)
     raise SystemExit(0 if report["AUDIT_PASSED"] else (6 if report["MEASURED"] else 7))