Blame for src/scripts/twolayer_full.py

"""Job 187 (runs/CSEL/TWOLAYER_FULL_PREREG.md): the two-layer canvas (job 186's HIERARCHICAL model, SYLLABLE commit
unit) through the production selection (64 candidates per topic, Typhoon top-8 rank 0, NO refill), two arms:
  A  = 186's decode unchanged (twolayer_canvas.py decode --unit syllable)
  B  = A + the STUTTER FIX at decode (--stutter-fix): a cell may not take a token whose SOUND equals the sound of its
       committed (or forced) left or right neighbour in the same วรรค, unless the two cells are consecutive syllables of
       one word (the word is in the lexicon: attested as a word) or the WORD PAIR (left word, right word) is an attested
       reduplication / same-sound join in the training grids, >= --allow-min times (default 2; a single occurrence can
       be a typo). Default OFF: with the flag off this module's decode is byte-identical to twolayer_canvas.py decode.
New file; nothing shared is patched. twolayer_canvas.Decoder is replaced by StutterDecoder only inside this process and
only with --stutter-fix.

  allow DATA_DIR OUT.json [--min N]     allowlist from DATA_DIR/poems.jsonl (twolayer_data.py build; 186's data)
  decode --model RUN --prompts P --out O [--seed S] [--limit N] [--stutter-fix --allow A.json]  (= 186's decode + flag)
  merge D NT S0 NS                      D/cands/c{S0..}.jsonl -> D/cands.jsonl (+ seed); coverage / flag asserts
  formgate D                            candidate form >= 98% (exit 1 = arm dropped before Typhoon)
  picks D                               D/typhoon_top8.jsonl rank 0 -> D/pick.jsonl (id, prompt, poem, seed, all_pass)
  attested TRAIN OUT.json [--min N]     written syllables the training text repeats adjacently >= N x (gate metric)
  gates OUT ATT CTL_CANDS CTL_FINAL CTL_PICK A=D B=D   pre-judge gates (prereg) -> OUT/GATES.json; prints PASSING ...
  pools OUT FB_DIR ARM=D ...            RAW-text judge inputs per topic half: fixfinal, fixpick, arms (all 49 each)
  score OUT FB_DIR ARM=D ...            after judging runs/TLF_{CLEAN,ALIGN}0{h}/run1.json: paired stats + decision
  test                                  CPU tests of the pure parts (no model)
"""
import os, sys, json, collections, argparse, random
sys.path[:0] = [".", "scripts"]

POOL_PREFIX = "TLF"


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)


# ---------------------------------------------------------------- allowlist

def same_sound_pairs(grids):
    """Counter {(left word, right word): n} over adjacent cells of one วรรค with the SAME sound that are not consecutive
    syllables of one word occurrence. grids: iterable of วรรค, each a list of cells [word, k, n, phon]."""
    c = collections.Counter()
    for line in grids:
        for a, b in zip(line, line[1:]):
            if a[3] != b[3]: continue
            if a[0] == b[0] and a[1] + 1 == b[1] and a[2] == b[2]: continue        # inside one word (นานา)
            c[(a[0], b[0])] += 1
    return c


def build_allow(data_dir, out, min_n=2):
    c = same_sound_pairs(line for l in open(f"{data_dir}/poems.jsonl", encoding="utf-8") for line in json.loads(l)["grid"])
    keep = {f"{a}|{b}": n for (a, b), n in sorted(c.items()) if n >= int(min_n)}
    json.dump({"min": int(min_n), "source": f"{data_dir}/poems.jsonl", "pairs": keep}, open(out, "w", encoding="utf-8"),
              ensure_ascii=False, indent=0)
    print(f"  allowlist {out}: {len(keep)} of {len(c)} same-sound word pairs attested >= {min_n}x "
          f"({sum(keep.values())} of {sum(c.values())} occurrences)")


def load_allow(path, X):
    """{left word id: set(right word ids)} and the reverse, for words in the lexicon"""
    import torch
    d = json.load(open(path, encoding="utf-8"))["pairs"]
    fwd, rev = collections.defaultdict(set), collections.defaultdict(set)
    for k in d:
        a, b = k.split("|")
        ia, ib = X.wid.get(a), X.wid.get(b)
        if ia is None or ib is None or ia < 4 or ib < 4: continue
        fwd[ia].add(ib); rev[ib].add(ia)
    t = lambda s: torch.tensor(sorted(s), dtype=torch.long)
    return {k: t(v) for k, v in fwd.items()}, {k: t(v) for k, v in rev.items()}


