src/scripts/confirm2_select.py

"""Confirmation 2 (job 189, runs/CSEL/CONFIRM2_PREREG.md), stage 0: fresh held-out titles + human controls.

The TEST side of the title split (data/modern_bot_test_prov.jsonl) is the only source of titles the models never saw.
fresh120 (scripts/fresh_titles.py) took EVERY test title that passed its filters: re-running those filters today,
minus fresh120, leaves 0 titles. Confirmation 2 therefore uses the remaining test titles under a RELAXED rule,
written before any decode (the relaxations are listed in the prereg and in `rule` below):

  exclusions  every key in data/prompts/fresh_exclusion_keys.json (frozen 2026-09-28: all prompt files, judge batch
              topics, runs/*.jsonl prompts) + a CURRENT rescan with fresh_titles.used_keys() (picks up everything
              decoded or judged since) + all 106 fresh120 titles (incl. the user's exclusions F017/F039/F076) +
              dev49_pilot / test150 / test30_pilot (asserted) + every TRAIN title key. Frozen at first build to
              data/prompts/confirm2_exclusion_keys.json so a rebuild cannot drift.
  title       decoration-free key (split_titles.key); every raw spelling of the key: no forum-noise keyword (fresh120's
              list + KW_EXTRA), no letter outside the Thai block (Latin, digits, CJK ... - the stripped topic would
              differ from what the poet wrote to); no doubled or leading combining mark (typing noise). RELAXED: 2-12 syllables (fresh120: 6-22 characters).
  control     the human first stanza under the title, spaces stripped, 4 lines. RELAXED tiers (fresh120: tier 1 only):
                1 = opening post, stanza 0 (fresh120's rule)
                2 = opening post, its first 4-line klon8 stanza (the opener began with a non-klon8 stanza)
                3 = the thread's first 4-line klon8 stanza in any post (lowest post, then stanza)
              ties by lowest source id. The tier is recorded; per-tier rates are a secondary, descriptive split.
  order       sorted keys, seeded shuffle (seed 11), ids G001.. ; all eligible titles are kept (no cap).

Usage: confirm2_select.py build      (asserts >= 40 titles; prints count + sha256)
       confirm2_select.py test [env C2_SEL_FROZEN=1: skip the fresh rescan, as inside the job]                                 (CPU checks on real data, writes only to a temp dir)
"""
import os, sys, re, json, random, hashlib, argparse, collections, unicodedata, tempfile
sys.path[:0] = [".", "scripts"]

KW_EXTRA = ("สวัสดี", "ทักทาย", "แต่งโดย", "วิจารณ์", "บทละคร", "กลอนเจ็ด", "คำกลอน", "หัววรรค", "บ้านกลอน",
            "อิอิ", "ครับ", "ค่ะ")
NON_THAI_LETTER = lambda t: any(unicodedata.category(c).startswith(("L", "N")) and not ("฀" <= c <= "๿")
                                for c in unicodedata.normalize("NFC", t))
THAI_DIGIT = re.compile(r"[๐-๙]")
# typing noise: a key that starts with a combining mark or doubles one is not a clean topic
MARK_NOISE = re.compile(r"^[ัิ-ฺ็-๎]|([ัิ-ฺ็-๎])\1")
SEED, MIN_SYL, MAX_SYL = 11, 2, 12
OUT, HUMAN, EXCL = "data/prompts/confirm2.jsonl", "data/pools/confirm2_human.jsonl", "data/prompts/confirm2_exclusion_keys.json"
PRIOR_SETS = ("data/prompts/fresh120.jsonl", "data/prompts/dev49_pilot.jsonl", "data/prompts/test150.jsonl",
              "data/prompts/test30_pilot.jsonl")


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


def prior_keys():
    import fresh_titles as FT
    ks = set()
    for f in PRIOR_SETS:
        for r in rows(f): ks.add(FT.key(r.get("text") or r.get("prompt")))
    return ks


def exclusion_keys(path):
    """frozen at first build; the rescan skips this confirmation's own files"""
    import fresh_titles as FT
    if os.path.exists(path): return set(json.load(open(path, encoding="utf-8")))
    hide = [OUT]                                  # our own prompt file must not exclude itself on the first build
    moved = [(p, p + ".c2hide") for p in hide if os.path.exists(p)]
    for a, b in moved: os.replace(a, b)
    try: now = FT.used_keys()
    finally:
        for a, b in moved: os.replace(b, a)
    ex = set(json.load(open("data/prompts/fresh_exclusion_keys.json", encoding="utf-8"))) | now | prior_keys()
    ex.discard("")
    json.dump(sorted(ex), open(path, "w", encoding="utf-8"), ensure_ascii=False, indent=0)
    return ex


