Upload build/audit_contamination.py with huggingface_hub

build/audit_contamination.py

@@ -248,15 +248,48 @@ def main():
                          "build_mix --exclude-dir)")
     ap.add_argument("--out", default="", help="write the report JSON here (publish_mix expects "
                                              "mixroot/audit.json)")
+    ap.add_argument("--ref-cache", default="",
+                    help="directory to keep the built reference gram table in. Streaming the eight "
+                         "reference sets is the slow half of this script, and a re-audit after a filter "
+                         "pass rebuilds exactly the same table -- so it is reused only if the task list, "
+                         "item cap, k and tokenizer sha all still match.")
     a = ap.parse_args()
 
     from tokenizers import Tokenizer
     from huggingface_hub import hf_hub_download
-    tk = Tokenizer.from_file(hf_hub_download("HuggingFaceTB/SmolLM2-135M", "tokenizer.json"))
+    tjp = hf_hub_download("HuggingFaceTB/SmolLM2-135M", "tokenizer.json")
+    tk = Tokenizer.from_file(tjp)
     tk.no_truncation()
-
-    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)
+    import hashlib
+    tk_sha = hashlib.sha256(open(tjp, "rb").read()).hexdigest()
+
+    want_meta = {"k": K, "max_items": a.max_items_per_split, "tokenizer_sha256": tk_sha,
+                 "refs": [[t[0], t[1], list(t[2]), list(t[3])] for t in REFS]}
+    ref = refmask = None
+    refcounts = errors = {}
+    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:
+                ref = np.load(os.path.join(a.ref_cache, "ref.npy"))
+                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")))
+                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",
+                  flush=True)
+            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)
+        if a.ref_cache:
+            os.makedirs(a.ref_cache, exist_ok=True)
+            np.save(os.path.join(a.ref_cache, "ref.npy"), ref)
+            np.save(os.path.join(a.ref_cache, "refmask.npy"), refmask)
+            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"))
+            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})
     print(f"reference: {ref.size:,} grams, {len(tasks_covered)}/8 tasks covered, "