src/scripts/fixon_raw.py

"""RAW-text judge inputs, scoring and a channel check for the fix arms on the bug-fixed base
(runs/CSEL/FIX_TRAINING_PREREG.md amendments 3-4): job 184 (FIX 21-23, prefix jr, pools FXR) and job 185 (FIX 10/5/7
of jobs 173/175/176, prefix jq, pools FXQ; replaces job 177's ๆ-written pools).

The judge sees the poem exactly as the model wrote it: NO ๆ write-back (runs/CSEL/FIXED_BASE_PREREG.md amendment 1 -
fixbase_pipe.judged_text contracts any two equal adjacent syllables, which turned model stutters into plausible
reduplication for 6.6% of candidates). scripts/fixon.py (rawjudge/rawpools) still writes ๆ back and is not used here.

  rawpools OUT A-B PFX ARM=DIR...  judge inputs, seeds A..B: the control (170's FIXED_b0 candidates) + <= 3 arms; per
                                seed and topic half (fixbase_pipe.halves, 25/24 topics) OUT/{PFX}_{fixbase,ARM...}_s{seed}h{h}.jsonl
                                = ALL candidates (form-failing ones too, so every arm has the same count per pool and
                                build_pool's balance assert holds by construction; form enters the endpoint through
                                the candidate's all_pass), id TOPIC_sSEED, poem = the raw candidate text (asserted 4 lines)
  score OUT A-B PREFIX ARM...   after judging runs/{PREFIX}_{CLEAN,ALIGN}{seed}{half}/run1.json: per arm vs the control,
                                paired by (topic, seed) over ALL candidates (every candidate is judged; the endpoint
                                counts a form-failing candidate as not clean / not aligned; judged-only rates reported): clean, aligned, all-three, topic-clustered 95% CI, decision
                                (pre-registered rule); a pool whose human anchors are < 70% clean is void -> exit 3
  phonuse CTL_MODEL ARM_MODEL   row norms of token / rime / tone embeddings and the standard valid loss with rime+tone
                                zeroed, for both models (reported by job 179; not a gate)
  test                          CPU tests of the pure parts
"""
import os, sys, json, collections
sys.path.insert(0, "."); sys.path.insert(0, "scripts")

FB = "runs/FIXBASE"


def rows(p): return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()]


def wrows(p, R):
    with open(p + ".tmp", "w", encoding="utf-8") as f:
        for r in R: f.write(json.dumps(r, ensure_ascii=False) + "\n")
    os.replace(p + ".tmp", p)


def halves(ids):   # = fixbase_pipe.halves (copied so the pools do not depend on fixbase_pipe)
    ids = sorted(ids); h = (len(ids) + 1) // 2
    return [ids[:h], ids[h:]]


def arms_of(specs):
    A = [("fixbase", f"{FB}/FIXED_b0/cands")] + [tuple(x.split("=", 1)) for x in specs]
    assert len({n for n, _ in A}) == len(A) and len(A) <= 4, "arm names must be unique; at most control + 3 arms per pool"
    return A


def rawpools(out, spec, pfx, *arms):
    assert pfx in ("jr", "jq"), f"prefix {pfx}: jr (job 184) or jq (job 185)"
    a, b = map(int, spec.split("-"))
    A = arms_of(arms)
    os.makedirs(out, exist_ok=True); n = nf = 0
    for s in range(a, b + 1):
        R = {name: {r["id"]: r for r in rows(f"{d}/c{s}.jsonl")} for name, d in A}
        ids = sorted(R["fixbase"])
        assert all(set(R[k]) == set(ids) for k in R), f"seed {s}: arms cover different topics"
        assert all(R[k][i]["prompt"] == R["fixbase"][i]["prompt"] for k in R for i in ids), f"seed {s}: prompts differ"
        for h, part in enumerate(halves(ids)):
            for name, _ in A:
                J = [{"id": f"{i}_s{s}", "prompt": R[name][i]["prompt"], "poem": R[name][i]["poem"]} for i in part]
                bad = [r["id"] for r in J if len(r["poem"].split("\n")) != 4]
                assert not bad, f"{name} seed {s}: poems without exactly 4 lines (build_pool would drop them silently): {bad}"
                wrows(f"{out}/{pfx}_{name}_s{s}h{h}.jsonl", J); n += len(J); nf += sum(not R[name][i]["all_pass"] for i in part)
    print(f"  raw pools (no ๆ write-back): arms {[x for x, _ in A]}, seeds {a}-{b}, {n} poems ({nf} form-failing, judged too) in {2 * (b - a + 1)} pool pairs")