def tier_pick(rs):
    t1 = [r for r in rs if r["post"] == 0 and r["stanza"] == 0]
    if t1: return 1, min(t1, key=lambda r: r["id"])
    t2 = [r for r in rs if r["post"] == 0]
    if t2: return 2, min(t2, key=lambda r: (r["stanza"], r["id"]))
    return 3, min(rs, key=lambda r: (r["post"], r["stanza"], r["id"]))


def eligible(prov, excl, train_keys, g2p_len):
    import fresh_titles as FT
    by_key = collections.defaultdict(list)
    for r in prov: by_key[FT.key(r["title"])].append(r)
    why, keep = collections.Counter(), {}
    for k in sorted(by_key):
        rs = by_key[k]; raws = {r["title"] for r in rs}
        if not k: why["empty key"] += 1; continue
        if k in excl: why["used before"] += 1; continue
        if k in train_keys: why["train title"] += 1; continue
        if any(any(w in t for w in FT.KW + KW_EXTRA) for t in raws): why["forum keyword"] += 1; continue
        if any(NON_THAI_LETTER(t) or THAI_DIGIT.search(t) for t in raws): why["non-Thai letter/digit"] += 1; continue
        if not FT.THAI_ONLY.fullmatch(k): why["not Thai-only"] += 1; continue
        if MARK_NOISE.search(k): why["doubled/leading combining mark"] += 1; continue
        n = g2p_len(k)
        if not MIN_SYL <= n <= MAX_SYL: why[f"syllables outside {MIN_SYL}-{MAX_SYL}"] += 1; continue
        keep[k] = rs
    return keep, why, len(by_key)


def build(out=OUT, human=HUMAN, excl_path=EXCL, prov_path="data/modern_bot_test_prov.jsonl"):
    import fresh_titles as FT
    from klon8 import prosody as P
    train_keys = {FT.key(json.loads(l)["title"]) for l in open("data/modern_bot_train.jsonl", encoding="utf-8")}
    excl = exclusion_keys(excl_path)
    assert prior_keys() <= excl, "a prior set is not excluded"
    keep, why, n_keys = eligible(rows(prov_path), excl, train_keys, lambda k: len(P.g2p_line(k)))
    ks = sorted(keep); random.Random(SEED).shuffle(ks)
    assert len(ks) >= 40, f"only {len(ks)} titles"
    assert not (set(ks) & excl) and not (set(ks) & train_keys) and not (set(ks) & prior_keys())
    tiers = collections.Counter()
    os.makedirs(os.path.dirname(human), exist_ok=True)
    with open(out, "w", encoding="utf-8") as f, open(human, "w", encoding="utf-8") as h:
        for i, k in enumerate(ks, 1):
            tier, op = tier_pick(keep[k]); tiers[tier] += 1
            poem = op["poem"].replace(" ", "")
            assert len(poem.split("\n")) == 4 and all(poem.split("\n")), f"{k}: control is not 4 lines"
            gid = f"G{i:03d}"
            f.write(json.dumps({"id": gid, "form": "topic", "text": k, "raw_title": op["title"], "set": "CONFIRM2",
                                "syl": len(P.g2p_line(k)), "control_tier": tier}, ensure_ascii=False) + "\n")
            h.write(json.dumps({"id": gid, "prompt": k, "poem": poem, "src_id": op["id"], "board": op.get("board"),
                                "thread": op["thread"], "post": op["post"], "stanza": op["stanza"], "tier": tier},
                               ensure_ascii=False) + "\n")
    P.save_cache()
    print(f"  test-split title keys {n_keys}; excluded: " + ", ".join(f"{k} {v}" for k, v in why.most_common()))
    print(f"  kept {len(ks)} titles (seed {SEED}); control tiers {dict(sorted(tiers.items()))}; exclusion keys {len(excl)}")
    print(f"  {out} sha256 {sha(out)}\n  {human} sha256 {sha(human)}\n  {excl_path} sha256 {sha(excl_path)}")
    return len(ks)