# ---------------------------------------------------------------- decoder with the stutter fix

def make_decoder_class():
    import torch
    import twolayer_canvas as TC

    class StutterDecoder(TC.Decoder):
        """TC.Decoder + the stutter ban, applied in lp() after every mask of the base decoder (so span admissibility
        in span_scores / the tilings sees it). ALLOW = (fwd, rev) from load_allow, set on the class."""
        ALLOW = ({}, {})
        n_banned = 0

        def neighbour_tok(self, k):
            if k in self.committed: return int(self.tok[k])
            return self.forced.get(k)

        def ban_stutters(self, lp):
            X = self.X; fwd, rev = self.ALLOW
            ts, tw, tk, tn = X.tok_sound, X.tok_word, X.tok_k, X.tok_n
            for i in range(TC.L):
                for j in range(self.lens[i]):
                    k = self.idx(i, j)
                    if k in self.committed: continue
                    bad = None
                    if j > 0:
                        t = self.neighbour_tok(self.idx(i, j - 1))
                        if t is not None and t >= 4 and int(ts[t]) != 0:
                            b = ts == int(ts[t])
                            if int(tk[t]) < int(tn[t]) - 1: b[t + 1] = False             # next syllable of the same word
                            w = int(tw[t])
                            if w in fwd: b &= ~torch.isin(tw, fwd[w])
                            bad = b
                    if j + 1 < self.lens[i]:
                        t = self.neighbour_tok(self.idx(i, j + 1))
                        if t is not None and t >= 4 and int(ts[t]) != 0:
                            b = ts == int(ts[t])
                            if int(tk[t]) > 0: b[t - 1] = False                          # previous syllable of the same word
                            w = int(tw[t])
                            if w in rev: b &= ~torch.isin(tw, rev[w])
                            bad = b if bad is None else (bad | b)
                    if bad is not None:
                        self.n_banned += int((bad & (lp[k] > -1e8)).sum())
                        lp[k] += torch.where(bad, -1e9, 0.0)
            return lp

        def lp(self, hold=True):
            return self.ban_stutters(super().lp(hold))

    return StutterDecoder


def stutter_pairs(D):
    """(line, j) of adjacent same-sound cells of a finished Decoder that are neither inside one word nor allowed"""
    X = D.X; fwd, _ = type(D).ALLOW if hasattr(type(D), "ALLOW") else ({}, {})
    out = []
    for i in range(len(D.lens)):
        for j in range(1, D.lens[i]):
            a, b = int(D.tok[D.idx(i, j - 1)]), int(D.tok[D.idx(i, j)])
            if a < 4 or b < 4 or int(X.tok_sound[a]) != int(X.tok_sound[b]): continue
            if b == a + 1 and int(X.tok_k[b]) == int(X.tok_k[a]) + 1: continue
            wa, wb = int(X.tok_word[a]), int(X.tok_word[b])
            if wa in fwd and bool((fwd[wa] == wb).any()): continue
            out.append((i, j))
    return out


