Upload build/audit_contamination.py with huggingface_hub

build/audit_contamination.py

@@ -47,13 +47,15 @@ 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, "plain_text"], ["train", "validation"]),
+    ("piqa", "baber/piqa", [None], ["train", "validation"]),
     ("winogrande", "allenai/winogrande", ["winogrande_xl"], ["train", "validation"]),
     ("mmlu", "cais/mmlu", ["all"], ["dev", "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")
+TEXT_KEYS = ("question", "text", "choices", "story", "sentence", "prompt", "ctx", "input", "query",
+             "answer", "solution", "reasoning", "explanation", "goal", "sol1", "sol2", "abstractions",
+             "nonabstractions", "mc1_targets", "mc2_targets")
 # 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".
@@ -102,57 +104,141 @@ def flatten(v, out):
             flatten(x, out)
 
 
+SERVER = "https://datasets-server.huggingface.co"
+_SIZE_CACHE = {}
+
+
+def server_sizes(repo, cfg):
+    """Row counts the dataset server reports for each split of one config, or {} if it cannot be asked.
+
+    This exists because a streaming reader can stop early without raising, and because a loader that yields
+    0 rows of a 16,113-row split (exactly what `baber/piqa` did under the Kaggle image's `datasets`) is
+    indistinguishable from a clean result unless the expected count is checked. Never trust the count from
+    memory: the first version of this file expected 14,042 MMLU validation rows, and 14,042 is MMLU's
+    *test* split -- the one §3.3 forbids reading. `validation` is 1,531.
+    """
+    import urllib.parse
+    import urllib.request
+    ck = (repo, cfg)
+    if ck in _SIZE_CACHE:
+        return _SIZE_CACHE[ck]
+    q = {"dataset": repo}
+    if cfg:
+        q["config"] = cfg
+    out = {}
+    try:
+        with urllib.request.urlopen(SERVER + "/size?" + urllib.parse.urlencode(q), timeout=90) as r:
+            d = json.load(r)
+        out = {s["split"]: int(s.get("num_rows", 0)) for s in d.get("size", {}).get("splits", [])}
+    except Exception as e:
+        print(f"  sizes API {repo}/{cfg or ''}: {type(e).__name__}", flush=True)
+    _SIZE_CACHE[ck] = out
+    return out
+
+
+def api_rows(repo, cfg, sp, cap):
+    """Page the public rows endpoint (100 per request, full=true so no cell is silently truncated)."""
+    import urllib.parse
+    import urllib.request
+    off = 0
+    while off < cap:
+        q = {"dataset": repo, "split": sp, "offset": off, "limit": 100, "full": "true"}
+        if cfg:
+            q["config"] = cfg
+        try:
+            with urllib.request.urlopen(SERVER + "/rows?" + urllib.parse.urlencode(q),
+                                        timeout=120) as r:
+                d = json.load(r)
+        except Exception as e:
+            print(f"  rows API {repo}/{cfg or ''}/{sp} @{off}: {type(e).__name__}", flush=True)
+            return
+        batch = d.get("rows") or []
+        if not batch:
+            return
+        for item in batch:
+            yield item.get("row", item)
+        off += len(batch)
+
+
+def ingest(tk, rows, cap):
+    """Encode each row's text fields into gram windows. Returns (items kept, list of gram arrays)."""
+    g, n = [], 0
+    for row in rows:
+        if n >= cap:
+            break
+        parts = []
+        for k in TEXT_KEYS:
+            if k in row:
+                flatten(row[k], parts)
+        txt = " ".join(parts)
+        if not txt.strip():
+            continue
+        n += 1
+        ids = np.asarray(tk.encode(txt, add_special_tokens=False).ids, dtype=np.uint16)
+        h = window_hashes(ids)
+        if h.size:
+            g.append(h)
+    return n, g
+
+
 def build_reference(tk, max_items_per_split):
-    """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
+    """Returns (sorted uint64 reference grams, parallel uint16 task bitmask, per-split item counts,
+    errors, per-split expected row counts). Only counts and hashes are ever held or printed.
 
+    A split is accepted only if the rows actually read reach 90 % of what the dataset server says the
+    split holds; short reads are retried through the HTTP rows endpoint and, if still short, recorded as
+    loader problems -- which fail MEASURED rather than scoring as "no contamination found".
+    """
     grams, masks = [], []
-    counts, errors = {}, {}
+    counts, errors, srv = {}, {}, {}
+    try:
+        import datasets
+    except Exception as e:
+        datasets = None
+        errors["import datasets"] = f"{type(e).__name__}: {str(e)[:120]}"
     for ti, (task, repo, configs, splits) in enumerate(REFS):
         bit = np.uint16(1 << ti)
         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
+            key = f"{task}::{repo}::{sp}"
+            best = 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[ckey] = f"{type(e).__name__}: {str(e)[:140]}"
-                    print(f"  ref {ckey}: UNREADABLE {type(e).__name__}: {str(e)[:80]}", flush=True)
-                    continue
-                n = 0
-                for row in ds:
-                    if n >= max_items_per_split:
-                        break
-                    parts = []
-                    for k in TEXT_KEYS:
-                        if k in row:
-                            flatten(row[k], parts)
-                    txt = " ".join(parts)
-                    if not txt.strip():
+                want = server_sizes(repo, cfg).get(sp)
+                srv[key] = want
+                n, g = 0, []
+                if datasets is not None:
+                    try:
+                        ds = (datasets.load_dataset(repo, cfg, split=sp, streaming=True) if cfg
+                              else datasets.load_dataset(repo, split=sp, streaming=True))
+                        n, g = ingest(tk, ds, max_items_per_split)
+                    except Exception as e:
+                        errors[ckey] = f"{type(e).__name__}: {str(e)[:140]}"
+                        print(f"  ref {ckey}: UNREADABLE {type(e).__name__}: {str(e)[:80]}", flush=True)
                         continue
-                    n += 1
-                    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 want and n < int(0.9 * want):
+                    print(f"  ref {ckey}: {n} of {want} rows via datasets -> rows-API fallback",
+                          flush=True)
+                    n2, g2 = ingest(tk, api_rows(repo, cfg, sp, max_items_per_split),
+                                    max_items_per_split)
+                    if n2 > n:
+                        n, g = n2, g2
                 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
+                print(f"  ref {ckey}: {n} items of {want} -> {sum(a.size for a in g):,} grams",
+                      flush=True)
+                if best is None or n > best[1]:
+                    best = (cfg, n, g, want)
+                if n >= MIN_ITEMS_PER_SPLIT and (not want or n >= int(0.9 * want)):
                     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}"
+            if best is None or best[1] < MIN_ITEMS_PER_SPLIT:
+                errors[key] = (f"no config of {repo} yielded >= {MIN_ITEMS_PER_SPLIT} rows for {sp} "
+                               f"(best {(best[1] if best else 0)} of {srv.get(key)})")
                 continue
-            per_task = [a for a in g]
-            if per_task:
-                arr = np.concatenate(per_task)
+            cfg, n, g, want = best
+            if want and n < int(0.9 * want):
+                errors[key] = f"covered {n} of {want} rows in {sp} -- under-read, not calling it clean"
+                continue
+            if g:
+                arr = np.concatenate(g)
                 grams.append(arr)
                 masks.append(np.full(arr.shape, bit, dtype=np.uint16))
     if not grams:
@@ -160,7 +246,8 @@ def build_reference(tk, max_items_per_split):
     allg = np.concatenate(grams)
     allm = np.concatenate(masks)
     order = np.argsort(allg, kind="stable")
-    return allg[order], allm[order], counts, errors
+    return allg[order], allm[order], counts, errors, srv
+
 
 
 
@@ -264,9 +351,10 @@ def main():
     tk_sha = hashlib.sha256(open(tjp, "rb").read()).hexdigest()
 
     want_meta = {"k": K, "max_items": a.max_items_per_split, "tokenizer_sha256": tk_sha,
+                 "text_keys": list(TEXT_KEYS),
                  "refs": [[t[0], t[1], list(t[2]), list(t[3])] for t in REFS]}
     ref = refmask = None
-    refcounts = errors = {}
+    refcounts = errors = srv = {}
     if a.ref_cache and os.path.exists(os.path.join(a.ref_cache, "refmeta.json")):
         try:
             if json.load(open(os.path.join(a.ref_cache, "refmeta.json"))) == want_meta:
@@ -274,6 +362,7 @@ def main():
                 refmask = np.load(os.path.join(a.ref_cache, "refmask.npy"))
                 refcounts = json.load(open(os.path.join(a.ref_cache, "refcounts.json")))
                 errors = json.load(open(os.path.join(a.ref_cache, "errors.json")))
+                srv = json.load(open(os.path.join(a.ref_cache, "srv.json")))
                 print(f"reference: loaded from cache, {ref.size:,} grams", flush=True)
         except Exception as e:
             print(f"reference cache unusable ({type(e).__name__}: {str(e)[:120]}); rebuilding",
@@ -281,7 +370,7 @@ def main():
             ref = None
     if ref is None:
         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)
+        ref, refmask, refcounts, errors, srv = build_reference(tk, a.max_items_per_split)
         if a.ref_cache:
             os.makedirs(a.ref_cache, exist_ok=True)
             np.save(os.path.join(a.ref_cache, "ref.npy"), ref)
@@ -289,6 +378,7 @@ def main():
             json.dump(want_meta, open(os.path.join(a.ref_cache, "refmeta.json"), "w"))
             json.dump(refcounts, open(os.path.join(a.ref_cache, "refcounts.json"), "w"))
             json.dump(errors, open(os.path.join(a.ref_cache, "errors.json"), "w"))
+            json.dump(srv, open(os.path.join(a.ref_cache, "srv.json"), "w"))
             print(f"reference cached under {a.ref_cache}", flush=True)
     tasks_covered = sorted({k.split("::")[0] for k, v in refcounts.items()
                             if v >= MIN_ITEMS_PER_SPLIT})
@@ -387,6 +477,7 @@ def main():
         "k": K,
         "fingerprint": "multiplicative rolling hash mod 2^64, P=0x100000001B3",
         "reference_items_by_split": refcounts,
+        "reference_rows_expected_by_split": srv,
         "reference_loader_problems": errors,
         "reference_grams_distinct": int(np.unique(ref).size),
         "reference_grams_total": int(ref.size),