Upload source/files/in2_gate.py with huggingface_hub

source/files/in2_gate.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""
+in2_gate.py -- the EXACT layout gate for in2_gen.py output. No GPU, no CUDA.
+
+This is the only exact pre-flight check in the IN2 pilot. The step-0 loss check
+was demoted to a diagnostic because its two hypotheses sit only ~0.22 nats apart
+(no-induction 4.147 vs strong-induction 3.925, against a 3.901 plain baseline),
+which cannot validate a layout. These assertions can.
+
+Checks, per arm dir:
+  1  file size is exactly examples*T*2 bytes (headerless uint16)
+  2  every token in [0, VOCAB)
+  3  no answer span intrudes into the last PAD_TAIL offsets of its example
+  4  planted spans (needle / probe+answer / decoy) never overlap
+  5  RETRIEVAL IS ACTUALLY POSSIBLE: the tokens at the probe equal the needle's
+     key, and the tokens at the answer equal the needle's value. If this fails
+     the task is unsolvable and the whole experiment is void.
+  6  batch alignment: re-derive x/y exactly as DataLoaderLite::next_batch does
+     (x=buf[i], y=buf[i+1], stride B*T*dp) and assert every example boundary
+     lands on a row boundary
+  7  arm-specific distance invariants (b: all == span, c: all >= d_min)
+
+Cross-arm:
+  8  arms B and C have BIT-IDENTICAL answer_start histograms -- the discipline
+     the whole B-vs-C contrast rests on
+
+Usage
+  python3 in2_gate.py --dirs /tmp/in2b /tmp/in2c --B 2 --dp 2
+"""
+import argparse, array, json, os, sys
+
+VOCAB = 50304
+fails = []
+
+
+def check(cond, msg):
+    print(("  ok   " if cond else "  FAIL ") + msg)
+    if not cond:
+        fails.append(msg)
+
+
+def load(d, split="train"):
+    js = [f for f in os.listdir(d) if f.endswith(".json") and f"_{split}_" in f]
+    if len(js) != 1:
+        raise SystemExit(f"{d}: expected exactly one {split} manifest, found {js}")
+    man = json.load(open(os.path.join(d, js[0])))
+    toks = array.array("H")
+    for sp in man.get("shards", [man["bin"]]):
+        with open(sp, "rb") as f:
+            toks.fromfile(f, os.path.getsize(sp) // 2)
+    return man, toks
+
+
+def gate_one(d, B, dp, split="train"):
+    man, toks = load(d, split)
+    T, n_ex, arm, span = man["T"], man["examples"], man["arm"], man["span"]
+    kl, vl, pad = man["key_len"], man["val_len"], man["pad_tail"]
+    print(f"\n=== arm {arm}  T={T}  examples={n_ex}  K={man['K']} ===")
+
+    tot = sum(os.path.getsize(sp) for sp in man.get("shards", [man["bin"]]))
+    check(tot == n_ex * T * 2,
+          f"1 total size across {len(man.get('shards',[1]))} shard(s) == examples*T*2 "
+          f"({n_ex*T*2} bytes, got {tot})")
+    lo, hi = min(toks), max(toks)
+    check(lo >= 0 and hi < VOCAB, f"2 tokens in [0,{VOCAB}) (min {lo}, max {hi})")
+
+    bad_pad = bad_ovl = bad_ret = bad_dist = 0
+    for ei, recs in enumerate(man["records"]):
+        base = ei * T
+        spans = []
+        for r in recs:
+            a0, a1 = r["answer_start"], r["answer_start"] + vl
+            if a1 > T - pad:
+                bad_pad += 1
+            spans.append((r["needle_start"], r["needle_start"] + span))
+            spans.append((r["probe_start"], r["probe_start"] + span))
+            # 5: the probe must reproduce the needle's key, the answer its value
+            nk = toks[base + r["needle_start"]: base + r["needle_start"] + kl]
+            nv = toks[base + r["needle_start"] + kl: base + r["needle_start"] + span]
+            pk = toks[base + r["probe_start"]: base + r["probe_start"] + kl]
+            av = toks[base + a0: base + a1]
+            if nk != pk or nv != av:
+                bad_ret += 1
+            if arm == "b" and r["distance"] != span:
+                bad_dist += 1
+            if arm == "c" and r["distance"] < man["d_min"]:
+                bad_dist += 1
+        spans.sort()
+        for (s0, e0), (s1, _) in zip(spans, spans[1:]):
+            if s1 < e0:
+                bad_ovl += 1
+    check(bad_pad == 0, f"3 no answer within the last {pad} offsets ({bad_pad} bad)")
+    check(bad_ovl == 0, f"4 planted spans never overlap ({bad_ovl} bad)")
+    check(bad_ret == 0, f"5 RETRIEVAL POSSIBLE: probe==key and answer==value ({bad_ret} bad)")
+    check(bad_dist == 0, f"7 arm-{arm} distance invariant ({bad_dist} bad)")
+
+    # 6: DataLoaderLite windowing. next_batch reads B*T+1 tokens and sets
+    # y[i]=buf[i+1]; pos_ advances by B*T*world. Example boundaries must land on
+    # row boundaries or an example would straddle two rows.
+    stride = B * T * dp
+    check(stride % T == 0, f"6a batch stride {stride} is a multiple of T={T}")
+    check((n_ex * T) % stride == 0,
+          f"6b total {n_ex*T} tokens is a whole number of batches (stride {stride})")
+    if len(toks) >= B * T + 1:
+        x0 = toks[0:B * T]
+        y0 = toks[1:B * T + 1]
+        rows_ok = all(x0[r * T] == toks[r * T] for r in range(B))
+        check(rows_ok, f"6c row r starts at token r*T for all B={B} rows")
+        check(y0[T - 1] == toks[T], "6d y[T-1] targets the next example's first token "
+                                    "(why PAD_TAIL exists)")
+    return man
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dirs", nargs="+", required=True)
+    ap.add_argument("--B", type=int, default=2)
+    ap.add_argument("--dp", type=int, default=2)
+    ap.add_argument("--split", default="train", choices=["train", "val"])
+    args = ap.parse_args()
+
+    mans = {}
+    for d in args.dirs:
+        m = gate_one(d, args.B, args.dp, args.split)
+        mans[m["arm"]] = m
+
+    if "b" in mans and "c" in mans:
+        print("\n=== cross-arm ===")
+        hb, hc = mans["b"]["answer_start_hist"], mans["c"]["answer_start_hist"]
+        check(hb == hc, "8 arms B and C have BIT-IDENTICAL answer_start histograms "
+                        f"({len(hb)} vs {len(hc)} distinct offsets)")
+        check(mans["b"]["seed"] == mans["c"]["seed"], "8b same seed")
+        check(mans["b"]["answer_tokens"] == mans["c"]["answer_tokens"],
+              "8c same answer-token count")
+
+    print(f"\n{'ALL PASS' if not fails else str(len(fails)) + ' FAILURE(S)'}")
+    return 1 if fails else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())