def decode_cli(argv):
    """= twolayer_canvas.py decode (same records, same RNG streams) + --stutter-fix; extra keys only with the flag"""
    import torch, random as _r
    import twolayer_canvas as TC
    from klon8 import prosody as P
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True); ap.add_argument("--prompts", required=True); ap.add_argument("--out", required=True)
    ap.add_argument("--seed", type=int, default=0); ap.add_argument("--temp", type=float, default=0.8); ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--unit", choices=["syllable", "word", "wak", "bot"], default="word")
    ap.add_argument("--refine", type=int, default=3); ap.add_argument("--topic-boost", type=float, default=1.5)
    ap.add_argument("--stutter-fix", action="store_true", help="ban same-sound neighbours unless attested (default off)")
    ap.add_argument("--allow", default="", help="allowlist from `allow` (required with --stutter-fix)")
    a = ap.parse_args(argv)
    model, X = TC.load(a.model)
    if a.stutter_fix:
        assert a.allow, "--stutter-fix needs --allow"
        SD = make_decoder_class(); SD.ALLOW = load_allow(a.allow, X)
        TC.Decoder = SD                                      # decode_one builds TC.Decoder (this process only)
    rows_ = [json.loads(l) for l in open(a.prompts, encoding="utf-8")]
    if a.limit: rows_ = rows_[:a.limit]
    rng_len = _r.Random(a.seed)
    with open(a.out + ".tmp", "w", encoding="utf-8") as f:
        for n, r in enumerate(rows_):
            text = r.get("text") or r.get("prompt") or r.get("title") or r.get("topic")
            assert text, f"prompt row without text: {r}"
            want = TC.pick_lens(rng_len)
            lines, D, tries = TC.decode_one(model, X, text, unit=a.unit, temp=a.temp, lens=want, refine=a.refine,
                                            topic_boost=a.topic_boost, gen=torch.Generator().manual_seed(a.seed * 100003 + n))
            poem = "\n".join(lines); sc = P.score(poem, n_bot_expected=1)
            got = [len(P.g2p_line(x)) for x in lines]
            rec = {"id": r.get("id"), "prompt": text, "sample": 0, "poem": poem, "unit": a.unit, "syl_target": want,
                   "syl": got, "retries": tries, "forward_passes": D.n_fwd, "all_pass": sc["all_pass"],
                   "structure": sc["structure"], "ext_rhyme": sc["ext_rhyme"], "tone": sc["tone"], "internal": sc["internal"]}
            if a.stutter_fix:
                rec["stutter_fix"] = True; rec["n_banned"] = D.n_banned; rec["stutters_left"] = len(stutter_pairs(D))
            f.write(json.dumps(rec, ensure_ascii=False) + "\n"); f.flush()
            print(f"{n}/{len(rows_)} all_pass={sc['all_pass']} syl={got} want={want} retries={tries}", file=sys.stderr, flush=True)
    os.replace(a.out + ".tmp", a.out)
    P.save_cache()


# ---------------------------------------------------------------- pipeline

def merge(d, nt, s0, ns, fix=None):
    nt, s0, ns = int(nt), int(s0), int(ns)
    rs = []
    for s in range(s0, s0 + ns):
        x = rows(f"{d}/cands/c{s}.jsonl"); assert len(x) == nt, f"{d}/cands/c{s}.jsonl: {len(x)} rows, want {nt}"
        rs += [dict(r, seed=s) for r in x]
    assert all(r["unit"] == "syllable" for r in rs), "not the syllable unit"
    if fix is not None:
        want = fix == "1"
        assert all(bool(r.get("stutter_fix")) == want for r in rs), f"rows decoded with stutter_fix != {want}"
    assert len({r["id"] for r in rs}) == nt and all(sum(1 for r in rs if r["id"] == i) == ns for i in {r["id"] for r in rs})
    wrows(f"{d}/cands.jsonl", rs)
    print(f"  {d}: {len(rs)} candidates for {nt} topics, seeds {s0}-{s0 + ns - 1}; form {100 * sum(r['all_pass'] for r in rs) / len(rs):.1f}%")


def picks(d, nt=49):
    top = [r for r in rows(f"{d}/typhoon_top8.jsonl") if r["sel_rank"] == 0]
    assert len({r["topic_id"] for r in top}) == len(top) == int(nt), "rank-0 picks: missing or duplicate topics"
    wrows(f"{d}/pick.jsonl", [{"id": r["topic_id"], "prompt": r["prompt"], "poem": r["poem"], "seed": r["seed"],
                               "all_pass": r["all_pass"]} for r in sorted(top, key=lambda r: r["topic_id"])])
    print(f"  {d}: {len(top)} Typhoon picks (no refill)")


def formgate(d, thr=98.0):
    """pre-registered form gate, run right after merge (before Typhoon): candidate form >= 98% or the arm is dropped
    (exit 1); writes D/FORMGATE"""
    rs = rows(f"{d}/cands.jsonl"); f = 100 * sum(bool(r["all_pass"]) for r in rs) / len(rs)
    ok = f >= float(thr)
    open(f"{d}/FORMGATE", "w").write(f"{f:.2f} {'PASS' if ok else 'FAIL'}\n")
    print(f"  {d}: candidate form {f:.1f}% -> " + ("form gate PASS" if ok else "KILLED before Typhoon: form < 98%"))
    return ok


