"""Pre-judge gates for FIX 5 (mask schedule) and FIX 7 (PAD attention mask), runs/CSEL/FIX_TRAINING_PREREG.md. Compares one fix arm with the shared control runs/CANVAS_FIXU_CTL (job 152: HEAD trainer, WangchanBERTa body, fixed 12-epoch OneCycle, last epoch) on: 1. LOSS-CURVE gate (pre-registered). fix5 ("trajectory"): the average valid loss over mask ratios {1.0, 0.9, ..., 0.1, 0.05} must drop >= 0.03 nats vs CTL AND the standard valid loss must be no worse than CTL + 0.03. fix7 ("standard"): the trainer's standard masking must drop >= 0.03 nats. Both conditions are computed here on the trainer's 1,000 validation poems (same split: seed 0 shuffle of load_poems), identical masks for both models (seed 1 for standard). 2. raw candidates on the dev49 pilot (the shared CTL decodes runs/DEV49_CTL vs the arm's), paired by seed: form >= 98% and not below CTL by more than 1pp; line-final syllable entropy within 0.3 bits; distinct-2 over syllables not below CTL by more than 0.03; immediate syllable repeats (stutters) per poem not above CTL + 0.05; line-final word (newmm) share unseen in the train poems not above CTL + 0.02, and mean line-final word log-frequency within 0.3 of CTL (strange or bland). Writes RUN/gates_ARM.json; exits 2 if any gate fails (the arm is then not judged). Usage: fixtrain_gates.py RUN_DIR CTL_MODEL ARM_MODEL ARM_NAME {trajectory|standard} CTL_CANDS ARM_CANDS [SEEDS A-B] Amendment 2026-09-29 (fixes re-run on the bug-fixed base, CTL_MODEL = runs/CANVAS_FIXBASE): the validation poems are loaded exactly as the control model's trainer loaded them (segmenter, title cleaning and ๆ handling read from its vocab.json; needs 170's patched canvas_bot.py for such models), both models must record the same data flags, and SEEDS restricts the candidate files to c{A..B}.jsonl (170's control directory holds seeds 0-63). """ import os, sys, json, glob, math, random, collections sys.path.insert(0, "."); sys.path.insert(0, "scripts") import torch, torch.nn.functional as F import canvas_bot as B from syl_align import syls_aligned as SA RATIOS = (1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0.05) def valid_losses(model_dir, VY_cache={}): """standard: the trainer's validation masking (seed 1); r=X: every poem cell masked with probability X (seed 2), the trajectory an iterative decoder visits; identical masks for every model.""" model, V = B.load_model(model_dir) meta = json.load(open(f"{model_dir}/vocab.json", encoding="utf-8")) kw = {"seg": meta.get("seg", "aligned")} kw.update({k: True for k in ("clean_titles", "yamok") if meta.get(k)}) # the fixed base's data flags if "VY" not in VY_cache: random.seed(0); torch.manual_seed(0) # the trainer's split: seed 0 shuffle poems = B.load_poems("data/modern_bot_train.jsonl", **kw); random.shuffle(poems) VY_cache["VY"] = B.build_tok(poems[:1000], V); VY_cache["itos"] = V.itos; VY_cache["kw"] = kw assert V.itos == VY_cache["itos"], "control and arm use different vocabularies" assert kw == VY_cache["kw"], f"control and arm load their data differently: {VY_cache['kw']} vs {kw}" VY = VY_cache["VY"]; seg = model.seg_id.cpu(); poem = (seg[None, :] == 1) & (VY != B.PAD) conds = [("standard",) + B.mask_batch(VY, seg, torch.Generator().manual_seed(1))] g = torch.Generator().manual_seed(2); u = torch.rand(VY.shape, generator=g) for r in RATIOS: m = poem & (u < r) if r == 1.0: m = poem.clone() conds.append((f"r={r}", torch.where(m, torch.full_like(VY, B.MASK), VY), m)) out = {} with torch.no_grad(): for name, x, m in conds: tot = cnt = 0 for i in range(0, len(VY), 50): mi = m[i:i + 50] if not mi.any(): continue lg = model(x[i:i + 50].to(B.DEV)).float().cpu() tot += float(F.cross_entropy(lg[mi], VY[i:i + 50][mi], reduction="sum")); cnt += int(mi.sum()) out[name] = tot / max(cnt, 1) out["trajectory"] = sum(out[f"r={r}"] for r in RATIOS) / len(RATIOS) del model if str(B.DEV) == "mps": torch.mps.empty_cache() return out WF = None def word_logfreq(): """log frequency of words in the TRAIN poems (newmm), for the line-final strange-word check""" global WF if WF is None: from pythainlp.tokenize import word_tokenize c = collections.Counter() for l in open("data/modern_bot_train.jsonl", encoding="utf-8"): for ln in json.loads(l)["poem"].split("\n"): c.update(w for w in word_tokenize(ln.replace(" ", ""), engine="newmm") if w.strip()) tot = sum(c.values()); WF = (c, tot) return WF def stats(rows): from pythainlp.tokenize import word_tokenize c_, tot_ = word_logfreq() finals = collections.Counter(); bi = set(); nbi = 0; stut = 0; lf = [] for r in rows: for ln in r["poem"].split("\n"): sy = SA(ln, 12) if sy: finals[sy[-1]] += 1 pairs = list(zip(sy, sy[1:])); nbi += len(pairs); bi.update(pairs); stut += sum(a == b for a, b in pairs) ws = [w for w in word_tokenize(ln, engine="newmm") if w.strip()] if ws: lf.append(math.log((c_[ws[-1]] + 1) / (tot_ + len(c_)))) tot = sum(finals.values()) return {"n": len(rows), "form": 100 * sum(bool(r["all_pass"]) for r in rows) / len(rows), "final_bits": -sum(c / tot * math.log2(c / tot) for c in finals.values()), "distinct2": len(bi) / max(nbi, 1), "stutters": stut / len(rows), "final_word_logfreq": sum(lf) / max(len(lf), 1), "final_word_unseen": sum(x < math.log(1.5 / (tot_ + len(c_))) for x in lf) / max(len(lf), 1)} def seed_files(d, spec=None): fs = sorted(glob.glob(f"{d}/c*.jsonl")) if spec: a, b = map(int, spec.split("-")) fs = [f"{d}/c{s}.jsonl" for s in range(a, b + 1)] assert all(os.path.exists(f) for f in fs), f"{d}: missing seed files in {spec}" return fs def main(run, ctl_model, arm_model, arm, cond, ctl_cands, arm_cands, seeds=None): assert cond in ("trajectory", "standard") vc, va = valid_losses(ctl_model), valid_losses(arm_model) S = {} for a, d in (("CTL", ctl_cands), (arm, arm_cands)): rows = [json.loads(l) for f in seed_files(d, seeds) for l in open(f, encoding="utf-8")] assert rows, f"no candidates in {d}" S[a] = stats(rows) c, f = S["CTL"], S[arm] dead = [] if cond == "trajectory": if va["trajectory"] > vc["trajectory"] - 0.03: dead.append("loss-curve (trajectory-average valid loss not >= 0.03 below CTL)") if va["standard"] > vc["standard"] + 0.03: dead.append("loss-curve (standard valid loss worse than CTL + 0.03)") elif va[cond] > vc[cond] - 0.03: dead.append(f"loss-curve ({cond} valid loss not >= 0.03 below CTL)") if f["form"] < c["form"] - 1: dead.append("form (below CTL - 1pp)") if abs(f["final_bits"] - c["final_bits"]) > 0.3: dead.append("rhyme variety") if f["distinct2"] < c["distinct2"] - 0.03: dead.append("distinct-2") if f["stutters"] > c["stutters"] + 0.05: dead.append("stutters") # line-final strange-word check (both directions): rarer = strange words at line ends; more frequent = blander if f["final_word_unseen"] > c["final_word_unseen"] + 0.02: dead.append("line-final unseen words") if abs(f["final_word_logfreq"] - c["final_word_logfreq"]) > 0.3: dead.append("line-final word frequency (strange or bland)") print(f" valid loss CTL standard {vc['standard']:.3f} trajectory {vc['trajectory']:.3f} | {arm} standard {va['standard']:.3f} trajectory {va['trajectory']:.3f} (gate: {cond})") for a, s in S.items(): print(f" {a}: n {s['n']}, form {s['form']:.1f}%, final-syllable bits {s['final_bits']:.2f}, distinct-2 {s['distinct2']:.3f}, stutters/poem {s['stutters']:.3f}, line-final word log-freq {s['final_word_logfreq']:.2f}, unseen {s['final_word_unseen']:.3f}") json.dump({"valid": {"CTL": vc, arm: va}, "arms": S, "dead": dead}, open(f"{run}/gates_{arm}.json", "w"), indent=1) print(f" {arm} KILLED before judging: " + ", ".join(dead) if dead else f" {arm} pre-judge gates PASS") if dead: sys.exit(2) if __name__ == "__main__": main(*sys.argv[1:9])