def decide(clean, aligned):
    """Pre-registered (FIX_TRAINING_PREREG.md): CONTINUE if clean >= +5pp with the lower 95% bound > 0 and aligned
    >= -3pp; KILL if clean < +2pp; otherwise GREY."""
    if clean["diff"] >= 5 and clean["ci"][0] > 0 and aligned["diff"] >= -3: return "CONTINUE"
    if clean["diff"] < 2: return "KILL"
    return "GREY"


def labels(prefix, a, b, root_pools="runs"):
    """{(src, topic, seed): {clean, aligned}} for judged system poems; anchor clean share per clean pool"""
    lab = {"clean": {}, "aligned": {}}; anchors = {}
    for s in range(a, b + 1):
        for h in (0, 1):
            for kind, inst in (("clean", "CLEAN"), ("aligned", "ALIGN")):
                d = f"{root_pools}/{prefix}_{inst}{s}{h}"
                km = json.load(open(f"{d}/keymap.json", encoding="utf-8"))
                res = {r["item"]: r for r in json.load(open(f"{d}/run1.json", encoding="utf-8"))}
                assert set(res) == set(km), f"{d}: judge output does not cover the pool"
                hum = []
                for k, v in km.items():
                    r = res[k]
                    val = (not any(e.get("severity") in ("critical", "major") for e in r["errors"])) if kind == "clean" else int(r["align"]) >= 4
                    if v["src"] == "human": hum.append(val); continue
                    t, ss = v["id"].rsplit("_s", 1); assert int(ss) == s, f"{d}: {v} is not seed {s}"
                    assert (v["src"], t, s) not in lab[kind], f"{v} judged twice"
                    lab[kind][(v["src"], t, s)] = val
                if kind == "clean": anchors[d] = sum(hum) / max(len(hum), 1)
    return lab, anchors


def score(out, spec, prefix, *arm_specs, root_pools="runs"):
    import fixbase_pipe
    a, b = map(int, spec.split("-"))
    A = arms_of(arm_specs)
    form = {}
    for name, d in A:
        for s in range(a, b + 1):
            for r in rows(f"{d}/c{s}.jsonl"): form[(name, r["id"], s)] = bool(r["all_pass"])
    lab, anchors = labels(prefix, a, b, root_pools)
    void = {d: v for d, v in anchors.items() if v < 0.70}
    for kind in ("clean", "aligned"):
        assert set(lab[kind]) == set(form), f"{kind}: judged poems differ from the candidates"
    full = {kind: {k: form[k] and lab[kind][k] for k in form} for kind in ("clean", "aligned")}
    full["all-three"] = {k: full["clean"][k] and full["aligned"][k] for k in form}
    full["form"] = dict(form); full["judged-clean"] = dict(lab["clean"]); full["judged-aligned"] = dict(lab["aligned"])
    R = {"anchor_clean_share": anchors, "void_pools": sorted(void)}
    for name, _ in A[1:]:
        R[name] = {}
        for kind, L in full.items():
            c = {k[1:]: v for k, v in L.items() if k[0] == "fixbase"}; t = {k[1:]: v for k, v in L.items() if k[0] == name}
            st = R[name][kind] = fixbase_pipe.paired_stats(c, t)
            print(f"  {name:8s} {kind:9s} control {st['a']:5.1f}%  arm {st['b']:5.1f}%  diff {st['diff']:+.1f}pp "
                  f"({st['up']} up / {st['down']} down, n={st['n']})  95% CI [{st['ci'][0]:+.1f}, {st['ci'][1]:+.1f}]")
        R[name]["decision"] = decide(R[name]["clean"], R[name]["aligned"])
        print(f"  {name}: DECISION (pre-registered) {R[name]['decision']}")
    os.makedirs(out, exist_ok=True)
    json.dump(R, open(f"{out}/RESULT_{prefix}.json", "w"), indent=1)
    if void:
        print(f"  VOID pools (anchors < 70% clean, judge_frozen §2.2): {void} - re-judge those pools before any decision"); sys.exit(3)


def phonuse(ctl, arm):
    import torch
    import canvas_bot as B, fixtrain_gates as G
    G.RATIOS = (1.0,)          # only the standard condition is reported (and r=1.0, which valid_losses needs)
    real = B.load_model
    def zeroed(path):
        m, V = real(path)
        with torch.no_grad(): m.rime.weight.zero_(); m.tone.weight.zero_()
        return m, V
    for d in (ctl, arm):
        m, _ = real(d); n = lambda w: float(w.detach().float().norm(dim=1).mean())
        print(f"  {d}: row norms token {n(m.tok.weight):.3f}, rime {n(m.rime.weight):.3f}, tone {n(m.tone.weight):.3f}")
        del m
        full = G.valid_losses(d)["standard"]
        B.load_model = zeroed
        try: z = G.valid_losses(d)["standard"]
        finally: B.load_model = real
        print(f"  {d}: standard valid loss {full:.4f}; rime+tone zeroed {z:.4f} (use = {z - full:+.4f} nats)")