def form_of(poem):
    from klon8 import prosody as P
    return bool(P.score(poem, n_bot_expected=1)["all_pass"])


def attested_syllables(train, out, min_n=2):
    """written syllables X (v1 segmenter, as every stutter gate) that the TRAINING TEXT repeats adjacently (XX, with ๆ
    written out as the repeated word) >= min_n times: an independent text-level check, not the decoder's own word-pair
    allowlist. Writes {syllable: count} for all repeats and the kept set."""
    from syl_align import syls_aligned as SA
    import twolayer_data as TD
    c = collections.Counter()
    for l in open(train, encoding="utf-8"):
        for ln in json.loads(l)["poem"].split("\n"):
            s = SA("".join(TD.line_words(ln)), 40)
            c.update(a for a, b in zip(s, s[1:]) if a == b)
    keep = sorted(k for k, v in c.items() if v >= int(min_n))
    json.dump({"min": int(min_n), "source": train, "kept": keep, "counts": dict(c)}, open(out, "w", encoding="utf-8"), ensure_ascii=False)
    print(f"  attested reduplication syllables {out}: {len(keep)} of {len(c)} repeated syllables (>= {min_n}x)")


def stutters_unattested(poem, att):
    import fixbase_pipe as FP
    return FP.stutters(poem, att)


def gate_fail(c, f, stutter_gate):
    """pre-registered (TWOLAYER_FULL_PREREG.md), on the 64 x 49 raw candidates: form >= 98% (every arm); UNATTESTED
    stutters (adjacent identical written syllables, v1 segmenter, excluding syllables the training text repeats >= 2x)
    <= control + 0.10 per poem (arm B only - arm A is reported, not killed). Raw stutters (186's metric) are reported."""
    dead = []
    if f["form"] < 98: dead.append("form < 98%")
    if stutter_gate and f["stutters_unattested"] > c["stutters_unattested"] + 0.10: dead.append("unattested stutters")
    return dead


def reported_186(c, f):
    """186's other gates, REPORTED only (no kill here): repeats, distinct-2, rhyme-word variety, unseen, final-word freq"""
    import twolayer_pipe as TP
    return [x for x in TP.gate_fail(c, f) if x not in ("form < 98%", "stutters")]


def gates(out, att_json, ctl_cands, ctl_final, ctl_pick, *arms):
    import twolayer_pipe as TP
    A = dict(x.split("=", 1) for x in arms)
    assert set(A) <= {"A", "B"} and A, "arms: A=DIR and/or B=DIR"
    att = frozenset(json.load(open(att_json, encoding="utf-8"))["kept"]); assert len(att) > 100, "attested set too small"
    def extra(R):
        return {"stutters_unattested": sum(stutters_unattested(r["poem"], att) for r in R) / len(R)}
    cand = {"CTL": rows(ctl_cands)}; cand.update({k: rows(f"{d}/cands.jsonl") for k, d in A.items()})
    fin = {"CTL_final": [dict(r, all_pass=r["form"]) for r in rows(ctl_final)],
           "CTL_pick": [dict(r, all_pass=form_of(r["poem"])) for r in rows(ctl_pick)]}
    fin.update({f"{k}_pick": rows(f"{d}/pick.jsonl") for k, d in A.items()})
    S = {k: dict(TP.stats(R), **extra(R)) for k, R in list(cand.items()) + list(fin.items())}
    res = {"stats": S, "dead": {}, "reported_186": {}}
    for n, s in S.items():
        print(f"  {n:10s} n {s['n']:4d}, form {s['form']:.1f}%, stutters raw {s['stutters']:.3f} unattested {s['stutters_unattested']:.3f}, "
              f"repeats {s['repeats']:.2f}, distinct-2 {s['distinct2']:.3f}, final bits {s['final_bits']:.2f}, "
              f"final-word log-freq {s['final_word_logfreq']:.2f}, unseen {s['final_word_unseen']:.3f}")
    for k in A:
        res["dead"][k] = gate_fail(S["CTL"], S[k], stutter_gate=(k == "B"))
        res["reported_186"][k] = reported_186(S["CTL"], S[k])
        for m in ("stutters", "stutters_unattested"):
            if S[k][m] > S["CTL"][m] + 0.10: print(f"  {k}: {m} {S[k][m]:.3f} above control + 0.10 ({S['CTL'][m]:.3f})" + (" - reported, not a kill" if k == "A" or m == "stutters" else ""))
        print(f"  {k}: " + ("KILLED before judging: " + ", ".join(res["dead"][k]) if res["dead"][k] else "pre-judge gates PASS")
              + (f"; 186's other gates (reported): {', '.join(res['reported_186'][k])}" if res["reported_186"][k] else "; 186's other gates: all within"))
    os.makedirs(out, exist_ok=True); json.dump(res, open(f"{out}/GATES.json", "w"), indent=1)
    print("PASSING " + " ".join(k for k, d in res["dead"].items() if not d))


