Upload build/audit_contamination.py with huggingface_hub

build/audit_contamination.py

@@ -39,18 +39,26 @@ P = np.uint64(0x100000001B3)          # FNV-1a prime, odd, so invertible mod 2^6
 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)
+# task -> (repo, config(s) to try in order, splits treated as reference material; 'test' deliberately
+# absent). Split names were read off the dataset server on 2026-09-19 rather than assumed: gsm8k has
+# train/test only, truthfulqa's repo is `truthfulqa/truthful_qa` and exposes `validation` only, and
+# `sylinrl/TruthfulQA` -- the name the first version of this file used -- no longer exists.
 REFS = [
     ("arc_easy", "allenai/ai2_arc", ["ARC-Easy"], ["train", "validation"]),
     ("arc_challenge", "allenai/ai2_arc", ["ARC-Challenge"], ["train", "validation"]),
     ("hellaswag", "Rowan/hellaswag", [None], ["train", "validation"]),
-    ("piqa", "baber/piqa", [None], ["train", "validation"]),
+    ("piqa", "baber/piqa", [None, "plain_text"], ["train", "validation"]),
     ("winogrande", "allenai/winogrande", ["winogrande_xl"], ["train", "validation"]),
     ("mmlu", "cais/mmlu", ["all"], ["dev", "validation"]),
-    ("truthfulqa", "sylinrl/TruthfulQA", [None], ["validation"]),
-    ("gsm8k", "openai/gsm8k", ["main"], ["train", "validation"]),
+    ("truthfulqa", "truthfulqa/truthful_qa", ["multiple_choice", "generation"], ["validation"]),
+    ("gsm8k", "openai/gsm8k", ["main"], ["train"]),
 ]
 TEXT_KEYS = ("question", "text", "choices", "story", "sentence", "prompt", "ctx", "input", "query")
+# A reference split that yields fewer rows than this did not load, it merely failed to raise: the first
+# run of this audit reported `piqa::default::train: 0 items -> 0 grams` with no exception, which a
+# coverage-blind verdict would have scored as "no contamination found in PIQA".
+MIN_ITEMS_PER_SPLIT = 200
+
 
 
 def window_hashes(ids, k=K):
@@ -95,7 +103,7 @@ def flatten(v, out):
 
 
 def build_reference(tk, max_items_per_split):