def _test():
    import tempfile
    global FB
    t = tempfile.mkdtemp(); fb0, FB = FB, t
    os.makedirs(f"{t}/FIXED_b0/cands"); os.makedirs(f"{t}/x")
    stut = "ห้องห้องนี้\nฝนตก\nลมพัด\nใจเย็น"   # a model stutter: must reach the judge as written (no ๆ)
    for d, bad in ((f"{t}/FIXED_b0/cands", 3), (f"{t}/x", 1)):
        wrows(f"{d}/c0.jsonl", [{"id": f"T{i:02d}", "prompt": "ฝน", "poem": stut if i == 0 else "ฝนตก\nลม\nพัด\nมา", "all_pass": i != bad} for i in range(5)])
    rawpools(f"{t}/p", "0-0", "jr", f"fix21={t}/x")
    h0c, h0a, h1a = rows(f"{t}/p/jr_fixbase_s0h0.jsonl"), rows(f"{t}/p/jr_fix21_s0h0.jsonl"), rows(f"{t}/p/jr_fix21_s0h1.jsonl")
    # ALL candidates (form-failing ones too): equal counts per arm and pool by construction
    assert [r["id"] for r in h0c] == [r["id"] for r in h0a] == ["T00_s0", "T01_s0", "T02_s0"], (h0c, h0a)
    assert [r["id"] for r in h1a] == ["T03_s0", "T04_s0"] and h0c[0]["poem"] == stut and "ๆ" not in h0a[0]["poem"]
    # a candidate that is not 4 lines stops the export (build_pool would drop it silently)
    os.makedirs(f"{t}/y"); wrows(f"{t}/y/c0.jsonl", [dict(r, poem="ฝน\nตก\nลม") if r["id"] == "T02" else r for r in rows(f"{t}/x/c0.jsonl")])
    try: rawpools(f"{t}/q", "0-0", "jq", f"fix5={t}/y"); raise AssertionError("3-line poem not caught")
    except AssertionError as e: assert "4 lines" in str(e), e
    # scoring: fake judged pools; fix21 clean on everything it passed, control fails T00
    for inst in ("CLEAN", "ALIGN"):
        for h in (0, 1):
            d = f"{t}/runs/FXT_{inst}0{h}"; os.makedirs(d)
            km, res = {}, []
            for name in ("fixbase", "fix21"):
                for r in rows(f"{t}/p/jr_{name}_s0h{h}.jsonl"):
                    k = f"{name}{r['id']}"; km[k] = {"src": name, "id": r["id"]}
                    bad = name == "fixbase" and r["id"] == "T00_s0"
                    res.append({"item": k, "errors": [{"severity": "major"}] if bad else [], "align": 2 if bad else 5})
            for j in range(4):
                km[f"h{j}"] = {"src": "human", "id": f"H{j}"}
                res.append({"item": f"h{j}", "errors": [{"severity": "major"}] if (h == 1 and inst == "CLEAN" and j < 2) else [], "align": 5})
            json.dump(km, open(f"{d}/keymap.json", "w")); json.dump(res, open(f"{d}/run1.json", "w"))
    try:
        score(f"{t}/o", "0-0", "FXT", f"fix21={t}/x", root_pools=f"{t}/runs"); raise AssertionError("void pool not caught")
    except SystemExit as e:
        assert e.code == 3
    R = json.load(open(f"{t}/o/RESULT_FXT.json"))
    assert R["void_pools"] == [f"{t}/runs/FXT_CLEAN01"], R["void_pools"]
    c = R["fix21"]["clean"]   # control clean: T01,T02,T04 (T00 judged bad, T03 form-fail); arm: T00,T02,T03,T04 (T01 form-fail)
    assert R["fix21"]["judged-clean"]["a"] == 80 and R["fix21"]["judged-clean"]["b"] == 100 and R["fix21"]["form"]["b"] == 80
    assert c["n"] == 5 and abs(c["a"] - 60) < 1e-9 and abs(c["b"] - 80) < 1e-9 and c["up"] == 2 and c["down"] == 1, c
    assert decide({"diff": 6, "ci": (1, 11)}, {"diff": -2}) == "CONTINUE" and decide({"diff": 6, "ci": (1, 11)}, {"diff": -4}) == "GREY"
    assert decide({"diff": 1.9, "ci": (-3, 7)}, {"diff": 0}) == "KILL" and decide({"diff": 3, "ci": (-3, 9)}, {"diff": 0}) == "GREY"
    FB = fb0
    print("fixon_raw: tests pass")


if __name__ == "__main__":
    cmd, args = sys.argv[1], sys.argv[2:]
    {"rawpools": rawpools, "score": score, "phonuse": phonuse, "test": _test}[cmd](*args)