def halves(ids):   # = fixbase_pipe.halves (170's pools: topics split by sorted id, 25/24)
    ids = sorted(ids); h = (len(ids) + 1) // 2
    return [ids[:h], ids[h:]]


LABEL = {"A": "tl_a", "B": "tl_b"}


def pool_sources(fb, arms):
    """{label: {topic: record}}: 170's final and pick, and each arm's pick; every label covers the same 49 topics"""
    S = {"fixfinal": {r["id"]: r for r in rows(f"{fb}/final.jsonl")}, "fixpick": {r["id"]: r for r in rows(f"{fb}/pick.jsonl")}}
    for x in arms:
        k, d = x.split("=", 1)
        dead = json.load(open(f"{os.path.dirname(d.rstrip('/'))}/GATES.json")).get("dead", {}).get(k) \
            if os.path.exists(f"{os.path.dirname(d.rstrip('/'))}/GATES.json") else None
        if not os.path.exists(f"{d}/pick.jsonl") or dead:
            print(f"  arm {k}: killed by its gate ({dead or 'form, before Typhoon'}) - skipped"); continue
        S[LABEL[k]] = {r["id"]: r for r in rows(f"{d}/pick.jsonl")}
    ids = sorted(S["fixfinal"])
    for lab, R in S.items():
        assert sorted(R) == ids, f"{lab}: covers different topics"
        assert all(R[i]["prompt"] == S["fixfinal"][i]["prompt"] for i in ids), f"{lab}: prompts differ"
        assert all(len(R[i]["poem"].split("\n")) == 4 for i in ids), f"{lab}: poem without 4 lines"
    return S, ids


def pools(out, fb, *arms):
    """RAW text (no ๆ write-back), ALL 49 poems of every label (form-failing too; form enters the endpoint), topic halves"""
    S, ids = pool_sources(fb, arms)
    os.makedirs(out, exist_ok=True); n = 0
    for h, part in enumerate(halves(ids)):
        for lab, R in S.items():
            wrows(f"{out}/jp_{lab}_h{h}.jsonl", [{"id": f"{i}_s0", "prompt": R[i]["prompt"], "poem": R[i]["poem"]} for i in part]); n += len(part)
    same = sum(S["fixfinal"][i]["poem"] == S["fixpick"][i]["poem"] for i in ids)
    print(f"  pools: labels {list(S)}, {len(ids)} topics each, {n} system poems; fixpick == fixfinal (not refilled) for {same}")
    return list(S)


def decide(allthree):
    """pre-registered primary rule (B vs 170 final, all-three): CONTINUE if diff >= +5pp and the CI lower bound > 0;
    STOP if the CI upper bound < +5pp; otherwise GREY"""
    if allthree["diff"] >= 5 and allthree["ci"][0] > 0: return "CONTINUE"
    if allthree["ci"][1] < 5: return "STOP"
    return "GREY"


