Upload build/build_mix.py with huggingface_hub

build/build_mix.py

@@ -392,13 +392,20 @@ def merge(args, st, tokenizer):
     audit and the merge replay the source's own record.json shard order.
     """
     vocab = tokenizer.get_vocab_size()
-    masks, dropped = {}, {}
+    masks, dropped, filter_digest = {}, {}, None
     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", {})
+        filter_digest = hashlib.sha256(open(fj, "rb").read()).hexdigest()
+        fj_obj = json.load(open(fj))
+        if not fj_obj.get("scanned_sources"):
+            print("FATAL: filter.json records no scanned_sources -- the audit that wrote it predates "
+                  "the completeness field, so an unscanned source could be silently unfiltered. Re-run "
+                  "audit_contamination --write-filter.", flush=True)
+            sys.exit(5)
+        dropped = fj_obj.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:
@@ -446,6 +453,42 @@ def merge(args, st, tokenizer):
         # silent success in this pipeline, so it is a failure and not a return.
         print("FATAL merge: nothing staged. Refusing to emit an empty mix.", flush=True)
         sys.exit(4)
+    if args.exclude_dir:
+        # A missing mask must never mean "this source is clean" (review finding): it means the audit
+        # never scanned it, and the default behaviour would have been to merge every contaminated doc.
+        unscanned = [s["key"] for s in streams if s["key"] not in masks]
+        if unscanned:
+            print(f"FATAL merge: --exclude-dir has no mask for {unscanned}. Refusing to merge sources "
+                  "the contamination audit did not read.", flush=True)
+            sys.exit(5)
+
+    # A re-merge is only correct if it starts from nothing. The loop below stops when merged_tokens()
+    # reaches the input total, so on a mixroot that already holds 141 shards it would emit ZERO new
+    # shards, exit 0, and let the manifest claim a filter that never happened (found in review while
+    # Gate 2c was running against exactly this path). Resume-after-a-kill stays legal -- it is the same
+    # inputs -- so the test is on a digest of the merge's inputs, not on whether shards exist.
+    digest = hashlib.sha256(json.dumps({
+        "streams": [[s["key"], round(s["weight"]) ] for s in streams], "seed": SHUFFLE_SEED,
+        "tokens_per_shard": TOKENS_PER_SHARD, "val_stride": VAL_STRIDE,
+        "filter": filter_digest}, sort_keys=True).encode()).hexdigest()
+    prior = st.d.get("merge_inputs")
+    if prior is not None and prior != digest and st.d["merged_shards"]:
+        if not (args.rebuild_merged or args.exclude_dir):
+            print("FATAL merge: the mix root already holds shards merged from DIFFERENT inputs "
+                  f"(recorded {prior[:12]}, now {digest[:12]}). Continuing would be a silent no-op. "
+                  "Pass --rebuild-merged to discard them deliberately.", flush=True)
+            sys.exit(6)
+        n_old = len(st.d["merged_shards"]) + len(st.d.get("val_shards", []))
+        for grp in ("merged_shards", "val_shards"):
+            for rec in st.d.get(grp, []):
+                p = os.path.join(st.root, rec["file"])
+                if os.path.exists(p):
+                    os.remove(p)
+        st.d["merged_shards"] = []
+        st.d["val_shards"] = []
+        print(f"merge: inputs changed ({prior[:12]} -> {digest[:12]}); discarded {n_old} old shard(s) "
+              "and rebuilding from the staged sources", flush=True)
+    st.d["merge_inputs"] = digest
     total_w = sum(s["weight"] for s in streams)
     print(f"merge: {len(streams)} sources, "
           f"{sum(s['weight'] for s in streams) / 1e6:,.0f}M staged tokens", flush=True)
@@ -568,6 +611,9 @@ def main():
     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)")
+    ap.add_argument("--rebuild-merged", action="store_true",
+                    help="discard shards already recorded in build_state.json before merging. Implied by "
+                         "--exclude-dir, because a filter changes every shard's contents")
     args = ap.parse_args()
 
     if args.only: