src/scripts/vetomore.py

"""Experiment 1 (director, 2026-09-28): more evidence for the collocation veto, on extra development topics.

Same pipeline as 129 (runs/CSEL/VETOREF_PREREG.md) on topics that never touched test150: the user's everyday
benchmark topics (data/prompts/everyday_bench.jsonl, U001..U222) minus corpus titles (train or test split),
minus every evaluation set, normalised to Thai letters. Arms, per topic:
  T  = Typhoon pick (rank 0 of the 095 top-8 ranking)          R1 = T + round-1 located refill
  V  = veto pick (122: bottom half by WC line PLL vetoed)       VR = V + round-1 located refill
Refill = done/105 verbatim: weakest line = argmin copy-masked canvas line mean, K=16 re-draws with seeds
sha256(id|line|k), best topic-keeping form-passing variant replaces the line only if >= 0.5 higher.
Where V == T the VR poem is the R1 poem by construction (same id, same seeds); only V != T picks are refilled
a second time.

Subcommands (D = run dir, default runs/VETOMORE):
  topics  write data/prompts/vetomore_topics.jsonl (the disjointness asserts live here)
  picks   D/typhoon_picks.jsonl from D/typhoon_top8.jsonl (rank 0, as runs/CSEL/typhoon_picks.jsonl)
  arms    after colloc_veto.py: D/pick_T.jsonl (all topics) and D/pick_V.jsonl (V != T only)
  locmap  D/{T,V}_scores.jsonl -> D/{T,V}_line.json (argmin canvas line mean)
  select  accept rule for both arms -> D/r1.jsonl, D/vr.jsonl; changed = VR != R1 (form-driven and unrefilled
          excluded, reported); drift gates on changed VR vs R1 (form >= 98%, stutters <= +0.10, cross-line
          repeats <= +0.50); writes D/changed_ids.json and D/PASS only if the gates pass (exit 2 = killed,
          exit 3 = nothing changed)
  export  judge inputs on the changed topics: D/judge_{round1,vetoref}.jsonl (id/prompt/poem only), plus the
          6 test30_pilot topics whose VR != R1 in runs/VETOREF (129 excluded them as tuning topics):
          D/judge_pilot_{round1,vetoref}.jsonl
"""
import os, re, sys, json, glob, collections
sys.path.insert(0, "."); sys.path.insert(0, "scripts")

BENCH = "data/prompts/everyday_bench.jsonl"
OUT_TOPICS = "data/prompts/vetomore_topics.jsonl"
EVAL_SETS = ["fresh120", "test150", "test30_pilot", "dev30_pilot", "test_topics", "seen_topics", "ul_train_titles",
             "test", "train", "dev"]


def rows(p): return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()]
def norm(t): return re.sub(r"[^ก-๙]", "", t or "")


def topics(d=None):
    corpus = set()
    for f in ("data/modern_bot_train.jsonl", "data/modern_bot_test.jsonl"):
        for r in rows(f):
            for k in ("topic", "title"):
                if r.get(k): corpus.add(norm(r[k]))
    ev = set()
    for s in EVAL_SETS:
        p = f"data/prompts/{s}.jsonl"
        if os.path.exists(p): ev |= {norm(r["text"]) for r in rows(p)}
    U = rows(BENCH)
    keep = [u for u in U if norm(u["text"]) and norm(u["text"]) not in corpus and norm(u["text"]) not in ev]
    keys = [norm(u["text"]) for u in keep]
    assert len(set(keys)) == len(keys), "duplicate normalised topics"
    # asserts (tests): no overlap with corpus titles or any evaluation set, fresh120 in particular
    fresh = {norm(r["text"]) for r in rows("data/prompts/fresh120.jsonl")}
    assert not (set(keys) & corpus) and not (set(keys) & ev) and not (set(keys) & fresh)
    with open(OUT_TOPICS, "w", encoding="utf-8") as f:
        for u in keep: f.write(json.dumps({"id": u["id"], "form": "topic", "text": u["text"], "set": "VETOMORE",
                                           "group": u.get("group")}, ensure_ascii=False) + "\n")
    print(f"  {len(U)} benchmark topics; corpus titles {sum(norm(u['text']) in corpus for u in U)}, "
          f"evaluation sets {sum(norm(u['text']) in ev for u in U)} -> {len(keep)} extra development topics")


