Upload build/build_mix.py with huggingface_hub

build/build_mix.py

@@ -375,23 +375,59 @@ def stage(args, st, tokenizer):
 
 # --------------------------------------------------------------------- merge
 def merge(args, st, tokenizer):
-    """Interleave staged documents proportionally to each source's staged token total."""
+    """Interleave staged documents proportionally to each source's staged token total.
+
+    --exclude-dir applies per-source uint8 masks written by audit_contamination --write-filter: a 1 byte
+    means "this document's 13-grams appear in a benchmark train/validation/dev split, do not emit it"
+    (§3.3). Masks are indexed by document ordinal within the source, which is stable because both the
+    audit and the merge replay the source's own record.json shard order.
+    """
     vocab = tokenizer.get_vocab_size()
+    masks, dropped = {}, {}
+    if args.exclude_dir:
+        fj = os.path.join(args.exclude_dir, "filter.json")
+        if not os.path.exists(fj):
+            print(f"FATAL: --exclude-dir {args.exclude_dir} contains no filter.json", flush=True)
+            sys.exit(5)
+        dropped = json.load(open(fj)).get("sources", {})
+        for key in sorted(os.listdir(args.exclude_dir)):
+            if key.endswith(".u8"):
+                with open(os.path.join(args.exclude_dir, key), "rb") as fh:
+                    masks[key[:-3]] = fh.read()
+        print(f"exclude: {len(masks)} mask(s), "
+              f"{sum(len(m) for m in masks.values()):,} docs covered, "
+              f"{sum(int(v.get('dropped_docs', 0)) for v in dropped.values()):,} flagged", flush=True)
     streams = []
     for src in SOURCES:
         key = src_key(src)
         rec = st.d["staged"].get(key)
         if not rec or not rec.get("shards"):
             continue
+        staged_docs = sum(int(s.get("docs", 0)) for s in rec["shards"])
+        m = masks.get(key)
+        if m is not None and len(m) != staged_docs:
+            # A short mask would silently keep exactly the documents the audit did not reach.
+            print(f"FATAL: mask for {key} covers {len(m)} docs but {staged_docs} are staged", flush=True)
+            sys.exit(5)
         weight = rec.get("tokens") or sum(s["tokens"] for s in rec["shards"])
+        weight = weight - int((dropped.get(key) or {}).get("dropped_tokens", 0))
         if weight <= 0:
+            print(f"  {key}: fully excluded, dropped from the merge", flush=True)
             continue
         iters = [iter_shard_docs(os.path.join(st.root, s["file"])) for s in rec["shards"]]
 
-        def gen(its=iters):
+        def gen(its=iters, m=m):
+            if m is None:
+                for it in its:
+                    for d in it:
+                        yield d
+                return
+            i = 0
             for it in its:
                 for d in it:
-                    yield d
+                    if not m[i]:
+                        yield d
+                    i += 1
 
         streams.append({"key": key, "weight": float(weight), "credit": 0.0, "it": gen(iters),
                         "pending": None, "emitted_tokens": 0, "exhausted": False})
@@ -520,6 +556,9 @@ def main():
                          "empty string disables sync")
     ap.add_argument("--stage-only", action="store_true")
     ap.add_argument("--merge-only", action="store_true")
+    ap.add_argument("--exclude-dir", default="",
+                    help="directory of per-source <key>.u8 contamination masks + filter.json written by "
+                         "audit_contamination --write-filter; documents flagged there are not merged (§3.3)")
     args = ap.parse_args()
 
     if args.only:
@@ -608,6 +647,22 @@ def main():
                                               default=0),
         "seconds": round(time.monotonic() - t0, 1),
     }
+    if args.exclude_dir:
+        # Provenance: a filtered mix is only reproducible if the exclusion set travels with it, so the
+        # counts go in the manifest and the mask bytes go in mixroot/filter/ (publish_mix ships them).
+        try:
+            fj = json.load(open(os.path.join(args.exclude_dir, "filter.json")))
+            manifest["contamination_filter"] = {
+                "criterion": fj.get("criterion"), "k": fj.get("k"),
+                "docs_dropped": sum(int(v.get("dropped_docs", 0)) for v in fj["sources"].values()),
+                "tokens_dropped": sum(int(v.get("dropped_tokens", 0) or 0)
+                                      for v in fj["sources"].values()),
+                "per_source": {k2: {"docs": v2.get("dropped_docs"),
+                                    "tokens": v2.get("dropped_tokens")}
+                               for k2, v2 in fj["sources"].items()}}
+        except Exception as e:
+            print(f"FATAL: --exclude-dir set but filter.json unreadable: {e}", flush=True)
+            sys.exit(5)
     with open(os.path.join(args.root, "manifest.json"), "w") as f:
         json.dump(manifest, f, indent=1)
     # Nothing may be printed between the JSON and its END marker: the caller parses this span with