-    """Returns (sorted uint64 array of every reference gram, parallel uint16 task-bitmask, per-task
+    """Returns (sorted uint64 array of every reference gram, parallel uint16 task-bitmask, per-split
     item counts, per-reference errors). Only counts and hashes are ever held or printed."""
     import datasets
 
@@ -103,17 +111,21 @@ def build_reference(tk, max_items_per_split):
     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}"
+        for sp in splits:
+            # One split, tried against each candidate config in order, until a loader yields rows. The
+            # alternative -- trusting the first config that does not raise -- is what scored PIQA as
+            # clean while reading 0 items from it.
+            key, g, loaded = f"{task}::{repo}::{sp}", [], None
+            for cfg in configs:
+                ckey = f"{key}[{cfg or 'default'}]"
                 try:
                     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:
-                    errors[key] = f"{type(e).__name__}: {str(e)[:140]}"
-                    print(f"  ref {key}: UNREADABLE {type(e).__name__}", flush=True)
+                    errors[ckey] = f"{type(e).__name__}: {str(e)[:140]}"
+                    print(f"  ref {ckey}: UNREADABLE {type(e).__name__}: {str(e)[:80]}", flush=True)
                     continue
-                g, n = [], 0
+                n = 0
                 for row in ds:
                     if n >= max_items_per_split:
                         break
@@ -129,12 +141,20 @@ def build_reference(tk, max_items_per_split):
                     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)
+                counts[ckey] = n
+                print(f"  ref {ckey}: {n} items -> {sum(a.size for a in g):,} grams", flush=True)
+                if n >= MIN_ITEMS_PER_SPLIT:
+                    loaded = ckey
+                    break
+                errors[ckey] = f"only {n} rows (needs >= {MIN_ITEMS_PER_SPLIT}); trying next config"
+            if loaded is None:
+                errors[key] = f"no config of {repo} yielded >= {MIN_ITEMS_PER_SPLIT} rows for {sp}"
+                continue
+            per_task = [a for a in g]
+            if per_task:
+                arr = np.concatenate(per_task)
+                grams.append(arr)
+                masks.append(np.full(arr.shape, bit, dtype=np.uint16))
     if not grams:
         raise SystemExit("no reference grams built at all -- the audit would be a pass over nothing")
     allg = np.concatenate(grams)
@@ -143,10 +163,89 @@ def build_reference(tk, max_items_per_split):
     return allg[order], allm[order], counts, errors
 
 
+
+def gram_lookup(g, ref, refmask):
+    """(number of hits, the task bitmask of the matched reference grams) for one document.
+
+    numpy only, because this runs ~1.1 M times. `searchsorted` into the sorted reference table plus the
+    parallel bitmask is what makes per-task attribution possible without holding eight sets.
+    """
+    if g.size == 0:
+        return 0, None
+    pos = np.clip(np.searchsorted(ref, g), 0, ref.size - 1)
+    eq = ref[pos] == g
+    if not eq.any():
+        return 0, None
+    return int(eq.sum()), refmask[pos[eq]]
+
+
+def per_task_counts(bits):
+    return [int((bits & np.uint16(1 << ti)).astype(np.uint16).sum()) for ti in range(len(REFS))]
+
+
+def scan_docs(ref, refmask, docs):
+    """Fold an iterator of (ordinal_within_source, raw_bytes) into counts.
+
+    The ordinal is the only identity a document ever gets here: a filter pass can drop it by index later
+    without anything in this process knowing what it said. `hit_docs_over_half` is the column that decides
+    interpretation -- contamination by verbatim inclusion shows up as documents whose grams are *mostly*
+    matched, while one shared boilerplate sentence shows up as a single gram in a handful of documents.
+    """
+    import time
+    import hashlib
+    t0 = time.time()
+    s = {"docs": 0, "tokens": 0, "grams": 0, "hit_docs": 0, "hits": 0, "hit_docs_over_half": 0,
+         "hit_tokens": 0, "dup_docs": 0, "dup_tokens": 0, "per_task": [0] * len(REFS), "ordinals": []}
+    seen = set()
+    names = [t[0] for t in REFS]
+    for idx, raw in docs:
+        b = bytes(raw)
+        ids = np.frombuffer(b, dtype=np.uint16)
+        s["docs"] += 1
+        s["tokens"] += int(ids.size)
+        dg = hashlib.blake2b(b, digest_size=8).digest()
+        if dg in seen:                      # counted, never removed: de-duplication is out of scope (D-014)
+            s["dup_docs"] += 1
+            s["dup_tokens"] += int(ids.size)
+        else:
+            seen.add(dg)
+        g = window_hashes(ids)
+        s["grams"] += int(g.size)
+        h, bits = gram_lookup(g, ref, refmask)
+        if h:
+            s["hit_docs"] += 1
+            s["hits"] += h
+            s["hit_tokens"] += int(ids.size)
+            s["per_task"] = [a + c for a, c in zip(s["per_task"], per_task_counts(bits))]
+            s["ordinals"].append(idx)
+            if h * 2 >= int(g.size):
+                s["hit_docs_over_half"] += 1
+    s["seconds"] = round(time.time() - t0, 1)
+    s["per_task_named"] = {n: c for n, c in zip(names, s["per_task"]) if c}
+    return s
+
+
+def doc_stream(BM, shard_paths):
+    """(ordinal, raw document bytes) over a fixed list of shards. Ordinal order is the shard order in
+    the source's own record.json, which is exactly what the merge replays -- so an index computed here
+    still means the same document after a session restart."""
+    n = 0
+    for p in shard_paths:
+        for d in BM.iter_shard_docs(p):
+            yield n, d
+            n += 1
+
+
 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("--stage-subdir", default="stage",
+                    help="per-source staged shards live here under --root; scanning them is what makes "
+                         "attribution possible, because every merged shard contains all sources at once")
+    ap.add_argument("--max-items-per-split", type=int, default=50000)
+    ap.add_argument("--write-filter", default="",
+                    help="directory for <source>.u8 exclusion masks + filter.json (read back by "
+                         "build_mix --exclude-dir)")
     ap.add_argument("--out", default="", help="write the report JSON here (publish_mix expects "
                                              "mixroot/audit.json)")
     a = ap.parse_args()
@@ -158,77 +257,138 @@ def main():
 
     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)
+    tasks_covered = sorted({k.split("::")[0] for k, v in refcounts.items()
+                            if v >= MIN_ITEMS_PER_SPLIT})
+    print(f"reference: {ref.size:,} grams, {len(tasks_covered)}/8 tasks covered, "
+          f"{len(errors)} loader problems", flush=True)
 
-    man = json.load(open(os.path.join(a.root, "manifest.json")))
     import build_mix as BM
 
+    # ---- per-source attribution, over the staged (pre-merge) shards ------------------------------
+    # The first version attributed merged hits as `c // len(rec["sources"])` to every source named in the
+    # shard. All 141 shards name all 15 sources, so every source received an almost identical number and
+    # the table looked like a finding while being arithmetically meaningless.
+    stage_root = os.path.join(a.root, a.stage_subdir)
+    by_source, masks = {}, {}
+    if os.path.isdir(stage_root):
+        for key in sorted(os.listdir(stage_root)):
+            sd = os.path.join(stage_root, key)
+            rd = os.path.join(sd, "record.json")
+            if not os.path.isdir(sd) or not os.path.exists(rd):
+                continue
+            rec = json.load(open(rd))
+            paths = [os.path.join(a.root, s["file"]) for s in rec.get("shards", [])]
+            absent = [p for p in paths if not os.path.exists(p)]
+            if absent:
+                errors[f"stage::{key}"] = f"{len(absent)} staged shard(s) missing on disk"
+                print(f"  stage {key}: {len(absent)} shard(s) missing, skipped", flush=True)
+                continue
+            s = scan_docs(ref, refmask, doc_stream(BM, paths))
+            s["shards"] = len(paths)
+            s["stage_complete"] = bool(rec.get("complete"))
+            by_source[key] = s
+            masks[key] = (np.asarray(sorted(s["ordinals"]), dtype=np.int64), s["docs"])
+            print(f"  stage {key}: {s['docs']:,} docs {s['tokens']:,} tok -> hits={s['hits']:,} in "
+                  f"{s['hit_docs']:,} docs ({s['hit_docs_over_half']:,} of them >50 % matched) "
+                  f"[{s['seconds']}s] {json.dumps(s['per_task_named'])}", flush=True)
+            s.pop("ordinals")
+    else:
+        print(f"no staged sources under {stage_root}: attribution unavailable, totals only", flush=True)
+
+    # ---- the verdict is always measured on the bytes that will actually train ---------------------
+    man_path = os.path.join(a.root, "manifest.json")
     hits = {}
-    hits_by_source = {}
-    docs = grams_seen = 0
+    docs = grams_seen = mix_tokens = hit_docs = hit_docs_over_half = 0
+    dup_docs = dup_tokens = 0
     missing = []
-    for rec in man["shards"]:
-        p = os.path.join(a.root, rec["file"])
-        if not os.path.exists(p):
-            # 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
-            ids = np.frombuffer(bytes(d), dtype=np.uint16)
-            g = window_hashes(ids)
-            if g.size == 0:
+    if os.path.exists(man_path):
+        man = json.load(open(man_path))
+        for rec in man["shards"]:
+            p = os.path.join(a.root, rec["file"])
+            if not os.path.exists(p):
+                # A shard not on disk is unmeasured, not clean (E-017's shape).
+                missing.append(rec["file"])
                 continue
-            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)
+            s = scan_docs(ref, refmask, doc_stream(BM, [p]))
+            docs += s["docs"]
+            mix_tokens += s["tokens"]
+            grams_seen += s["grams"]
+            hit_docs += s["hit_docs"]
+            hit_docs_over_half += s["hit_docs_over_half"]
+            dup_docs += s["dup_docs"]
+            dup_tokens += s["dup_tokens"]
+            for n, c in s["per_task_named"].items():
+                hits[n] = hits.get(n, 0) + c
+        total_overlap = sum(hits.values())
+        print(f"  mix: {docs:,} docs, {grams_seen:,} grams, {total_overlap:,} hits in {hit_docs:,} docs "
+              f"({hit_docs_over_half:,} majority-hit)", flush=True)
+    else:
+        total_overlap = 0
+        print(f"no manifest at {a.root}: merged mix not scanned", flush=True)
+
+    expected_fp = grams_seen * int(np.unique(ref).size) / float(1 << 64)
+    noise_allow = max(10.0, 10.0 * expected_fp)
+    per_task_named = {t[0]: hits.get(t[0], 0) for t in REFS}
+
+    if a.write_filter:
+        os.makedirs(a.write_filter, exist_ok=True)
+        summary = {"k": K, "criterion": "any 13-gram of the document appears in a reference "
+                                        "train/validation/dev split", "sources": {}}
+        for key, (ord_sorted, n_docs) in masks.items():
+            m = np.zeros(n_docs, dtype=np.uint8)
+            if ord_sorted.size:
+                m[ord_sorted] = 1
+            m.tofile(os.path.join(a.write_filter, key + ".u8"))
+            s = by_source[key]
+            summary["sources"][key] = {"docs": s["docs"], "tokens": s["tokens"],
+                                       "dropped_docs": int(ord_sorted.size),
+                                       "dropped_tokens": int(s["hit_tokens"]),
+                                       "kept_docs": int(n_docs - ord_sorted.size)}
+        with open(os.path.join(a.write_filter, "filter.json"), "w") as f:
+            json.dump(summary, f, indent=1, default=str)
+        print(f"filter masks written to {a.write_filter} for {len(masks)} source(s)", flush=True)
+
+    contaminated_sources = {k: v["hit_docs"] / max(1, v["docs"]) for k, v in by_source.items()
+                            if v["hit_docs"]}
     report = {
         "k": K,
         "fingerprint": "multiplicative rolling hash mod 2^64, P=0x100000001B3",
         "reference_items_by_split": refcounts,
-        "reference_unreadable": errors,
+        "reference_loader_problems": errors,
         "reference_grams_distinct": int(np.unique(ref).size),
         "reference_grams_total": int(ref.size),
-        "shards_expected": len(man["shards"]),
+        "tasks_covered": tasks_covered,
         "shards_missing": missing,
         "mix_documents_audited": docs,
+        "mix_tokens_audited": mix_tokens,
         "mix_grams_audited": grams_seen,
-        "overlap_by_task": hits,
-        "overlap_by_source": hits_by_source,
-        "expected_false_positives": round(expected_fp, 3),
+        "mix_documents_with_any_hit": hit_docs,
+        "mix_documents_majority_hit": hit_docs_over_half,
+        "overlap_by_task": per_task_named,
         "overlap_total": total_overlap,
-        "overlap_exceeds_noise": total_overlap > max(10.0, 3.0 * expected_fp),
+        "by_source": by_source,
+        "source_hit_document_fraction": {k: round(v, 6) for k, v in sorted(
+            contaminated_sources.items(), key=lambda kv: -kv[1])},
+        "duplicate_documents": dup_docs,
+        "duplicate_tokens": dup_tokens,
+        "expected_false_positives": round(expected_fp, 3),
+        "noise_allowance": round(noise_allow, 3),
+        "overlap_exceeds_noise": bool(total_overlap > noise_allow),
         "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)),
+                 "per task -- never test, per §3.3. overlap_by_task sums gram hits, so one verbatim "
+                 "benchmark document contributes as many hits as it has windows; read "
+                 "mix_documents_majority_hit and source_hit_document_fraction to judge severity. "
+                 "Duplicate counting is diagnostic only: de-duplication is out of scope (D-014), "
+                 "contamination removal is not (§3.3)."),
+        "COVERED_ALL_SHARDS": bool(not missing and docs > 0),
+        "MEASURED": bool(docs > 0 and grams_seen > 0 and not errors and len(tasks_covered) == 8),
     }
+    report["AUDIT_PASSED"] = bool(report["MEASURED"] and total_overlap <= noise_allow)
     txt = json.dumps(report, indent=1, default=str)
-    print(f"VERDICT AUDIT_PASSED={report['AUDIT_PASSED']} docs={docs:,} overlap={total_overlap} "
-          f"expected_fp={expected_fp:.2f} missing_shards={len(missing)} refs_ok={len(refcounts)} "
-          f"refs_failed={len(errors)}")
+    print(f"VERDICT AUDIT_PASSED={report['AUDIT_PASSED']} MEASURED={report['MEASURED']} "
+          f"docs={docs:,} overlap={total_overlap} hit_docs={hit_docs:,} "
+          f"majority_hit_docs={hit_docs_over_half:,} expected_fp={expected_fp:.3f} "
+          f"tasks={len(tasks_covered)}/8 problems={len(errors)}")
     print("AUDIT_JSON_BEGIN")
     print(txt)
     print("AUDIT_JSON_END")