def picks(d):
    top = [r for r in rows(f"{d}/typhoon_top8.jsonl") if r["sel_rank"] == 0]
    with open(f"{d}/typhoon_picks.jsonl", "w", encoding="utf-8") as f:
        for r in sorted(top, key=lambda r: r["topic_id"]):
            rec = {k: v for k, v in r.items() if k not in ("topic_id", "seed")}; rec["id"] = r["topic_id"]
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
    print(f"  {len(top)} Typhoon picks")


def arms(d):
    T = {r["id"]: r for r in rows(f"{d}/typhoon_picks.jsonl")}
    V = {r["id"]: r for r in rows(f"{d}/veto_picks.jsonl")}
    assert set(T) == set(V), "veto picks do not cover the topics"
    assert all(V[i]["prompt"] == T[i]["prompt"] for i in T)
    with open(f"{d}/pick_T.jsonl", "w", encoding="utf-8") as f:
        for i in sorted(T): f.write(json.dumps({"id": i, "prompt": T[i]["prompt"], "poem": T[i]["poem"]}, ensure_ascii=False) + "\n")
    ch = sorted(i for i in T if V[i]["poem"] != T[i]["poem"])
    with open(f"{d}/pick_V.jsonl", "w", encoding="utf-8") as f:
        for i in ch: f.write(json.dumps({"id": i, "prompt": V[i]["prompt"], "poem": V[i]["poem"]}, ensure_ascii=False) + "\n")
    fd = sorted(i for i in ch if V[i].get("form_driven"))
    json.dump({"veto_changed": ch, "form_driven": fd}, open(f"{d}/veto_changed.json", "w"), indent=0)
    print(f"  {len(T)} topics; veto pick differs on {len(ch)} ({len(fd)} form-driven)")