def _test():
    import fresh_titles as FT
    # tier rule on hand-built rows
    R = lambda i, p, s: {"id": i, "post": p, "stanza": s}
    assert tier_pick([R("b", 0, 0), R("a", 0, 0), R("c", 1, 0)]) == (1, R("a", 0, 0))
    assert tier_pick([R("b", 0, 3), R("a", 0, 5), R("c", 1, 0)]) == (2, R("b", 0, 3))
    assert tier_pick([R("b", 2, 0), R("a", 1, 4), R("c", 1, 2)]) == (3, R("c", 1, 2))
    # title rule on hand-built titles
    assert NON_THAI_LETTER("รัก VS ใจ") and NON_THAI_LETTER("恭喜 รวย") and not NON_THAI_LETTER("…รัก(ใจ)!!")
    fake = [dict(title=t, id=f"x{n}", post=0, stanza=0, poem="a\nb\nc\nd", thread=1) for n, t in enumerate(
        ["รักเธอ", "ครูครับ", "วันเกิดเธอ", "ใจ VS ใจ", "ฝนตก", "ฝน", "used", "ที่่รัก", "่่รัก"])]
    keep, why, _ = eligible(fake, {"ฝนตก"}, {"รักเธอ"}, lambda k: 1 if k == "ฝน" else 3)
    assert set(keep) == set() and why["used before"] == 1 and why["train title"] == 1 and why["forum keyword"] == 2 \
        and why["non-Thai letter/digit"] == 1 and why["syllables outside 2-12"] == 1 \
        and why["doubled/leading combining mark"] == 2, (keep, why)
    assert not MARK_NOISE.search("น้ำตาพราวที่ราวฟ้า") and not MARK_NOISE.search("เก็บมาลัยไว้กราบแม่")
    # the real build into a temp dir (the exclusion list is REBUILT there: it must equal a fresh derivation)
    d = tempfile.mkdtemp(prefix="c2sel_", dir=os.environ.get("C2_TMP") or None)
    if os.environ.get("C2_SEL_FROZEN"):        # inside job 189: runs/ now holds this confirmation's own outputs, so a
        import shutil                          # fresh rescan would exclude its titles; check the frozen derivation only
        shutil.copy(EXCL, f"{d}/ex.json")
    n = build(f"{d}/p.jsonl", f"{d}/h.jsonl", f"{d}/ex.json")
    P_, H_ = rows(f"{d}/p.jsonl"), rows(f"{d}/h.jsonl")
    assert len(P_) == len(H_) == n and [r["id"] for r in P_] == [r["id"] for r in H_]
    assert all(p["text"] == h["prompt"] for p, h in zip(P_, H_))
    ex = set(json.load(open(f"{d}/ex.json", encoding="utf-8")))
    train_keys = {FT.key(json.loads(l)["title"]) for l in open("data/modern_bot_train.jsonl", encoding="utf-8")}
    for f in PRIOR_SETS + ("data/prompts/fresh_exclusion_keys.json",):
        ks = set(json.load(open(f))) if f.endswith(".json") else {FT.key(r.get("text") or r.get("prompt")) for r in rows(f)}
        assert not ({r["text"] for r in P_} & ks), f"overlap with {f}"
    fresh = rows("data/prompts/fresh120.jsonl"); assert len(fresh) == 106
    assert not ({r["text"] for r in P_} & {r["text"] for r in fresh} | set()) and {"F017", "F039", "F076"} <= {r["id"] for r in fresh}
    assert not ({r["text"] for r in P_} & train_keys) and {r["text"] for r in P_} <= set(r["text"] for r in P_) and not ({r["text"] for r in P_} & ex)
    test_keys = {FT.key(r["title"]) for r in rows("data/modern_bot_test_prov.jsonl")}
    assert {r["text"] for r in P_} <= test_keys, "a title is not from the test split"
    prov = {r["id"]: r for r in rows("data/modern_bot_test_prov.jsonl")}
    for h in H_[:3] + H_[-3:]:                 # controls by content: the source row, under the same key, spaces stripped
        src = prov[h["src_id"]]
        assert FT.key(src["title"]) == h["prompt"] and src["poem"].replace(" ", "") == h["poem"]
    if os.path.exists(OUT):                    # the frozen list must be reproducible from the frozen exclusion file
        d2 = tempfile.mkdtemp(prefix="c2sel2_", dir=os.environ.get("C2_TMP") or None)
        import shutil; shutil.copy(EXCL, f"{d2}/ex.json")
        build(f"{d2}/p.jsonl", f"{d2}/h.jsonl", f"{d2}/ex.json")
        assert open(f"{d2}/p.jsonl", "rb").read() == open(OUT, "rb").read(), "frozen title list does not reproduce"
        assert open(f"{d2}/h.jsonl", "rb").read() == open(HUMAN, "rb").read(), "frozen human controls do not reproduce"
        print("  frozen list reproduces byte-for-byte from the frozen exclusion keys")
    print(f"confirm2_select: tests pass ({n} titles, temp {d})")


if __name__ == "__main__":
    if sys.argv[1] == "build": build()
    elif sys.argv[1] == "test": _test()
    else: sys.exit("usage: confirm2_select.py build|test")