def score(out, fb, *arms, root_pools="runs"):
    import fixon_raw, fixbase_pipe
    S, ids = pool_sources(fb, arms)
    lab, anchors = fixon_raw.labels(POOL_PREFIX, 0, 0, root_pools)
    void = {d: v for d, v in anchors.items() if v < 0.70}
    form = {(l, i, 0): (bool(S[l][i]["form"]) if l == "fixfinal" else form_of(S[l][i]["poem"])) for l in S for i in ids}
    for kind in ("clean", "aligned"):
        assert set(lab[kind]) == set(form), f"{kind}: judged poems differ from the pipeline output"
    full = {k: {x: form[x] and lab[k][x] for x in form} for k in ("clean", "aligned")}
    full["all-three"] = {x: full["clean"][x] and full["aligned"][x] for x in form}
    full["form"] = form
    R = {"anchor_clean_share": anchors, "void_pools": sorted(void), "contrasts": {}}
    labs = [l for l in S if l.startswith("tl_")]
    C_ = [(f"{l} vs fixfinal", "fixfinal", l) for l in labs] + [(f"{l} vs fixpick", "fixpick", l) for l in labs] + \
         [("fixfinal vs fixpick (refill)", "fixpick", "fixfinal")] + ([("tl_b vs tl_a (stutter fix)", "tl_a", "tl_b")] if len(labs) == 2 else [])
    for name, a_, b_ in C_:
        R["contrasts"][name] = {}
        print(f"  {name}")
        for kind, L in full.items():
            a = {x[1:]: v for x, v in L.items() if x[0] == a_}; b = {x[1:]: v for x, v in L.items() if x[0] == b_}
            st = R["contrasts"][name][kind] = fixbase_pipe.paired_stats(a, b)
            print(f"    {kind:9s} {a_} {st['a']:5.1f}%  {b_} {st['b']:5.1f}%  diff {st['diff']:+.1f}pp ({st['up']} up / {st['down']} down, "
                  f"n={st['n']})  95% CI [{st['ci'][0]:+.1f}, {st['ci'][1]:+.1f}]")
    same = [i for i in ids if S["fixfinal"][i]["poem"] == S["fixpick"][i]["poem"]]
    agree = {k: sum(lab[k][("fixfinal", i, 0)] == lab[k][("fixpick", i, 0)] for i in same) for k in ("clean", "aligned")}
    R["repeat_agreement_identical_poems"] = {"n": len(same), **agree}
    print(f"  judge repeatability on {len(same)} identical fixfinal/fixpick poems: clean {agree['clean']}/{len(same)}, aligned {agree['aligned']}/{len(same)}")
    prim = "tl_b vs fixfinal"
    R["decision"] = decide(R["contrasts"][prim]["all-three"]) if prim in R["contrasts"] else "B not judged (killed at the pre-judge gates)"
    os.makedirs(out, exist_ok=True); json.dump(R, open(f"{out}/RESULT.json", "w"), indent=1, default=str)
    print(f"  anchors clean: {anchors}")
    if void:
        print(f"  VOID pools (anchors < 70% clean): {void} - re-judge before any decision"); sys.exit(3)
    print(f"  DECISION (pre-registered, primary = B vs 170 final, all-three): {R['decision']}")


# ---------------------------------------------------------------- tests (CPU, no model)