def locmap(d):
    for a in ("T", "V"):
        P_ = rows(f"{d}/pick_{a}.jsonl")
        if not P_:                                     # no veto change at all: arm V is skipped (review)
            json.dump({}, open(f"{d}/{a}_line.json", "w")); print(f"  arm {a}: no picks, skipped"); continue
        S = rows(f"{d}/{a}_scores.jsonl")
        assert [p["id"] for p in P_] == [s["id"] for s in S] and all(p["poem"] == s["poem"] for p, s in zip(P_, S)), \
            f"{a}: scores do not belong to the picks"
        m = {s["id"]: min(range(4), key=lambda i: s["score"]["line_mean"][i]) + 1 for s in S if s["score"]}
        miss = [s["id"] for s in S if not s["score"]]
        print(f"  arm {a}: {len(m)} located, {len(miss)} without a score (kept unchanged, as round 1)")
        assert len(miss) <= max(3, len(S) // 20), "too many picks fail to round-trip"
        json.dump(m, open(f"{d}/{a}_line.json", "w"), indent=0)


def accept(cands, scores, loc):
    """done/105's accept rule, verbatim (also vetoref_select.select)."""
    from syl_align import syls_aligned as SA
    from canvas_mlm import STOP_SYL
    assert len(cands) == len(scores) and all(c["id"] == s["id"] and c["poem"] == s["poem"] for c, s in zip(cands, scores)), \
        "scores do not belong to these candidates"
    topic_of = {c["id"]: {x for x in SA(c.get("prompt") or "", 12) if x not in STOP_SYL and len(x) > 1} for c in cands}
    orig, var, unscored = {}, collections.defaultdict(list), []
    for c, s in zip(cands, scores):
        if c["variant"] == "orig":
            if s["score"]: orig[c["id"]] = (c, s["score"]["line_mean"])
            else: unscored.append(c)
        elif s["score"]: var[c["id"]].append((c, s["score"]["line_mean"]))
    best = {}
    for i, (o, om) in orig.items():
        L = loc[i] - 1
        keep = set(SA(o["poem"].split("\n")[L], 12)) & topic_of[i]
        ok = [(c, m) for c, m in var[i] if keep <= set(SA(c["poem"].split("\n")[L], 12))]
        b = o
        if ok:
            c, m = max(ok, key=lambda x: x[1][L])
            if m[L] >= om[L] + 0.5: b = c
        best[i] = b
    for c in unscored: best[c["id"]] = c
    return best, {c["id"] for c in unscored}


def select(d):
    from klon8 import prosody as P
    from syl_align import syls_aligned as SA
    from canvas_mlm import STOP_SYL
    R1, unT = accept(rows(f"{d}/T_cands.jsonl"), rows(f"{d}/T_cands_scores.jsonl"), {k: int(v) for k, v in json.load(open(f"{d}/T_line.json")).items()})
    vc = json.load(open(f"{d}/veto_changed.json"))
    if vc["veto_changed"]:
        VRc, unV = accept(rows(f"{d}/V_cands.jsonl"), rows(f"{d}/V_cands_scores.jsonl"), {k: int(v) for k, v in json.load(open(f"{d}/V_line.json")).items()})
    else: VRc, unV = {}, set()
    VR = {i: VRc.get(i, R1[i]) for i in R1}           # V == T: the R1 poem by construction
    assert set(VRc) == set(vc["veto_changed"]) and len(R1) == len(rows(f"{d}/pick_T.jsonl"))
    excl = set(vc["form_driven"]) | unV | unT
    changed = sorted(i for i in R1 if VR[i]["poem"] != R1[i]["poem"] and i not in excl)
    def stut(p): return sum(a == b for ln in p.split("\n") for a, b in zip(SA(ln, 12), SA(ln, 12)[1:]))
    def xrep(p):
        ls = [set(x for x in SA(l, 12) if x not in STOP_SYL and len(x) > 1) for l in p.split("\n")]
        c = collections.Counter(x for s_ in ls for x in s_); return sum(v - 1 for v in c.values() if v > 1)
    m = lambda f, ps: sum(map(f, ps)) / max(len(ps), 1)
    new_p, old_p = [VR[i]["poem"] for i in changed], [R1[i]["poem"] for i in changed]
    form = 100.0 if not new_p else 100 * sum(bool(P.score(p, n_bot_expected=1)["all_pass"]) for p in new_p) / len(new_p)
    rep = {"stutters": (m(stut, new_p), m(stut, old_p)), "cross-line repeats": (m(xrep, new_p), m(xrep, old_p))}
    print(f"  {len(R1)} topics; R1 refilled {sum(r['variant'] != 'orig' for r in R1.values())}; VR differs from R1 on "
          f"{len(changed)} endpoint topics (excluded: form-driven {sorted(set(vc['form_driven']))}, unrefilled {sorted(unV | unT)})")
    print(f"  changed poems: form {form:.0f}%; " + "; ".join(f"{k} {a:.3f} (R1 {b:.3f})" for k, (a, b) in rep.items()))
    P.save_cache()
    for f_ in ("r1.jsonl", "vr.jsonl", "changed_ids.json", "PASS"):
        if os.path.exists(f"{d}/{f_}"): os.remove(f"{d}/{f_}")
    if not changed: print("  nothing changed: nothing to judge"); sys.exit(3)
    TOL = {"stutters": 0.10, "cross-line repeats": 0.50}
    dead = (["form"] if form < 98 else []) + [k for k, (a, b) in rep.items() if a > b + TOL[k]]
    print("  KILLED before judging: " + ", ".join(dead) if dead else "  pre-judge gates PASS")
    if dead: sys.exit(2)
    for name, A in (("r1", R1), ("vr", VR)):
        with open(f"{d}/{name}.jsonl", "w", encoding="utf-8") as f:
            for i in sorted(A): f.write(json.dumps({"id": i, "prompt": A[i]["prompt"], "poem": A[i]["poem"],
                                                    "variant": A[i]["variant"]}, ensure_ascii=False) + "\n")
    json.dump({"endpoint": changed, "form_driven": sorted(set(vc["form_driven"])), "unrefilled": sorted(unV | unT)},
              open(f"{d}/changed_ids.json", "w"), indent=0)
    open(f"{d}/PASS", "w").write("gates passed\n")


def export(d):
    assert os.path.exists(f"{d}/PASS"), "gates did not pass"
    ids = json.load(open(f"{d}/changed_ids.json"))["endpoint"]
    for name, src in (("round1", f"{d}/r1.jsonl"), ("vetoref", f"{d}/vr.jsonl")):
        A = {r["id"]: r for r in rows(src)}
        with open(f"{d}/judge_{name}.jsonl", "w", encoding="utf-8") as f:
            for i in ids: f.write(json.dumps({"id": i, "prompt": A[i]["prompt"], "poem": A[i]["poem"]}, ensure_ascii=False) + "\n")
    # the 6 pilot (test30) topics 129 left out as tuning topics: already decoded and refilled in runs/VETOREF
    pil = {r["id"] for r in rows("data/prompts/test30_pilot.jsonl")}
    ep = json.load(open("runs/VETOREF/endpoint.json"))
    excl = set(ep["form_driven"]) | set(ep["not_refilled"])
    vr = {r["id"]: r for r in rows("runs/VETOREF/best.jsonl")}; r1 = {r["id"]: r for r in rows("runs/LOCREF/best.jsonl")}
    pc = sorted(i for i in pil if vr[i]["poem"] != r1[i]["poem"] and i not in excl and i not in ep["endpoint"])
    # the same drift gates on the pilot pairs (review); failing pairs are dropped and reported, not judged
    from syl_align import syls_aligned as SA
    from canvas_mlm import STOP_SYL
    from klon8 import prosody as P
    def stut(p): return sum(a == b for ln in p.split("\n") for a, b in zip(SA(ln, 12), SA(ln, 12)[1:]))
    def xrep(p):
        ls = [set(x for x in SA(l, 12) if x not in STOP_SYL and len(x) > 1) for l in p.split("\n")]
        c = collections.Counter(x for s_ in ls for x in s_); return sum(v - 1 for v in c.values() if v > 1)
    m = lambda f, ps: sum(map(f, ps)) / max(len(ps), 1)
    np_, op_ = [vr[i]["poem"] for i in pc], [r1[i]["poem"] for i in pc]
    form = 100.0 if not np_ else 100 * sum(bool(P.score(q, n_bot_expected=1)["all_pass"]) for q in np_) / len(np_)
    dead = (["form"] if form < 98 else []) + (["stutters"] if m(stut, np_) > m(stut, op_) + 0.10 else []) + \
           (["cross-line repeats"] if m(xrep, np_) > m(xrep, op_) + 0.50 else [])
    print(f"  pilot pairs: form {form:.0f}%; stutters {m(stut, np_):.3f} vs {m(stut, op_):.3f}; repeats {m(xrep, np_):.3f} vs {m(xrep, op_):.3f}"
          + (f" - KILLED ({', '.join(dead)}): pilot pairs not exported" if dead else " - gates PASS"))
    if dead: pc = []
    for name, A in (("round1", r1), ("vetoref", vr)):
        with open(f"{d}/judge_pilot_{name}.jsonl", "w", encoding="utf-8") as f:
            for i in pc: f.write(json.dumps({"id": i, "prompt": A[i]["prompt"], "poem": A[i]["poem"]}, ensure_ascii=False) + "\n")
    print(f"  judge inputs: {len(ids)} extra-topic pairs, {len(pc)} pilot pairs {pc}")


if __name__ == "__main__":
    cmd, d = sys.argv[1], (sys.argv[2] if len(sys.argv) > 2 else "runs/VETOMORE")
    {"topics": topics, "picks": picks, "arms": arms, "locmap": locmap, "select": select, "export": export}[cmd](d)