def _test():
    import tempfile, torch
    t = tempfile.mkdtemp()
    # 1. allowlist counting: within-word repeats never counted; cross-word same-sound pairs counted per word pair
    g = [[["นานา", 0, 2, "naa0"], ["นานา", 1, 2, "naa0"], ["นาน", 0, 1, "naan0"], ["นาน", 0, 1, "naan0"]],
         [["นาน", 0, 1, "naan0"], ["นาน", 0, 1, "naan0"], ["เลย", 0, 1, "l@@j0"], ["เลย", 0, 1, "l@@j0"]],
         [["ใจเย็น", 0, 2, "caj0"], ["ใจเย็น", 1, 2, "jen0"], ["เย็น", 0, 1, "jen0"], ["ดี", 0, 1, "dii0"]],
         [["นาน", 0, 1, "naan0"], ["นาน", 0, 1, "naan0"], ["ใจเย็น", 0, 2, "caj0"], ["ใจเย็น", 1, 2, "jen0"]],
         [["ใจเย็น", 0, 2, "caj0"], ["ใจเย็น", 1, 2, "jen0"], ["เย็น", 0, 1, "jen0"], ["กุ้ง", 0, 1, "kuN2"]]]
    c = same_sound_pairs(g)
    assert c == {("นาน", "นาน"): 3, ("เลย", "เลย"): 1, ("ใจเย็น", "เย็น"): 2}, c
    os.makedirs(f"{t}/d"); wrows(f"{t}/d/poems.jsonl", [{"grid": g}])
    build_allow(f"{t}/d", f"{t}/allow.json", 2)
    assert set(json.load(open(f"{t}/allow.json"))["pairs"]) == {"นาน|นาน", "ใจเย็น|เย็น"}
    # 2. the ban on a toy lexicon, without a model: StutterDecoder.ban_stutters on a hand-built state
    import twolayer_canvas as TC
    lex = {"words": ["นาน", "เลย", "ใจเย็น", "เย็น", "ดี", "นานา", "กุ้ง"],
           "sounds": {"นาน": ["naan0"], "เลย": ["l@@j0"], "ใจเย็น": ["caj0", "jen0"], "เย็น": ["jen0"], "ดี": ["dii0"],
                      "นานา": ["naa0", "naa0"], "กุ้ง": ["kuN2"]}, "count": {}}
    X = TC.Lex(lex); SD = make_decoder_class(); SD.ALLOW = load_allow(f"{t}/allow.json", X)
    D = SD.__new__(SD); D.X = X; D.lens = [8] * TC.L; D.committed = set(); D.forced = {}
    D.tok = torch.full((TC.T,), TC.MASK, dtype=torch.long); D.n_banned = 0
    tokof = lambda w, k=0: X.tok(w, k)
    def put(i, j, w, k=0): kk = D.idx(i, j); D.tok[kk] = tokof(w, k); D.committed.add(kk)
    put(0, 0, "เลย"); put(0, 2, "นาน"); put(0, 5, "ใจเย็น", 0); D.forced[D.idx(0, 6)] = tokof("ใจเย็น", 1)
    put(1, 3, "กุ้ง")
    lp = torch.zeros(TC.T, X.NT); D.ban_stutters(lp)
    ok = lambda i, j, w, k=0: bool(lp[D.idx(i, j), tokof(w, k)] > -1e8)
    assert not ok(0, 1, "เลย"), "เลย|เลย (unattested) must be banned"
    assert ok(0, 1, "นาน"), "(0,1) sits between เลย and นาน: นาน|นาน is attested"
    assert ok(0, 3, "นาน"), "นาน|นาน attested (right of นาน)"
    assert ok(0, 1, "ดี") and ok(0, 3, "ดี"), "different sound never banned"
    assert ok(0, 7, "เย็น"), "ใจเย็น|เย็น (left is a FORCED cell) attested >= 2 -> allowed"
    assert not ok(0, 7, "ใจเย็น", 1), "a stray 2nd syllable jen0 after the forced jen0 -> banned"
    assert ok(0, 4, "ดี")
    assert not ok(1, 2, "กุ้ง") and not ok(1, 4, "กุ้ง"), "กุ้ง|กุ้ง unattested, both sides"
    assert ok(1, 2, "เลย") and ok(1, 4, "นาน")
    # within one word: a cell left of a committed 2nd syllable may take the SAME word's 1st syllable (นานา)
    put(2, 5, "นานา", 1)
    lp = torch.zeros(TC.T, X.NT); D.ban_stutters(lp)
    assert ok(2, 4, "นานา", 0), "consecutive syllables of one word are never banned"
    assert not ok(2, 6, "นานา", 0), "a new นานา starting right after นานา's last syllable: same sound naa0 -> banned"
    assert not ok(2, 4, "นานา", 1), "(นานา,1) at cell 4 is not the word's continuation -> banned"
    # stutter_pairs on a finished state: เลย|เลย counted, นาน|นาน not
    E = SD.__new__(SD); E.X = X; E.lens = [3, 3, 3, 3]; E.tok = torch.full((TC.T,), TC.MASK, dtype=torch.long)
    for j, w in enumerate(["เลย", "เลย", "ดี"]): E.tok[E.idx(0, j)] = tokof(w)
    for j, w in enumerate(["นาน", "นาน", "ดี"]): E.tok[E.idx(1, j)] = tokof(w)
    E.tok[E.idx(2, 0)] = tokof("นานา", 0); E.tok[E.idx(2, 1)] = tokof("นานา", 1); E.tok[E.idx(2, 2)] = tokof("ดี")
    for j, w in enumerate(["ดี", "กุ้ง", "กุ้ง"]): E.tok[E.idx(3, j)] = tokof(w)
    assert stutter_pairs(E) == [(0, 1), (3, 2)], stutter_pairs(E)
    # 3. gates rule
    base = {"form": 99.0, "stutters": 0.06, "stutters_unattested": 0.02}
    assert gate_fail(base, {"form": 98.4, "stutters": 0.28, "stutters_unattested": 0.20}, False) == []
    assert gate_fail(base, {"form": 98.4, "stutters": 0.28, "stutters_unattested": 0.13}, True) == ["unattested stutters"]
    assert gate_fail(base, {"form": 98.4, "stutters": 0.28, "stutters_unattested": 0.11}, True) == []   # raw is reported only
    assert gate_fail(base, {"form": 97.9, "stutters": 0.10, "stutters_unattested": 0.0}, True) == ["form < 98%"]
    # text-level attested set: รักรัก x2 kept, เลยเลย x1 not; ๆ written out
    wrows(f"{t}/train.jsonl", [{"poem": "รักรักนานนาน\nรักๆเลยเลย"}])
    attested_syllables(f"{t}/train.jsonl", f"{t}/att.json", 2)
    at = frozenset(json.load(open(f"{t}/att.json"))["kept"])
    assert "รัก" in at and "เลย" not in at, at
    assert stutters_unattested("ฉันรักรักเธอเลยเลยนะ", at) == 1
    # 4. decision rule
    assert decide({"diff": 6, "ci": (0.5, 12)}) == "CONTINUE" and decide({"diff": 6, "ci": (-1, 12)}) == "GREY"
    assert decide({"diff": -3, "ci": (-12, 4.9)}) == "STOP" and decide({"diff": 2, "ci": (-8, 11)}) == "GREY"
    # 5. pool_sources skips a gate-killed arm and an arm without picks (form gate), never asserts on it
    fb = f"{t}/fb"; os.makedirs(fb); os.makedirs(f"{t}/J/A"); os.makedirs(f"{t}/J/B")
    rec = lambda i: {"id": i, "prompt": "p" + i, "poem": "ก\nข\nค\nง", "form": True}
    for f_ in (f"{fb}/final.jsonl", f"{fb}/pick.jsonl", f"{t}/J/A/pick.jsonl"): wrows(f_, [rec("T1"), rec("T2")])
    json.dump({"dead": {"A": ["unattested stutters"]}}, open(f"{t}/J/GATES.json", "w"))
    S_, _ = pool_sources(fb, [f"A={t}/J/A", f"B={t}/J/B"]); assert set(S_) == {"fixfinal", "fixpick"}, set(S_)
    json.dump({"dead": {"A": []}}, open(f"{t}/J/GATES.json", "w"))
    S_, _ = pool_sources(fb, [f"A={t}/J/A", f"B={t}/J/B"]); assert set(S_) == {"fixfinal", "fixpick", "tl_a"}, set(S_)
    wrows(f"{t}/J/A/cands.jsonl", [{"all_pass": True}] * 49 + [{"all_pass": False}]); assert formgate(f"{t}/J/A")
    wrows(f"{t}/J/B/cands.jsonl", [{"all_pass": True}] * 97 + [{"all_pass": False}] * 3); assert not formgate(f"{t}/J/B")
    print("twolayer_full: tests pass")


if __name__ == "__main__":
    cmd, args = sys.argv[1], sys.argv[2:]
    if cmd == "decode": decode_cli(args)
    elif cmd == "allow": build_allow(args[0], args[1], *(args[2:3] or [2]))
    elif cmd == "attested": attested_syllables(args[0], args[1], *(args[2:3] or [2]))
    elif cmd == "merge": merge(*args)
    elif cmd == "picks": picks(*args)
    elif cmd == "formgate": sys.exit(0 if formgate(*args) else 1)
    elif cmd == "gates": gates(*args)
    elif cmd == "pools": pools(*args)
    elif cmd == "score": score(*args)
    elif cmd == "test": _test()
    else: sys.exit(f"unknown command {cmd}")