src/scripts/canvas_bot.py

"""Canvas with a pretrained Thai body: WangchanBERTa (105M, bidirectional, 78 GB of Thai) as the
encoder over the 8 x 10 syllable grid, with the phoneme channels, grid coordinates and บาท ROLE
embeddings (frame / complete / develop / close) added at the input. Same joint decoding as
scripts/canvas_mlm.py: one rime family per rhyme group, tone masks at line ends, iterative refinement,
per-บาท topic steering. Trained as a masked LM on the human corpus (PyTorch, MPS).

Why: the small canvas had the craft (human-likeness 0.79) but no language knowledge (sense 1.6/5
by the judge). This keeps the joint decision and adds the knowledge.

Usage: canvas_wangchan.py train --epochs 3 --out runs/CANVAS_WC
       canvas_wangchan.py decode --model runs/CANVAS_WC --prompts data/prompts/test.jsonl --out O.jsonl [--topic-boost 3]
"""
import collections, os, re, sys, json, math, time, random, argparse, collections, hashlib
sys.path.insert(0, "."); sys.path.insert(0, "scripts")
import torch, torch.nn as nn, torch.nn.functional as F
from klon8 import prosody as P
import syl_align
from canvas_mlm import (C, TL, PAD, MASK, UNK, SEP, EVEN, CHECKED, STOP_SYL,
                        phon, syls_of, Vocab, clauses_of)

L = 4                                              # one บท
GROUPS = [[(1, -1), (2, 2)], [(2, -1), (3, -1), (4, 2)]]
TONE_AT = {2: CHECKED, 3: EVEN, 4: EVEN}
N_BOT = 1

WC = "airesearch/wangchanberta-base-att-spm-uncased"
DEV = torch.device(os.environ.get("KLON8_DEVICE") or ("mps" if torch.backends.mps.is_available() else "cpu"))   # KLON8_DEVICE=cpu tests beside a GPU job
ROLE_OF_LINE = [1, 2, 3, 4]                        # วรรค roles: สดับ / รับ / รอง / ส่ง (0 = topic)

# PRONOUNCED syllables per วรรค, measured over 12,000 บท with klon8.prosody.g2p_line — the count the
# user confirmed by hand on six lines (2026-09-22). The grid segmenter under-counts: it reads ถวิล as
# one syllable where it is ถ|วิล, and ธนา as one where it is ธ|นา, so an earlier version of this table
# was built from cell counts and was wrong (it put 10-13% of วรรค at 7 syllables; the truth is ~1%,
# and the real alternative to 8 is 9). A grid cell is therefore an approximation of a syllable, and
# the decoder verifies the real count with g2p_line rather than trusting the cell count.
WAK_LEN_DIST = {0: {7: 16, 8: 764, 9: 213, 10: 7},
                1: {7: 6, 8: 782, 9: 207, 10: 6},
                2: {7: 10, 8: 740, 9: 240, 10: 10},
                3: {7: 7, 8: 714, 9: 267, 10: 11}}


def pick_lens(mode, rng):
    """Target PRONOUNCED syllables per วรรค. mode: '8', 'sample', or a pattern like '8,8,9,8'."""
    if mode == "sample":
        out = []
        for i in range(L):
            d = WAK_LEN_DIST[i]; tot = sum(d.values()); r = rng.random() * tot
            for k, v in sorted(d.items()):
                r -= v
                if r <= 0: out.append(k); break
            else: out.append(8)
        return out
    if "," in mode:
        v = [int(x) for x in mode.split(",")]
        assert len(v) == L and all(7 <= x <= 10 for x in v), f"--wak-lens {mode} must be {L} values in 7..10"
        return v
    n = int(mode); assert 7 <= n <= 10, f"--wak-lens {mode} out of range"
    return [n] * L


def syl_counts(poem):
    """What the poem actually says, per วรรค. The authority for length; the grid is not."""
    return [len(P.g2p_line(x)) for x in poem.split("\n")]



def load_poems(path, vocab_min=3, seg="aligned"):
    """Four-line บท on a grid of SYLLABLES.

    seg='aligned' uses scripts/syl_align.py, which cuts the written form where the pronunciation
    says a syllable begins (ถ|วิล, พัน|ธ|นา): 96.7% of human วรรค come out with the pronounced
    count, against 81.5% for the old character-level segmenter. A บท is kept only when every one of
    its four วรรค matches its pronounced count exactly, so a cell in the training data IS a syllable.
    seg='grid' is the old behaviour, kept as the control.
    """
    seg_fn = (lambda t, lim=None: syl_align.syls_aligned(t, lim)) if seg == "aligned" else syls_of
    poems, kept, seen = [], 0, 0
    for l in open(path, encoding="utf-8"):
        if not l.strip(): continue
        r = json.loads(l)
        lines = r["poem"].split("\n")
        if len(lines) != L: continue
        seen += 1
        grid = [seg_fn(ln) for ln in lines]
        if not all(7 <= len(g) <= C for g in grid): continue
        # verify against the JOINED cells, not the source line: the model emits syllables
        # concatenated with no spaces, and rejoining can change how newmm reads the line, so this
        # is the form whose count has to be right.
        if seg == "aligned" and any(len(g) != len(P.g2p_line("".join(g))) for g in grid):
            continue                                    # the cell count must be the syllable count
        kept += 1
        poems.append({"title": seg_fn(r.get("title", "") or "", TL), "grid": grid, "board": r.get("board") or ""})
    print(f"load_poems({path}, seg={seg}): {kept}/{seen} บท kept", flush=True)
    return poems


LAYOUT_BUFS = ("line_id", "pos_id", "seg_id", "role_id")


def load_body(model, path, old_itos=None, V=None, allow_missing_expert=False):
    """Warm start from an 8-วรรค checkpoint. The layout buffers are a function of the grid and must
    never be imported (WC14 stores role_id for a 12-cell title, which silently overwrote WCPLAN v1's
    บาท roles). self.line is Embedding(L+1) and no longer matches, so it is skipped by shape too."""
    sd = torch.load(path, map_location="cpu")
    for k in LAYOUT_BUFS: sd.pop(k, None)
    own = model.state_dict()
    if old_itos and V is not None and len(old_itos) != len(V.itos):
        # vocabulary changed (re-segmentation): copy the rows for syllables present in both
        old = {t: i for i, t in enumerate(old_itos)}
        for key in ("tok.weight", "out.weight"):
            if key not in sd or key not in own: continue
            w, fresh = sd[key], own[key].clone()
            n = 0
            for j, t in enumerate(V.itos):
                i = old.get(t)
                if i is not None and i < w.shape[0]:
                    fresh[j] = w[i]; n += 1
            sd[key] = fresh
            print(f"  {key}: {n}/{len(V.itos)} rows carried over from the old vocabulary", flush=True)
    shape_skip = [k for k, v in sd.items() if k in own and own[k].shape != v.shape]
    for k in shape_skip: sd.pop(k)
    r = model.load_state_dict(sd, strict=False)
    assert not r.unexpected_keys, f"unexpected keys in {path}: {r.unexpected_keys}"
    kept = [k for k in r.missing_keys if k not in LAYOUT_BUFS and k not in shape_skip
            and not (allow_missing_expert and (k.startswith("expert") or k.startswith("diction") or k.startswith("word_head") or k.startswith("det_head")))]   # only a warm start may lack the expert / diction channel / word head
    assert not kept, f"weights missing from the warm start: {kept}"
    print(f"warm start from {path}: {len(sd)} tensors loaded, "
          f"layout buffers rebuilt, re-initialised for the new grid: {shape_skip}", flush=True)
    return model


ZONE_TAIL = 2          # line-final cells in the zone (review 2026-09-27: 4 covered 56% of the poem)
EXPERT_AFTER = 7       # the expert block runs after body layer index 7 (8th of 12), so later layers see it

# Diction conditioning (user idea 2026-09-27: "About my poetic words injection, why don't you change the
# structure of the network"). Each training poem gets a DICTION LEVEL = how many of its 12 rhyme-zone cells
# (rhyme_zone's definition) hold a strong poetic-diction syllable (runs/LOC/poetic_syl.json, log-odds >= 2.5).
# Buckets by that count, from the train distribution (0-1: 39%, 2: 26%, 3: 18%, 4+: 17% of 6,962 บท):
# level 1 = 0-1, 2 = 2, 3 = 3, 4 = 4+; level 0 = "none" (condition dropped: 10% of training rows, and the
# default at decode). A learned embedding per level is added to EVERY cell's input through a per-dimension
# gate that starts at ZERO, so step 0 is exactly the base model. Every cell, not just the rhyme zone: the
# words leading into a poetic rhyme word have to collocate with it, and a zone-only signal would recreate
# the post-hoc bonus's problem (a poetic rhyme word bolted onto a line that was not written for it).
DICTION_LEVELS = 4
DICTION_LEX_MIN = 2.5
_DICTION_STRONG = None


def diction_count(grid, lex_path="runs/LOC/poetic_syl.json"):
    """How many of a poem's rhyme-zone cells hold a strong poetic syllable."""
    global _DICTION_STRONG
    if _DICTION_STRONG is None:
        _DICTION_STRONG = {s for s, lo in json.load(open(lex_path, encoding="utf-8")).items() if lo >= DICTION_LEX_MIN}
    k = 0
    for i, ln in enumerate(grid):
        n = len(ln); cells = set(range(max(0, n - ZONE_TAIL), n))
        if i in (1, 3): cells |= {j for j in (2, 4) if j < n}
        k += sum(ln[j] in _DICTION_STRONG for j in cells)
    return k


def diction_levels(poems, min_board=200):
    """1..4 per poem, WITHIN its forum board (reviews 2026-09-28: raw counts track register - lament and romance
    boards are poetic, jokes and technical titles are not - so a raw level would teach "level 4 = sad poem"
    and pull topic). Level = quartile of the poem's percentile of its count (ties split by a fixed hash) among poems of the same
    board; boards with < min_board poems are pooled. Order-independent, so the trainer and any later check
    that pass the same poem set get the same levels."""
    import collections
    ks = [diction_count(p["grid"]) for p in poems]
    size = collections.Counter(p.get("board", "") for p in poems)
    grp = [p.get("board", "") if size[p.get("board", "")] >= min_board else "_pooled" for p in poems]
    dist = collections.defaultdict(collections.Counter)
    for g, k in zip(grp, ks): dist[g][k] += 1
    out = []
    for p_, g, k in zip(poems, grp, ks):
        d = dist[g]; n = sum(d.values())
        # ties (counts are small integers) are split by a fixed hash of the poem, so every board gets ~25%
        # per level; a mid-rank left board means 2.22-2.76 because boards differ in how many poems have 0
        u = int(hashlib.sha256("".join("".join(l) for l in p_["grid"]).encode()).hexdigest()[:8], 16) / 16 ** 8
        pct = (sum(c for v, c in d.items() if v < k) + u * d[k]) / n
        out.append(1 + min(int(pct * DICTION_LEVELS), DICTION_LEVELS - 1))
    return out


def rhyme_zone(tok):
    """[B, T] 1.0 on the rhyme zone: the last ZONE_TAIL filled cells of each วรรค (the line-final word and
    the syllable before it, where bad spans end) and the receiving rhyme cells (3rd/5th of วรรค 2 and 4).
    User idea (2026-09-27): give more network to where the poems fail (asymmetry)."""
    B, T = tok.shape
    z = torch.zeros(B, T, device=tok.device)
    for i in range(L):
        cells = tok[:, TL + i * C: TL + (i + 1) * C]
        n = (cells != PAD).sum(1)                                   # line length per sample
        col = torch.arange(C, device=tok.device)[None, :]
        z[:, TL + i * C: TL + (i + 1) * C] = ((col >= (n[:, None] - ZONE_TAIL)) & (col < n[:, None])).float()
        if i in (1, 3):
            for j in (2, 4): z[:, TL + i * C + j] = torch.maximum(z[:, TL + i * C + j], (n > j).float())
    return z


# Whole-word rhyme slot (component C1, director 2026-09-27). 49.8% of the judge's bad spans reach the line
# end and are built from dictionary pieces whose phrase breaks (คงเกยใจ, มานระทม): the rhyme syllable is
# appended as a monosyllable after context it does not fit (last newmm token multi-syllable in 27% of bad
# line-end lines vs 46% of human lines). A head reads the hidden states of each line's last WORD_MAX_K cells
# and predicts the line's ENDING WORD (1-3 syllables) from a lexicon of words that end human TRAIN lines plus
# poetic/archaic entries (scripts/build_rhyme_lex.py). Trained as an auxiliary loss; at decode (--word-rhyme)
# the line end is drawn as a WORD inside the chosen rhyme family and all its syllables are committed jointly.
WORD_MAX_K = 3


def ending_word_idx(grid_line, lexmap):
    """lexicon index of a line's ending word (its last newmm token when that is exactly the last 1-3 cells), else -1"""
    from pythainlp.tokenize import word_tokenize
    ws = [w for w in word_tokenize("".join(grid_line), engine="newmm") if w.strip()]
    if not ws: return -1
    for k in range(1, WORD_MAX_K + 1):
        if k <= len(grid_line) and "".join(grid_line[-k:]) == ws[-1]:
            return lexmap.get(tuple(grid_line[-k:]), -1)
    return -1


class Canvas(nn.Module):
    def __init__(self, V, wc_tok, wc_model, expert="none", diction=False, word_lex=None, detector=False):
        super().__init__()
        self.expert_mode = expert
        self.diction = bool(diction)
        self.diction_level = None            # decode-time level (int); None = channel off
        if self.diction:
            _rng = torch.get_rng_state()
            self.diction_emb = nn.Embedding(DICTION_LEVELS + 1, wc_model.config.hidden_size)
            # std 0.1: an embedding row then has norm ~2.8, the size of the input it joins (pre-LN cell norm 2.76),
            # so a grown gate of ~0.15 gives a ~15% signal; at std 0.02 it stayed at 0.3% after 100 steps (scale check)
            nn.init.normal_(self.diction_emb.weight, std=0.1, generator=torch.Generator().manual_seed(11))
            self.diction_gate = nn.Parameter(torch.zeros(wc_model.config.hidden_size))
            torch.set_rng_state(_rng)
        self.word_lex = word_lex             # list of syllable lists, or None (no word head)
        if word_lex:
            d_ = wc_model.config.hidden_size
            # own seed, global RNG restored: the head's init must not shift the body's dropout stream, so CTL and
            # the head arm train on the same masks, order AND dropout (review rule from the diction channel)
            _rng = torch.get_rng_state(); torch.manual_seed(13)
            self.word_head = nn.Linear(WORD_MAX_K * d_, len(word_lex))
            torch.set_rng_state(_rng)
            ids = torch.full((len(word_lex), WORD_MAX_K), PAD, dtype=torch.long)
            for w, sy in enumerate(word_lex):
                for q, x in enumerate(sy): ids[w, q] = V.enc(x)
            assert not (ids == UNK).any(), "a lexicon syllable is outside the vocabulary"
            self.register_buffer("word_ids", ids, persistent=False)
            self.register_buffer("word_len", torch.tensor([len(sy) for sy in word_lex]), persistent=False)
        self.detector = bool(detector)
        if self.detector:
            # [B] per-cell binary 'replaced' logit on the body's last hidden state; own seed, global RNG restored,
            # so the body's dropout stream is the control's
            _rng = torch.get_rng_state(); torch.manual_seed(17)
            self.det_head = nn.Linear(wc_model.config.hidden_size, 1)
            torch.set_rng_state(_rng)
        if expert != "none":
            # one extra transformer block inserted after body layer EXPERT_AFTER (mid-stack), initialised
            # from the next layer; gated per dimension from ZERO, so at step 0 the model is exactly the base
            import copy
            self.expert = copy.deepcopy(wc_model.encoder.layer[EXPERT_AFTER + 1])
            self.expert_gate = nn.Parameter(torch.zeros(wc_model.config.hidden_size))
        n, d = len(V.itos), wc_model.config.hidden_size
        self.body = wc_model.encoder                          # 12 pretrained bidirectional layers
        self.ln_in = wc_model.embeddings.LayerNorm
        base = wc_model.embeddings.word_embeddings.weight.detach()
        init = torch.zeros(n, d)
        mask_id = wc_tok.mask_token_id
        for i, s in enumerate(V.itos):
            if i == MASK:
                init[i] = base[mask_id]; continue
            if i in (PAD, UNK, SEP):
                init[i] = base[wc_tok.pad_token_id if i == PAD else wc_tok.unk_token_id]; continue
            ids = [t for t in wc_tok(s, add_special_tokens=False)["input_ids"] if wc_tok.convert_ids_to_tokens(t) != "▁"]
            init[i] = base[ids].mean(0) if ids else base[wc_tok.unk_token_id]
        self.tok = nn.Embedding(n, d); self.tok.weight.data.copy_(init)
        self.out = nn.Linear(d, n); self.out.weight = self.tok.weight        # tied
        self.rime = nn.Embedding(V.n_rime, d); self.tone = nn.Embedding(6, d)
        self.line = nn.Embedding(L + 1, d); self.pos = nn.Embedding(max(TL, C), d)
        self.seg = nn.Embedding(2, d); self.role = nn.Embedding(5, d)
        for e in (self.rime, self.tone, self.line, self.pos, self.seg, self.role):
            nn.init.normal_(e.weight, std=0.02)
        self.register_buffer("rime_of", torch.tensor(V.rime)); self.register_buffer("tone_of", torch.tensor(V.tone))
        T = TL + L * C
        self.register_buffer("line_id", persistent=False, tensor=torch.tensor([0] * TL + [1 + i for i in range(L) for _ in range(C)]))
        self.register_buffer("pos_id", persistent=False, tensor=torch.tensor(list(range(TL)) + [j for _ in range(L) for j in range(C)]))
        self.register_buffer("seg_id", persistent=False, tensor=torch.tensor([0] * TL + [1] * (L * C)))
        self.register_buffer("role_id", persistent=False, tensor=torch.tensor([0] * TL + [ROLE_OF_LINE[i] for i in range(L) for _ in range(C)]))

    def forward(self, tok, level=None, return_hidden=False):
        known = tok != MASK
        r = torch.where(known, self.rime_of[tok], torch.zeros_like(tok)); t = torch.where(known, self.tone_of[tok], torch.full_like(tok, 5))
        x = (self.tok(tok) + self.rime(r) + self.tone(t) + self.line(self.line_id)[None] + self.pos(self.pos_id)[None]
             + self.seg(self.seg_id)[None] + self.role(self.role_id)[None])
        if self.diction:
            # no level given = level 0, the "no level" embedding the trainer uses for its dropped rows
            lv = level if level is not None else (self.diction_level if self.diction_level is not None else 0)
            if not torch.is_tensor(lv): lv = torch.full((tok.shape[0],), int(lv), dtype=torch.long, device=tok.device)
            self._x_norm = x.detach().norm(dim=-1).mean()       # for the injection-scale log
            x = x + self.diction_gate * self.diction_emb(lv)[:, None, :]    # gate starts at 0: base model
        x = self.ln_in(x)
        if self.expert_mode == "none":
            h = self.body(x).last_hidden_state
        else:
            # the expert sits MID-STACK so the layers after it attend to what it wrote - at the last layer
            # ASYM and SYM would compute the same function on the zone cells (purpose review 2026-09-27)
            h = x
            for li, layer in enumerate(self.body.layer):
                h = layer(h); h = h[0] if isinstance(h, tuple) else h
                if li == EXPERT_AFTER:
                    e = self.expert(h)
                    e = e[0] if isinstance(e, tuple) else e     # transformers >= 5 returns a tensor
                    assert e.shape == h.shape, f"expert output {tuple(e.shape)} != hidden {tuple(h.shape)}"
                    zone = rhyme_zone(tok)[..., None] if self.expert_mode == "asym" else torch.ones_like(h[..., :1])
                    h = h + zone * self.expert_gate * (e - h)   # ASYM: rhyme zone only; SYM: every cell
        return (self.out(h), h) if return_hidden else self.out(h)

    def word_logits(self, h, tok):
        """[B, L, n_words]: the ending-word head on each line's last WORD_MAX_K cells (by line length)"""
        B = tok.shape[0]; feats = []
        for i in range(L):
            n = (tok[:, TL + i * C: TL + (i + 1) * C] != PAD).sum(1)                     # [B]
            pos = TL + i * C + (n[:, None] - WORD_MAX_K + torch.arange(WORD_MAX_K, device=tok.device)[None, :]).clamp(min=0)
            feats.append(torch.gather(h, 1, pos[..., None].expand(-1, -1, h.shape[-1])).reshape(B, -1))
        return self.word_head(torch.stack(feats, 1))

    def detect_logits(self, tok):
        """[B, T] per-cell 'replaced' logit (experiment B). The input is a COMPLETE poem: no MASK cell and no
        corruption mask ever enters - the detector reads only the tokens, in training and at decode alike."""
        assert not (tok == MASK).any(), "the detector reads complete poems (no MASK cells)"
        _, h = self.forward(tok, return_hidden=True)
        return self.det_head(h)[..., 0]


def build_tok(poems, V):
    T = TL + L * C
    out = torch.full((len(poems), T), PAD, dtype=torch.long)
    for n, p in enumerate(poems):
        for j, s in enumerate(p["title"]): out[n, j] = V.enc(s)
        for i, line in enumerate(p["grid"]):
            for j, s in enumerate(line): out[n, TL + i * C + j] = V.enc(s)
    return out


def mask_batch(y, seg_id, gen):
    B, T = y.shape
    ratio = torch.rand(B, 1, generator=gen) * 0.75 + 0.15
    u = torch.rand(B, T, generator=gen)
    m = (seg_id[None, :] == 1) & (y != PAD) & (u < ratio)
    return torch.where(m, torch.full_like(y, MASK), y), m


def word_loss(model, hid, x, m, wt, acc=False):
    """Auxiliary ending-word CE on lines whose gold ending word is in the lexicon AND whose ending cells are
    all masked in x (otherwise the head could read the answer). Returns (mean CE, n lines[, n top-1 correct])."""
    wlog = model.word_logits(hid, x)                              # [B, L, W]
    B = x.shape[0]; sel = []
    for i in range(L):
        n = (x[:, TL + i * C: TL + (i + 1) * C] != PAD).sum(1)
        g = wt[:, i]; k = model.word_len[g.clamp(min=0)]
        ok = g >= 0
        for q in range(WORD_MAX_K):   # every cell of the gold word must be masked
            col = TL + i * C + (n - 1 - q).clamp(min=0)
            ok = ok & ((q >= k) | m.gather(1, col[:, None])[:, 0])
        sel.append(ok)
    sel = torch.stack(sel, 1)                                     # [B, L]
    nsel = int(sel.sum())
    if not nsel: return (torch.zeros((), device=x.device), 0) + ((0,) if acc else ())
    ce = F.cross_entropy(wlog[sel], wt[sel])
    if acc: return ce.detach(), nsel, int((wlog[sel].argmax(-1) == wt[sel]).sum())
    return ce, nsel


def stutter_ul(logits, x, y, m, seg_id, pos_id):
    """Unlikelihood on predicting a masked cell as its neighbour's syllable (C5, 6.3% of bad spans are
    stutters) - skipped where the gold itself repeats its neighbour (ค่อยค่อย is legitimate Thai)."""
    p = F.softmax(logits.float(), -1)
    tot, cnt = logits.new_zeros((), dtype=torch.float32), 0
    for sh, edge in ((1, 0), (-1, C - 1)):
        nb = torch.roll(x, sh, dims=1)            # the neighbour as the model SEES it: masked neighbours are skipped
        ok = m & (seg_id[None, :] == 1) & (pos_id[None, :] != edge) & (nb >= 4) & (nb != y)
        if ok.any():
            pn = p[ok].gather(1, nb[ok][:, None])[:, 0]
            tot = tot + (-torch.log((1 - pn).clamp(min=1e-6))).sum(); cnt += int(ok.sum())
    return tot / max(cnt, 1)


# [A'] unlikelihood on flagged word pairs of the model's own poems
def load_ul(path, V):
    """Rows from scripts/ul_prep.py flags: {"title": [syl], "grid": [[syl] x 4], "pairs": [[line, a, b, [sides]]]}
    where a, b are the junction cells of a word pair unseen in the train table and sides lists the cells that
    may be masked (never a rhyme cell). Returns (UY [N, T] token ids, PAIRS: per poem a list of
    (pos_a, pos_b, [eligible positions])). Rows with a syllable outside the vocabulary are dropped."""
    ys, prs, n_drop = [], [], 0
    for l in open(path, encoding="utf-8"):
        r = json.loads(l)
        y = build_tok([{"title": r["title"], "grid": r["grid"]}], V)[0]
        if (y[TL:] == UNK).any(): n_drop += 1; continue
        ys.append(y)
        prs.append([(TL + i * C + a_, TL + i * C + b_, [TL + i * C + q for q in sd]) for i, a_, b_, sd in r["pairs"]])
    print(f"unlikelihood data {path}: {len(ys)} poems, {sum(map(len, prs))} flagged pairs, "
          f"{n_drop} dropped (out-of-vocabulary syllable)", flush=True)
    return torch.stack(ys), prs


def ul_mask(prs, T, gen):
    """For each flagged pair mask ONE eligible side, its partner kept VISIBLE, so the loss targets the pairing
    (this syllable next to that one), not the syllable on its own. Pairs are visited in a random order; a pair
    is skipped when its chosen side or its partner is already masked, or its chosen side is another pair's partner."""
    m = torch.zeros(len(prs), T, dtype=torch.bool)
    for n, pl in enumerate(prs):
        keep = set()                   # partners of masked sides: must stay visible
        for k in torch.randperm(len(pl), generator=gen).tolist():
            a_, b_, sd = pl[k]
            s_ = sd[int(torch.randint(0, len(sd), (1,), generator=gen))]
            partner = b_ if s_ == a_ else a_
            if m[n, s_] or m[n, partner] or s_ in keep: continue
            m[n, s_] = True; keep.add(partner)
    return m


def flagged_ul(model, uy, uf):
    """Welleck et al. 2020 token unlikelihood: -log(1 - p(generated syllable)) at each masked flagged cell;
    everything else (title, the pair partner, the rest of the poem) is visible context. Unflagged cells get no
    loss, so the model is never trained to imitate its own output."""
    x = torch.where(uf, torch.full_like(uy, MASK), uy)
    p = F.softmax(model(x)[uf].float(), -1).gather(1, uy[uf][:, None])[:, 0]
    return (-torch.log((1 - p).clamp(min=1e-6))).mean(), p.detach().mean()


# [B] replaced-cell detection on the shared body
def load_det_bank(path, V):
    """The in-context corruption bank from scripts/det_prep.py: rows {title, grid, lab, line, split}. Replacements
    were written by the generator itself with the rest of the poem visible (ELECTRA-style), on human AND model-written
    host poems. Returns train and valid tensors (x token ids, lab 0/1 per cell, line corrupted or -1)."""
    out = {"train": ([], [], []), "valid": ([], [], [])}; n_drop = 0
    for l in open(path, encoding="utf-8"):
        r = json.loads(l)
        x = build_tok([{"title": r["title"], "grid": r["grid"]}], V)[0]
        if (x[TL:] == UNK).any(): n_drop += 1; continue
        lab = torch.zeros_like(x, dtype=torch.float)
        for i, li in enumerate(r["lab"]):
            for j, v in enumerate(li): lab[TL + i * C + j] = float(v)
        X, LB, LN = out[r["split"]]; X.append(x); LB.append(lab); LN.append(r["line"])
    res = {k: (torch.stack(X), torch.stack(LB), torch.tensor(LN)) for k, (X, LB, LN) in out.items()}
    print(f"detection bank {path}: train {len(res['train'][0])}, valid {len(res['valid'][0])} poems, "
          f"{n_drop} dropped (out of vocabulary); replaced share (train) "
          f"{float(res['train'][1].sum() / (res['train'][0][:, TL:] != PAD).sum()):.3f}", flush=True)
    return res


def det_valid_mask(x):
    return (x != PAD) & (torch.arange(x.shape[1], device=x.device)[None, :] >= TL)


def det_line_hit(lg, x, valid, line):
    """share of corrupted rows whose highest mean-probability line is the corrupted one (the locator rule)"""
    pr = torch.sigmoid(lg.float()); hit = n = 0
    for b in range(x.shape[0]):
        if line[b] < 0: continue
        s = [float(pr[b, TL + i * C:TL + (i + 1) * C][valid[b, TL + i * C:TL + (i + 1) * C]].mean()) for i in range(L)]
        hit += int(max(range(L), key=lambda i: s[i]) == int(line[b])); n += 1
    return hit, n


# [ULKL] KL anchor to a frozen copy of the starting model
def kl_all(cur, ref, x, excl=None):
    """KL(p_ref || p_cur) averaged over ALL non-PAD cells of input x (title and poem, masked or not): the anchor
    holds every prediction near the frozen starting model, not only the masked MLM targets. excl: cells left out
    (on the unlikelihood batch, the flagged cells themselves - the anchor must not pull back exactly what the
    unlikelihood term pushes down; review 2026-09-28)."""
    v = x != PAD
    if excl is not None: v = v & ~excl
    lr, lc = F.log_softmax(ref[v].float(), -1), F.log_softmax(cur[v].float(), -1)
    return (lr.exp() * (lr - lc)).sum(-1).mean()


def tripwire_sample(HY, HP):
    """fixed masks on held-out poems: one flagged side per pair (partner visible) and, separately, the same number
    of unflagged, non-rhyme cells"""
    g = torch.Generator().manual_seed(31)
    fm = ul_mask(HP, HY.shape[1], g)
    um = torch.zeros_like(fm)
    for n in range(len(HY)):
        flagged = {q for a_, b_, _ in HP[n] for q in (a_, b_)}
        rc = set()
        for i in range(L):
            ln = int((HY[n, TL + i * C: TL + (i + 1) * C] != PAD).sum())
            rc |= {TL + i * C + j for j in ({ln - 1} | ({2, 4} if i in (1, 3) else set()))}
        cand = [t for t in range(TL, HY.shape[1]) if HY[n, t] >= 4 and t not in flagged and t not in rc]
        for q in torch.randperm(len(cand), generator=g)[:int(fm[n].sum())].tolist(): um[n, cand[q]] = True
    return HY, fm, um


@torch.no_grad()
def tripwire_p(model, HY, fm, um):
    was = model.training; model.eval(); out = []
    for msk in (fm, um):
        x = torch.where(msk, torch.full_like(HY, MASK), HY).to(DEV)
        p = F.softmax(model(x)[msk.to(DEV)].float(), -1).gather(1, HY.to(DEV)[msk.to(DEV)][:, None])[:, 0]
        out.append(float(p.mean()))
    if was: model.train()
    return out


def flagged_ul_full(model, uy, uf):
    """flagged_ul with the full logits returned (for the KL term on the same forward); same loss definition."""
    x = torch.where(uf, torch.full_like(uy, MASK), uy)
    lg = model(x)
    p = F.softmax(lg[uf].float(), -1).gather(1, uy[uf][:, None])[:, 0]
    return (-torch.log((1 - p).clamp(min=1e-6))).mean(), p.detach().mean(), x, lg


def train(a):
    a._old_itos = json.load(open(a.init_vocab, encoding="utf-8"))["itos"] if getattr(a, "init_vocab", None) else None
    from transformers import AutoTokenizer, AutoModel
    random.seed(0); torch.manual_seed(0)
    poems = load_poems(a.data, seg=getattr(a, "seg", "aligned")); random.shuffle(poems)
    if getattr(a, "vocab", None):
        # rebuild the SAME Vocab object from a frozen itos so every derived table
        # (rime, tone, fam, n_rime) is recomputed by the constructor, not left stale
        frozen = json.load(open(a.vocab, encoding="utf-8"))["itos"]
        V = Vocab(frozen[4:] * 3)          # min_freq is 3; itos[:4] are the specials
        assert V.itos == frozen, f"frozen vocab not reproduced: {len(V.itos)} vs {len(frozen)}"
        print(f"frozen vocab from {a.vocab}: {len(V.itos)} entries, {V.n_rime-1} rimes", flush=True)
    else:
        V = Vocab([s for p in poems for line in p["grid"] for s in line] + [s for p in poems for s in p["title"]])
    os.makedirs(a.out, exist_ok=True)
    # record HOW the grid was segmented, so decode can use the same segmenter for the topic
    word_w = getattr(a, "word_loss_w", 0.0) or 0.0
    det_w = getattr(a, "det_weight", 0.0) or 0.0
    word_lex = None
    if word_w > 0:
        # the lexicon travels with the model: decode reads it from the model directory
        lex = json.load(open(a.word_lex, encoding="utf-8"))
        word_lex = lex["words"]
        json.dump(lex, open(os.path.join(a.out, "rhyme_words.json"), "w", encoding="utf-8"), ensure_ascii=False)
    json.dump({"itos": V.itos, "seg": getattr(a, "seg", "aligned"), "expert": getattr(a, "expert", "none"),
               **({"diction": True} if getattr(a, "diction", False) else {}),
               **({"word_rhyme": "rhyme_words.json"} if word_lex else {}),
               **({"detector": True} if det_w else {})},
              open(os.path.join(a.out, "vocab.json"), "w", encoding="utf-8"), ensure_ascii=False)
    wc_tok = AutoTokenizer.from_pretrained(WC); wc = AutoModel.from_pretrained(WC)
    model = Canvas(V, wc_tok, wc, expert=getattr(a, "expert", "none"), diction=getattr(a, "diction", False), word_lex=word_lex, detector=det_w > 0)
    if getattr(a, "init", None):
        sd_keys = torch.load(a.init, map_location="cpu").keys()
        load_body(model, a.init, old_itos=getattr(a, "_old_itos", None), V=V, allow_missing_expert=True)
        if getattr(model, "expert_mode", "none") != "none" and not any(k.startswith("expert") for k in sd_keys):
            # start the new block from the FINE-TUNED last layer, not the raw pretrained one
            model.expert.load_state_dict(model.body.layer[EXPERT_AFTER + 1].state_dict())
            print(f"  expert block initialised from the warm-started body.layer[{EXPERT_AFTER + 1}]; gate = 0", flush=True)
    model = model.to(DEV)
    kl_b = getattr(a, "kl_beta", 0.0) or 0.0
    if kl_b:
        assert getattr(a, "init", None), "--kl-beta anchors to the --init model"
        assert not (model.diction or word_lex or det_w), "the KL anchor is tested only on the plain body"
        import copy
        ref_model = copy.deepcopy(model).eval()      # frozen: no grad, never stepped, dropout off
        for q in ref_model.parameters(): q.requires_grad_(False)
        ref_sum = lambda: sum(float(q.detach().cpu().double().sum()) + float((q.detach().cpu().double() ** 2).sum()) for q in ref_model.parameters())   # MPS has no float64: sum on CPU
        ref_chk = ref_sum()
        print(f"KL anchor: beta {kl_b} to a frozen copy of {a.init} on all non-PAD cells", flush=True)
    trip = None
    if getattr(a, "ul_heldout", None):
        # in-training tripwire (review 2026-09-28): mean p(flagged side) and p(unflagged non-rhyme cells) on a fixed
        # sample of title-held-out generated poems, every 200 steps; the arm aborts if either falls below 0.5 x its
        # value at step 0 (A' drove both to ~0). Own RNG, no dropout: CTL's streams are untouched.
        HY, HP = load_ul(a.ul_heldout, V)
        trip = tripwire_sample(HY[:64], HP[:64])
        trip0 = tripwire_p(model, *trip)
        print(f"tripwire start: p(flagged) {trip0[0]:.4f} p(unflagged) {trip0[1]:.4f}", flush=True)
    nv = min(getattr(a, "valid", 500) or 500, max(1, len(poems) // 5))
    valid, tr = poems[:nv], poems[nv:]
    assert len(tr) >= a.batch, f"only {len(tr)} training poems after a {nv}-poem valid split"
    Y, VY = build_tok(tr, V), build_tok(valid, V)
    DL = VDL = None
    if model.diction:
        LV = torch.tensor(diction_levels(poems)); VDL, DL = LV[:nv], LV[nv:]
        dgen = torch.Generator().manual_seed(7)     # own RNG: the level dropout must not shift CTL's masks/order
        print(f"diction levels (train, within board): " + ", ".join(f"{k}: {int((DL == k).sum())}" for k in range(1, DICTION_LEVELS + 1)), flush=True)
    WT = VWT = None
    if word_lex:
        lexmap = {tuple(w): k for k, w in enumerate(word_lex)}
        WT = torch.tensor([[ending_word_idx(g, lexmap) for g in p["grid"]] for p in tr])
        VWT = torch.tensor([[ending_word_idx(g, lexmap) for g in p["grid"]] for p in valid])
        print(f"word-rhyme head: {len(word_lex)} words, loss weight {word_w}; line endings in the lexicon: "
              f"train {float((WT >= 0).float().mean()):.3f}, valid {float((VWT >= 0).float().mean()):.3f}", flush=True)
    assert not (word_lex and model.diction), "the word head and the diction channel are separate experiments (not combined yet)"
    stut_w = getattr(a, "stutter_ul", 0.0) or 0.0
    seg = model.seg_id.cpu()
    new = {n for n, _ in model.named_parameters() if n.startswith("expert") or n.startswith("diction") or n.startswith("word_head") or n.startswith("det_head")}
    gates = {"expert_gate", "diction_gate"}
    groups = [{"params": [p for n, p in model.named_parameters() if n not in new], "lr": a.lr, "weight_decay": 0.01}]
    if new:   # new parameters learn faster than the converged body; the gate has no weight decay
        groups.append({"params": [p for n, p in model.named_parameters() if n in new and n not in gates],
                       "lr": getattr(a, "expert_lr", a.lr), "weight_decay": 0.01})
        groups.append({"params": [p for n, p in model.named_parameters() if n in gates],
                       "lr": getattr(a, "gate_lr", a.lr), "weight_decay": 0.0})
    opt = torch.optim.AdamW(groups)
    steps = a.epochs * (len(tr) // a.batch)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=[g["lr"] for g in groups], total_steps=steps, pct_start=0.06)
    gen = torch.Generator().manual_seed(0)
    ul_w = getattr(a, "ul_weight", 0.0) or 0.0
    if ul_w:
        # own RNG and a dropout-free extra forward: the MLM batches, masks, order and dropout draws are exactly
        # CTL's; weight 0 never reaches this branch (byte-identical to the old trainer)
        UY, UP = load_ul(a.ul_data, V)
        ugen = torch.Generator().manual_seed(11)
    if det_w:
        # [B] own RNG and a dropout-free extra forward, as A': the MLM batches, masks, order and dropout draws
        # are exactly CTL's; weight 0 never reaches this branch
        DB = load_det_bank(a.det_bank, V)
        DX, DLAB, _ = DB["train"]; VDX, VDLAB, VDLINE = DB["valid"]; VDVAL = det_valid_mask(VDX)
        detgen = torch.Generator().manual_seed(23)
    print(f"poems {len(poems)} (train {len(tr)}), vocab {len(V.itos)}, rime {V.n_rime-1}, params {sum(p.numel() for p in model.parameters())/1e6:.0f}M, device {DEV}", flush=True)
    t0 = time.time(); it = 0
    best_vl, bad_epochs = float("inf"), 0
    for ep in range(a.epochs):
        model.train(); perm = torch.randperm(len(tr), generator=gen)
        for i in range(0, len(tr) - a.batch + 1, a.batch):
            y = Y[perm[i:i + a.batch]]
            x, m = mask_batch(y, seg, gen)
            x, y, m = x.to(DEV), y.to(DEV), m.to(DEV)
            if DL is not None:
                # condition dropout: 10% of rows train with level 0 ("none"), so an unconditioned decode is
                # in distribution (drawn from gen only when the channel exists: the old path is unchanged)
                lv = DL[perm[i:i + a.batch]].clone()
                lv[torch.rand(lv.shape[0], generator=dgen) < 0.1] = 0
                logits = model(x, level=lv.to(DEV))
            elif WT is not None:
                logits, hid = model(x, return_hidden=True)
            else:
                logits = model(x)
            zw = getattr(a, "zone_weight", 1.0)
            if zw != 1.0:   # the same weighting in every arm, so arms differ only in architecture
                w = 1.0 + (zw - 1.0) * rhyme_zone(y)
                loss = (F.cross_entropy(logits[m], y[m], reduction="none") * w[m]).sum() / w[m].sum()
            else:
                loss = F.cross_entropy(logits[m], y[m])
            if WT is not None:
                _mlm = float(loss)
                wl, _wn = word_loss(model, hid, x, m, WT[perm[i:i + a.batch]].to(DEV))
                if _wn: loss = loss + word_w * wl
                if (it + 1) % 50 == 0: print(f"   it {it + 1} MLM {_mlm:.3f} ending-word CE {float(wl):.3f} (x{word_w})", flush=True)
            if stut_w:
                loss = loss + stut_w * stutter_ul(logits, x, y, m, model.seg_id, model.pos_id)
            if kl_b:
                with torch.no_grad(): ref_lg = ref_model(x)
                klm = kl_all(logits, ref_lg, x)
                loss = loss + kl_b * klm
            if ul_w:
                ub = torch.randint(0, len(UY), (getattr(a, "ul_batch", 0) or a.batch,), generator=ugen).tolist()
                uf = ul_mask([UP[k] for k in ub], UY.shape[1], ugen)
                model.eval()                  # no dropout: the global RNG stream (dropout) stays CTL's
                if kl_b:
                    ul, up, xu, lgu = flagged_ul_full(model, UY[ub].to(DEV), uf.to(DEV))
                    with torch.no_grad(): ref_u = ref_model(xu)
                    klu = kl_all(lgu, ref_u, xu, excl=uf.to(DEV))
                    loss = loss + kl_b * klu
                    if (it + 1) % 50 == 0: print(f"   it {it + 1} KL mlm {float(klm):.4f} ul-batch {float(klu):.4f} (x{kl_b})", flush=True)
                else:
                    ul, up = flagged_ul(model, UY[ub].to(DEV), uf.to(DEV))
                model.train()
                if (it + 1) % 50 == 0: print(f"   it {it + 1} unlikelihood {float(ul):.3f} (x{ul_w}), mean p(flagged) {float(up):.4f}", flush=True)
                loss = loss + ul_w * ul
            if det_w:
                db = torch.randint(0, len(DX), (getattr(a, "det_batch", 0) or a.batch,), generator=detgen)
                xd, ld = DX[db], DLAB[db]; vd = det_valid_mask(xd)
                model.eval()                  # no dropout: the global RNG stream (dropout) stays CTL's
                dlg = model.detect_logits(xd.to(DEV))
                dl = F.binary_cross_entropy_with_logits(dlg[vd.to(DEV)].float(), ld.to(DEV)[vd.to(DEV)])
                model.train()
                if (it + 1) % 50 == 0: print(f"   it {it + 1} detection BCE {float(dl.detach()):.3f} (x{det_w}), replaced share {float(ld[vd].mean()):.3f}", flush=True)
                loss = loss + det_w * dl
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step(); sched.step(); it += 1
            if trip is not None and it % 200 == 0:
                tp = tripwire_p(model, *trip)
                print(f"   it {it} tripwire p(flagged) {tp[0]:.4f} p(unflagged) {tp[1]:.4f} (start {trip0[0]:.4f} / {trip0[1]:.4f})", flush=True)
                if tp[0] < 0.5 * trip0[0] or tp[1] < 0.5 * trip0[1]:
                    print(f"== TRIPWIRE: held-out p fell below half its start value - arm aborted", flush=True)
                    raise SystemExit(5)
            if it % 200 == 0 and str(DEV) == "mps": torch.mps.empty_cache()
            if it % 50 == 0:
                print(f"ep {ep} it {it}/{steps} loss {loss.item():.3f} {time.time()-t0:.0f}s mps {(torch.mps.driver_allocated_memory() if str(DEV) == "mps" else 0)/1e9:.1f}GB", flush=True)
        model.eval()
        with torch.no_grad():  # chunked: one forward over all 500 valid poems held ~3 GB of logits in the MPS cache
            x, m = mask_batch(VY, seg, torch.Generator().manual_seed(1))
            tot, cnt, zt, zc = 0.0, 0, 0.0, 0
            wt_, wn_, wc_ = 0.0, 0, 0
            for i in range(0, len(VY), 50):
                xi, yi, mi = x[i:i + 50].to(DEV), VY[i:i + 50].to(DEV), m[i:i + 50].to(DEV)
                lvi = VDL[i:i + 50].to(DEV) if VDL is not None else None
                if VWT is not None:
                    lg, hid = model(xi, return_hidden=True)
                    wl, wn, wk = word_loss(model, hid, xi, mi, VWT[i:i + 50].to(DEV), acc=True)
                    wt_ += float(wl) * wn; wn_ += wn; wc_ += wk
                else:
                    lg = model(xi, level=lvi)
                ce = F.cross_entropy(lg[mi], yi[mi], reduction="none"); zm = rhyme_zone(yi)[mi] > 0
                tot += ce.sum().item(); cnt += int(mi.sum()); zt += ce[zm].sum().item(); zc += int(zm.sum())
            vl = tot / max(cnt, 1)
            print(f"   valid loss on the rhyme zone {zt / max(zc, 1):.3f}, elsewhere {(tot - zt) / max(cnt - zc, 1):.3f}", flush=True)
            if det_w:
                dh = dn = 0; dbce = 0.0
                for i in range(0, len(VDX), 50):
                    lg_ = model.detect_logits(VDX[i:i + 50].to(DEV)).cpu()
                    vv = VDVAL[i:i + 50]
                    dbce += float(F.binary_cross_entropy_with_logits(lg_[vv].float(), VDLAB[i:i + 50][vv], reduction="sum"))
                    h_, n_ = det_line_hit(lg_, VDX[i:i + 50], vv, VDLINE[i:i + 50]); dh += h_; dn += n_
                print(f"   valid detection BCE {dbce / max(int(VDVAL.sum()), 1):.4f}; corrupted line found {dh}/{dn} = {dh / max(dn, 1):.3f}", flush=True)
            if VWT is not None:
                print(f"   valid ending-word CE {wt_ / max(wn_, 1):.3f}, top-1 {wc_ / max(wn_, 1):.3f} on {wn_} masked line endings", flush=True)
        if str(DEV) == "mps": torch.mps.empty_cache()
        if getattr(model, "expert_mode", "none") != "none":
            g = model.expert_gate.detach().abs()
            print(f"   gate |g| mean {g.mean().item():.4f} max {g.max().item():.4f}", flush=True)
        if model.diction:
            g = model.diction_gate.detach().abs()
            inj = (model.diction_gate * model.diction_emb.weight).detach().norm(dim=-1)
            print(f"   diction gate |g| mean {g.mean().item():.4f} max {g.max().item():.4f}; injected norm by level "
                  f"{[round(v, 3) for v in inj.tolist()]} vs input norm {float(model._x_norm):.3f} "
                  f"(ratio max {float(inj.max() / model._x_norm):.3f})", flush=True)
        if kl_b:
            assert ref_sum() == ref_chk, "the KL reference model changed during training"
            print(f"   KL reference unchanged (checksum {ref_chk:.6e})", flush=True)
        print(f"== epoch {ep} valid masked-loss {vl:.3f}" + ("  (best)" if vl < best_vl else ""), flush=True)
        if vl < best_vl - a.min_delta:
            best_vl, bad_epochs = vl, 0
            torch.save(model.state_dict(), os.path.join(a.out, "model.pt"))
        else:
            bad_epochs += 1
            if bad_epochs >= a.patience:
                print(f"== early stop at epoch {ep}: no improvement for {a.patience} epochs (best {best_vl:.3f})", flush=True); break


def load_model(path):
    from transformers import AutoTokenizer, AutoModel
    V = Vocab.__new__(Vocab)
    meta = json.load(open(os.path.join(path, "vocab.json"), encoding="utf-8"))
    V.itos = meta["itos"]
    V.seg = meta.get("seg", "grid")          # checkpoints written before this field used the old grid
    V.stoi = {t: i for i, t in enumerate(V.itos)}
    rimes = {}; V.rime, V.tone, V.fam = [0] * len(V.itos), [5] * len(V.itos), {}
    for i, t in enumerate(V.itos):
        if i < 4: continue
        r, tn = phon(t)
        if r is None: continue
        rimes.setdefault(r, len(rimes) + 1); V.rime[i], V.tone[i] = rimes[r], tn
        V.fam.setdefault(rimes[r], []).append(i)
    V.n_rime = len(rimes) + 1
    wc_tok = AutoTokenizer.from_pretrained(WC); wc = AutoModel.from_pretrained(WC)
    word_lex = json.load(open(os.path.join(path, meta["word_rhyme"]), encoding="utf-8"))["words"] if meta.get("word_rhyme") else None
    model = Canvas(V, wc_tok, wc, expert=meta.get("expert", "none"), diction=meta.get("diction", False), word_lex=word_lex,
                   detector=meta.get("detector", False))
    load_body(model, os.path.join(path, "model.pt"))
    return model.to(DEV).eval(), V


class CleanPrior:
    """Decode-time cleanness constraints from corpus syllable statistics (data/syl_ngrams.json).
    rhyme_lex: rhyme cells may only take syllables the corpus uses at that position >= N times.
    bigram_w:  open cells get w * log P(cell | committed left neighbour) + w * log P(committed right neighbour | cell),
               add-alpha smoothed adjacent-syllable bigrams within a วรรค."""
    def __init__(self, V, path="data/syl_ngrams.json", rhyme_lex=0, bigram_w=0.0, alpha=0.1):
        d = json.load(open(path, encoding="utf-8")); n = len(V.itos); self.n = n; self.w = bigram_w; self.alpha = alpha
        self.uni = torch.zeros(n)
        for s, c in d["uni"].items():
            i = V.enc(s)
            if i != UNK: self.uni[i] = c
        self.rows = collections.defaultdict(dict); self.cols = collections.defaultdict(dict)
        for k, c in d["bi"].items():
            a, b = k.split("\t"); ia, ib = V.enc(a), V.enc(b)
            if ia == UNK or ib == UNK: continue
            self.rows[ia][ib] = c; self.cols[ib][ia] = c
        self.final_mask = self.x_mask = None
        if rhyme_lex:
            self.final_mask = torch.full((n,), -1e9); self.x_mask = torch.full((n,), -1e9)
            for s, c in d["final"].items():
                if c >= rhyme_lex and V.enc(s) != UNK: self.final_mask[V.enc(s)] = 0.0
            for s, c in d["x"].items():
                if c >= rhyme_lex and V.enc(s) != UNK: self.x_mask[V.enc(s)] = 0.0
        self.denom = self.uni + alpha * n
        self._cache_l, self._cache_r = {}, {}
    def left(self, a):  # vector over b: log P(b | a)
        if a not in self._cache_l:
            v = torch.full((self.n,), math.log(self.alpha / float(self.denom[a])))
            for b, c in self.rows[a].items(): v[b] = math.log((c + self.alpha) / float(self.denom[a]))
            self._cache_l[a] = v
        return self._cache_l[a]
    def right(self, b):  # vector over x: log P(b | x)
        if b not in self._cache_r:
            v = torch.log(self.alpha / self.denom)
            for x, c in self.cols[b].items(): v[x] = math.log((c + self.alpha) / float(self.denom[x]))
            self._cache_r[b] = v
        return self._cache_r[b]



JUNK_CHARS = ".*๐"          # crawl artefacts that became vocabulary: '.กลอน', '..บท', '*คง.', '.....๐'


def junk_ids(V):
    """Vocabulary entries carrying corpus punctuation padding (725 of 8561); 2.71% of corpus วรรค and
    8 of 110 generated poems contain one, and one bad วรรค disqualifies a poem."""
    return [i for i, t in enumerate(V.itos) if i >= 4 and any(c in t for c in JUNK_CHARS)]


@torch.no_grad()
def rarity_gain(sy, freq, p_ref=1e-3, lo=1.0, hi=3.0):
    """How much harder a RARE topic syllable has to be pushed to compete.

    The boost is added to logits, so a flat constant is worth very different things to syllables
    whose priors span three orders of magnitude: one at 1e-5 needs about log(1e-3/1e-5) = 4.6 nats
    to reach the same odds as one at 1e-3, and a flat 1.5 does not get it there. Measured on the
    220 judged benchmark poems: the topic gate passes 4% when the topic's rarest content syllable
    is below 1e-5, 44% at 1e-5-1e-4, 70% at 1e-4-1e-3 and 83% above, while CLEANNESS is flat across
    the same split (54/41/45/55). Rarity breaks the topic gate specifically.
    """
    if freq is None: return 1.0
    p = freq["count"].get(sy, 0) / freq["total"]
    if p <= 0: return hi
    return min(hi, max(lo, math.log(p_ref / p)))


def decode_one(model, V, title_syls, T=12, temp=0.8, lens=None, fam_temp=0.5, topic=None, topic_boost=0.0, inner_boost=0.0, clean=None, fam_distinct=False, fam_size_alpha=0.0, fam_prior=None, fam_prior_w=0.0, ban_junk=False, seg=syls_of, rhyme_first=True, commit_n=0, rhyme_pos="3", fam_soft=0.0, fam_delay=0, topic_zero="global", topic_freq=None, prefill=None, orig=None, rhyme_last=False, trace=None, seam_repair=False, anaphora_penalty=0.0, poetic=None, word_rhyme=False, word_mlm_w=1.0, commit_log=None):
    """prefill: {line index: [token ids]} - those วรรค are committed before step 0 and never change
    (the refill mode). orig: the original poem's token lists for all four lines, used only to read
    which receiving position (3 or 5) its kept rhyme links used."""
    n = len(V.itos); lens = lens or [8] * L
    if rhyme_last: rhyme_first = False     # the two orders are opposites
    tok = torch.full((TL + L * C,), PAD, dtype=torch.long)
    for j, s in enumerate(title_syls[:TL]): tok[j] = V.enc(s)
    idx = {(i, j): TL + i * C + j for i in range(L) for j in range(C)}
    cells = [(i, j) for i in range(L) for j in range(lens[i])]
    for c in cells: tok[idx[c]] = MASK
    committed = set()
    for i, ids in (prefill or {}).items():
        assert len(ids) == lens[i], f"prefilled line {i} has {len(ids)} tokens for {lens[i]} cells"
        for j, t_ in enumerate(ids): tok[idx[(i, j)]] = t_; committed.add(idx[(i, j)])
    base_bad = torch.zeros(n); base_bad[:4] = -1e9
    if ban_junk:
        ids = junk_ids(V)
        assert ids, "--ban-junk but no vocabulary entry matches JUNK_CHARS"
        base_bad[torch.tensor(ids)] = -1e9
    tone_mask = {ln: torch.tensor([0.0 if (V.tone[k] in allowed and k >= 4) else -1e9 for k in range(n)]) for ln, allowed in TONE_AT.items()}
    fam_vec = {}
    for f, ks in V.fam.items():
        v = torch.full((n,), -1e9); v[torch.tensor(ks)] = 0.0; fam_vec[f] = v
    # The receiving cell of a rhyme link may sit at syllable 3 OR 5 (klon8.prosody.STRICT_POS);
    # human บท use 5 in 10.3% of สดับ->รับ and 12.4% of รับ->ส่ง links. GROUPS names syllable 3, and
    # the decoder only ever tried that one. With rhyme_pos="auto" the position is chosen per group
    # by which cell the model puts more family mass on, doubling the candidate mass at the most
    # constrained cell in the poem for nothing.
    def resolve(g, alt):
        out = []
        for (l, c) in g:
            if c == -1: out.append((l - 1, lens[l - 1] - 1))
            else: out.append((l - 1, 4 if (alt and 4 < lens[l - 1]) else c))
        return out
    groups = [resolve(g, False) for g in GROUPS]
    groups_alt = [resolve(g, True) for g in GROUPS]
    group_family, group_pos = {}, {}
    free_pos = set()                 # prefilled groups whose receiving cell sits in a refilled line
    if prefill:
        # A group touching a kept line inherits its family NOW, before step 0: the kept rhyme word is
        # fixed, so the refilled line must meet it. Left to the normal path, fam_delay would leave the
        # rhyme cells unmasked and rhyme_first could commit one outside the family at t=0.
        for gi, g in enumerate(GROUPS):
            recv_line = g[-1][0] - 1
            if recv_line in prefill:
                # the receiver is kept: use the position (3 or 5) where the ORIGINAL poem rhymed
                # (resolved on the ORIGINAL line lengths: a refilled line may be retried at another length)
                def shares(alt):
                    rs = {V.rime[orig[l - 1][-1] if c == -1 else orig[l - 1][4 if (alt and 4 < len(orig[l - 1])) else c]]
                          for (l, c) in g}
                    return len(rs) == 1 and 0 not in rs
                group_pos[gi] = 0 if shares(False) or not shares(True) else 1
            elif rhyme_pos == "auto":
                free_pos.add(gi)
            kept = [c for c in (groups_alt if group_pos.get(gi) else groups)[gi] if c[0] in prefill]
            if kept and V.rime[int(tok[idx[kept[0]]])]:
                group_family[gi] = V.rime[int(tok[idx[kept[0]]])]
    topic_vec = None
    if topic and topic_boost:
        rows_ = []
        # one clause per วรรค. The 8-วรรค grid appended each clause twice (one per บาท); keeping
        # that here dropped clauses 3 and 4 entirely and let clauses 1-2 steer two วรรค each.
        for clause in (list(topic) + [topic[-1]] * L)[:L]:   # short prompts: the last clause repeats
            v = torch.zeros(n)
            for sy in seg(clause):
                if sy not in STOP_SYL and V.enc(sy) != UNK and len(sy) > 1:
                    v[V.enc(sy)] = topic_boost * rarity_gain(sy, topic_freq)
            rows_ += [v]
        while len(rows_) < L: rows_.append(torch.zeros(n))
        topic_vec = torch.stack(rows_[:L])
    big_fams = [(f, torch.tensor(ks)) for f, ks in V.fam.items() if len(ks) >= 20]
    # สัมผัสใน steering: an open cell whose template partner (P.VOWEL_SLOTS, 1-indexed within the line)
    # is committed gets a bonus on every syllable sharing that partner's rime, except the identical syllable.
    fam_ind = {f: (v == 0).float() for f, v in fam_vec.items()}
    inner_pairs = {}
    if inner_boost:
        for i in range(L):
            for (p, q), w in P.VOWEL_SLOTS.items():
                if q <= lens[i]:
                    inner_pairs.setdefault(idx[(i, p - 1)], []).append((idx[(i, q - 1)], w))
                    inner_pairs.setdefault(idx[(i, q - 1)], []).append((idx[(i, p - 1)], w))
    cell_of = {v: c for c, v in idx.items()}
    repaired = set()
    # C1 word-rhyme: each open line's last WORD_MAX_K cells (except a receiving 3rd/5th cell) are HELD until the
    # line's ending is drawn as a whole word (all its syllables committed jointly), so no partial ending commits
    # first and the head sees the fully-masked ending it was trained on (blind review 2026-09-28). The word is
    # drawn only once the rest of its line exists (at most 2 of its other cells open) - never forced early
    # (purpose review: that recreated the old rhyme-first problem). A line with no admissible word is released to the normal path.
    assert not (word_rhyme and rhyme_last), "--word-rhyme and --rhyme-last are different orders"
    word_pending = {i for i in range(L) if idx[(i, lens[i] - 1)] not in committed} if word_rhyme else set()
    line_group = {}
    for gi, g in enumerate(GROUPS):
        for (l_, c_) in g:
            if c_ == -1: line_group[l_ - 1] = gi
    recv_cols = {i: {c[1] for g_ in (resolve(g, False) for g in GROUPS) for c in g_ if c[0] == i and c[1] != lens[i] - 1} |
                    {c[1] for g_ in (resolve(g, True) for g in GROUPS) for c in g_ if c[0] == i and c[1] != lens[i] - 1}
                 for i in range(L)}
    word_hold = {i: {idx[(i, j)] for j in range(max(0, lens[i] - WORD_MAX_K), lens[i]) if j not in recv_cols[i]}
                 for i in word_pending}
    for i in list(word_pending):
        if any(k in committed for k in word_hold[i]): word_pending.discard(i)   # refill: an ending already kept

    def rhyme_hold():
        """Rhyme cells still held: the resolved position of a group whose position is known, BOTH
        candidate positions (3 and 5) of a receiver whose position is not yet chosen."""
        h = set()
        for gi in range(len(groups)):
            cg = (groups_alt if group_pos.get(gi) else groups)[gi]
            if gi not in group_pos and rhyme_pos == "auto": cg = list(set(cg) | set(groups_alt[gi]))
            h |= {idx[c] for c in cg}
        h.add(idx[(L - 1, lens[L - 1] - 1)])   # the ส่ง line-final cell: no rhyme link, but a tone class,
        return h                               # and line-final spans were counted for all four วรรค

    def eligible(c):
        held = rhyme_hold()
        return all(idx[(c[0], j)] in committed or idx[(c[0], j)] in held for j in range(lens[c[0]]))

    for t in range(T):
        if word_pending:
            _out, _hid = model(tok[None].to(DEV), return_hidden=True)
            wlp = F.log_softmax(model.word_logits(_hid, tok[None].to(DEV))[0].float().cpu(), -1)   # [L, W]
            logits = _out[0].float().cpu() + base_bad[None, :]
            # canvas term on ONE scale for every cell of a candidate word: the model's own log-probs before the
            # family / tone masks (masked log-probs renormalise the final cell inside the family and made a
            # one-syllable word look better than a multi-syllable one - purpose review 2026-09-28)
            lp_raw = F.log_softmax(logits, dim=-1)
        else:
            logits = model(tok[None].to(DEV))[0].float().cpu() + base_bad[None, :]
        for ln, tm in tone_mask.items():
            logits[idx[(ln - 1, lens[ln - 1] - 1)]] += tm
        if topic_vec is not None:
            tv = topic_vec.clone()
            if committed:
                if topic_zero == "row":
                    # zero a topic syllable ONLY in the วรรค where it was committed. The original
                    # zeroed it in every row, so each topic syllable steered until its first
                    # commitment ANYWHERE - and a short title yields 2-3 usable syllables after the
                    # stopword/UNK filter, so the topic channel died within ~3 of ~33 commits,
                    # typically inside วรรค 1. Measured signature on 195 judged candidates: topic
                    # present in 67% of วรรค1 and 17-19% of วรรค2-4 (human 45% / 5-15%).
                    for k in committed:
                        if int(tok[k]) >= 4: tv[(k - TL) // C, int(tok[k])] = 0.0
                else:
                    tv[:, torch.tensor(sorted({int(tok[k]) for k in committed if tok[k] >= 4}))] = 0.0
            for i in range(L):
                for j in range(lens[i]): logits[idx[(i, j)]] += tv[i]
        if inner_boost:
            for k, partners in inner_pairs.items():
                if k in committed: continue
                for pk, w in partners:
                    if pk not in committed: continue
                    pt = int(tok[pk]); f = V.rime[pt]
                    if not f or f not in fam_ind: continue
                    b = fam_ind[f] * (inner_boost * w); b[pt] = 0.0
                    logits[k] += b
        if clean is not None:
            for (i, j) in cells:
                k = idx[(i, j)]
                if k in committed: continue
                if clean.final_mask is not None:
                    if j == lens[i] - 1: logits[k] += clean.final_mask
                    elif i % 2 == 1 and j in (2, 4): logits[k] += clean.x_mask
                if clean.w:
                    if j > 0 and idx[(i, j - 1)] in committed and int(tok[idx[(i, j - 1)]]) >= 4:
                        logits[k] += clean.w * clean.left(int(tok[idx[(i, j - 1)]]))
                    if j + 1 < lens[i] and idx[(i, j + 1)] in committed and int(tok[idx[(i, j + 1)]]) >= 4:
                        logits[k] += clean.w * clean.right(int(tok[idx[(i, j + 1)]]))
        lp = F.log_softmax(logits, dim=-1)
        for gi in list(free_pos):
            # inherited family, refilled receiver: choose 3 or 5 by the model's family mass, as the
            # normal path does when it picks the family
            ks = torch.tensor(V.fam[group_family[gi]]) if gi in group_family else None
            if ks is not None:
                def mass(cg): return sum(torch.logsumexp(lp[idx[c]][ks], 0).item() for c in cg if idx[c] not in committed)
                group_pos[gi] = 1 if mass(groups_alt[gi]) > mass(groups[gi]) else 0
            free_pos.discard(gi)
        for gi, cells_g in enumerate(groups):
            if gi in group_family: continue
            # The family used to be fixed at t=0, off a blank canvas, where the model holds only
            # ~2.9% of its mass on the family the human poet used and ranks the poet's own rhyme word
            # third. With the rest of the บท visible it ranks that word FIRST. fam_delay holds the
            # choice back until some content has committed, so the rhyme is chosen against a line
            # that exists. Until then the rhyme cells are simply left open.
            if rhyme_last:
                # choose a group's family only once one of its cells may commit, i.e. once the rest of
                # THAT cell's line exists (see rhyme_hold below). The group's other lines may still be
                # open (group 1 spans วรรค 2-4); they then have to meet a family chosen from one line.
                if not any(eligible(c) for c in set(cells_g) | set(groups_alt[gi])): continue
            elif t < fam_delay: continue
            done = [idx[c] for c in cells_g if idx[c] in committed]
            if done:
                f = V.rime[int(tok[done[0]])]
                if f: group_family[gi] = f
                continue
            used_f = collections.Counter(group_family.values())
            scores, fams = [], []
            for f, ks in big_fams:
                # per-cell mean of the family log-mass; size correction (alpha=1 -> mean log-prob per member);
                # human-frequency prior; reuse across the poem's 4 groups either penalised or forbidden
                sc = sum(torch.logsumexp(lp[idx[c]][ks], 0).item() - fam_size_alpha * math.log(len(ks)) for c in cells_g) / len(cells_g)
                if fam_prior is not None and fam_prior_w: sc += fam_prior_w * fam_prior.get(f, math.log(1e-4))
                if used_f.get(f, 0): sc += -1e9 if fam_distinct else -1.0
                scores.append(sc); fams.append(f)
            if fams:
                z = torch.tensor(scores) / max(fam_temp, 1e-6)
                group_family[gi] = fams[int(torch.multinomial(F.softmax(z, 0), 1))] if fam_temp > 0 else fams[int(torch.argmax(z))]
                if rhyme_pos == "auto" and gi not in group_pos:
                    ks = dict(big_fams).get(group_family[gi])
                    if ks is not None:
                        def mass(cg): return sum(torch.logsumexp(lp[idx[c]][ks], 0).item() for c in cg) / len(cg)
                        group_pos[gi] = 1 if mass(groups_alt[gi]) > mass(groups[gi]) else 0
        for gi in range(len(groups)):
            cells_g = (groups_alt if group_pos.get(gi) else groups)[gi]
            f = group_family.get(gi)
            if f in fam_vec:
                # A HARD family mask (-1e9) forces the rhyme cell into whatever the chosen family
                # allows, however badly it reads - and the family itself was picked at t=0 off a
                # blank canvas. --fam-soft replaces the wall with a finite penalty, so a much better
                # word can win and the checker (plus retries) decides whether the poem still scans.
                pen = fam_vec[f] if not fam_soft else (fam_vec[f] / 1e9) * fam_soft
                for cell in cells_g:
                    if idx[cell] not in committed: logits[idx[cell]] += pen
        for gi in range(len(groups)):
            cells_g = (groups_alt if group_pos.get(gi) else groups)[gi]
            for cell in cells_g:
                if idx[cell] in committed: continue
                for other in cells_g:
                    if other != cell and idx[other] in committed: logits[idx[cell], int(tok[idx[other]])] = -1e9
        if committed:
            rep = torch.tensor(sorted({int(tok[k]) for k in committed if tok[k] >= 4}))
            logits[:, rep] -= 2.0
        pbon = None
        # In refill (prefill given) every rhyme family is inherited from the kept lines before step 0, so
        # fam_delay's reason does not apply - and rhyme-first commits a refilled line-1/3 rhyme cell AT step 0,
        # which never got the bonus (reviews 2026-09-27: 60/150 targeted refills). Live from step 0 there.
        if poetic is not None and (t >= fam_delay or prefill):
            # poetic-diction steering (plan B, user idea 2026-09-27): the judge forgives imagery and poetic
            # diction, and 80% of dirty poems fail at the rhyme positions. Uncommitted rhyme-zone cells get a
            # bonus toward whole-word syllables far more common in poems than in prose (runs/LOC/poetic_syl.json).
            # Reviews 2026-09-27: (1) not before fam_delay - until then rhyme-first commits the line-1 rhyme
            # cell with NO family mask on a blank canvas, and the bonus would pick group 0's family; (2) a
            # lexicon syllable already committed anywhere loses its bonus (as --topic-zero row does), so +W
            # cannot outweigh the -2 repeat penalty; (3) the 3rd/5th receiving cell is boosted only once
            # its group's position is resolved, and only the chosen one; (4) the bonus is kept out of the
            # commit-order confidence below, so it changes WHICH word, not WHEN the cell commits.
            pv = poetic.clone()
            used = sorted({int(tok[k]) for k in committed if tok[k] >= 4})
            if used: pv[torch.tensor(used)] = 0.0
            recv = {idx[c] for gi in group_family for c in (groups_alt if group_pos.get(gi) else groups)[gi]}
            pbon = torch.zeros_like(logits)
            for i in range(L):
                for j in range(lens[i]):
                    k = idx[(i, j)]
                    if k in committed: continue
                    if j >= lens[i] - ZONE_TAIL or k in recv: pbon[k] = pv
            logits = logits + pbon
        if anaphora_penalty:
            # template breaker (2026-09-27): 32% of our poems open 3+ lines with the same two syllables
            # against 1% of human poems, and half of those are judged unclean. A syllable that two OTHER
            # lines already have at the same opening position (1st or 2nd) is penalised here.
            for i in range(L):
                for j in (0, 1):
                    k = idx[(i, j)]
                    if k in committed or j >= lens[i]: continue
                    seen = collections.Counter(int(tok[idx[(l, j)]]) for l in range(L)
                                               if l != i and j < lens[l] and idx[(l, j)] in committed)
                    for t_, c_ in seen.items():
                        if c_ >= 2 and t_ >= 4: logits[k, t_] -= anaphora_penalty
        word_step = False
        if word_pending and (t >= fam_delay or prefill):
            lpc = F.log_softmax(logits, dim=-1)                                  # masked: admissibility only
            wid, wlen = model.word_ids.cpu(), model.word_len.cpu()
            held_all = set().union(*(word_hold[i] for i in word_pending))
            normal_open = [idx[c] for c in cells if idx[c] not in committed and idx[c] not in held_all]
            cmask = torch.zeros(tok.shape[0], dtype=torch.bool)
            if committed: cmask[torch.tensor(sorted(committed))] = True
            used = torch.zeros(len(V.itos), dtype=torch.bool)
            if committed: used[torch.tensor(sorted({int(tok[k]) for k in committed if tok[k] >= 4}), dtype=torch.long)] = True
            cand = []
            for i in sorted(word_pending):
                gi = line_group.get(i)
                if gi is not None and gi not in group_family: continue            # family not chosen yet
                n_i = lens[i]
                rest = [idx[(i, j)] for j in range(n_i) if idx[(i, j)] not in word_hold[i]]
                n_open_rest = sum(k not in committed for k in rest)
                sc = wlp[i].clone()
                last = wid.gather(1, (wlen - 1)[:, None])[:, 0]
                # the line-final syllable must pass the family / tone / same-syllable-in-group masks
                ok = (wlen < n_i) & (lpc[idx[(i, n_i - 1)], last] > -1e8)
                lp_sum = torch.zeros(len(wlen)); n_open = torch.zeros(len(wlen), dtype=torch.long)
                n_rep = torch.zeros(len(wlen))
                for q in range(WORD_MAX_K):
                    in_word = q < wlen
                    j = (n_i - wlen + q).clamp(min=0, max=n_i - 1)               # cell of the word's q-th syllable
                    syl = wid[:, q]
                    cellk = TL + i * C + j                                        # == idx[(i, j)]
                    comm = cmask[cellk]
                    # a committed cell must already hold that syllable; an open one is scored under the canvas
                    ok = ok & (~in_word | ~comm | (tok[cellk] == syl))
                    lpq = lp_raw[cellk, syl.clamp(min=0)]
                    lp_sum = lp_sum + torch.where(in_word & ~comm, lpq, torch.zeros_like(lpq))
                    n_open = n_open + (in_word & ~comm).long()
                    n_rep = n_rep + (in_word & ~comm & used[syl.clamp(min=0)]).float()
                    for rj in recv_cols[i]:                                       # never cover a receiving cell
                        ok = ok & ~(in_word & (j == rj))
                # per-cell MEAN of the raw canvas log-probs (a sum charges each extra syllable); the head's
                # log-prob carries the word's own probability
                sc = sc + word_mlm_w * lp_sum / n_open.clamp(min=1)
                # the same -2 per already-used syllable every other cell pays (the raw scale above omits it;
                # smoke without it: cross-line repeats 2.07 vs CTL 1.10)
                sc = sc - 2.0 * n_rep
                sc[~ok] = -1e9
                if float(sc.max()) < -1e8:
                    word_pending.discard(i); continue                           # released to the normal path
                pr = F.softmax(sc / temp, 0)
                # context rule only: a lexicon softmax and a syllable softmax are not comparable, and a "more confident
                # than the best ordinary commit" clause drew endings EARLIER than the control (smoke: 0.24 vs 0.31 of the
                # schedule), the opposite of the purpose
                ready = n_open_rest <= 2 or t >= T - 1 - len(word_pending) or not normal_open
                if ready: cand.append((float(pr.max()), i, pr))
            # one word per step while ordinary cells remain; every ready line at once when only held cells are
            # left (otherwise a held ending would have to fall back to cell-by-cell commits)
            for _, i, pr in sorted(cand, reverse=True)[: (len(cand) if not normal_open or t >= T - 1 else 1)]:
                w = int(torch.multinomial(pr, 1)); k_ = int(wlen[w])
                if trace is not None: trace.append(("word", i, w, k_))
                for q in range(k_):
                    kk = idx[(i, lens[i] - k_ + q)]
                    if kk not in committed:
                        tok[kk] = int(wid[w, q]); committed.add(kk)
                        if trace is not None: trace.append(kk)
                        if commit_log is not None: commit_log[kk] = t
                # cells of the hold set the word did not cover (a shorter word) go back to the normal path
                word_pending.discard(i); word_step = True
        probs = F.softmax(logits / temp, dim=-1)
        conf = (probs if pbon is None else F.softmax((logits - pbon) / temp, dim=-1)).max(dim=-1).values
        open_cells = [idx[c] for c in cells if idx[c] not in committed]
        if word_pending:
            _wh = set().union(*(word_hold[i] for i in word_pending))
            open_cells = [k for k in open_cells if k not in _wh] or open_cells
        if rhyme_last:
            # Rhyme cells commit LAST in their line: 56% of judged bad spans run up to the line-final
            # rhyme word (23% by chance), i.e. the line bends to reach a syllable chosen before it
            # existed. A rhyme cell may commit only once every other cell of its line has.
            if t < T - 1:   # the last step commits whatever is left; a held cell must never be dropped
                open_cells = [k for k in open_cells if k not in rhyme_hold() or eligible(cell_of[k])] or open_cells
        elif t < fam_delay:                       # hold the rhyme cells open until the family is chosen
            held = {idx[c] for gi in range(len(groups)) for c in (groups_alt if group_pos.get(gi) else groups)[gi]}
            open_cells = [k for k in open_cells if k not in held] or open_cells
        n_pre = sum(len(v) for v in (prefill or {}).values())
        n_commit = max(1, math.ceil((len(cells) - n_pre) * (t + 1) / T) - (len(committed) - n_pre))
        order = sorted(open_cells, key=lambda k: -float(conf[k]))
        if rhyme_first and t == 0:
            # Asking for the rhyme word off a blank canvas: probed on 30 human บท, the model puts
            # 2.9% (median) of its mass on the family the poet used and ranks the poet's own rhyme
            # word 3rd. With the rest of the บท known the same model ranks it 1st. rhyme_first=False
            # lets confidence order the rhyme cells like any other.
            rh = [idx[c] for gi in range(len(groups)) for c in (groups_alt if group_pos.get(gi) else groups)[gi] if idx[c] not in committed]
            if word_pending: rh = [k for k in rh if k in open_cells]      # held line ends wait for their word
            order = list(dict.fromkeys(rh + order))
        if commit_n:
            n_commit = min(n_commit, commit_n)          # 1 removes the conditional-independence error
        held_now = rhyme_hold() if (rhyme_last and seam_repair) else set()
        if word_step: order = []          # the word was this step's commit (its neighbours are chosen against it next step)
        for k in order[:n_commit]:
            tok[k] = int(torch.multinomial(probs[k], 1)); committed.add(k)
            if commit_log is not None: commit_log[k] = t
            if trace is not None: trace.append(k)      # commit order, for tests
            if k in held_now:
                # seam repair (review 2026-09-26): with rhyme-last the rhyme cell's neighbours were fixed
                # while it was still MASK, so they never saw the family - the model's own likelihood of
                # the syllable before a line-final rhyme fell from -1.67 to -3.17. Once the rhyme word
                # exists, re-open its immediate neighbours (once each) so they are re-chosen against it.
                i, j = cell_of[k]
                for jj in (j - 1, j + 1):
                    if 0 <= jj < lens[i]:
                        nb = idx[(i, jj)]
                        if nb in committed and nb not in held_now and nb not in repaired and \
                                not (prefill and i in prefill):
                            committed.discard(nb); tok[nb] = MASK; repaired.add(nb)
        for k in open_cells:
            if k not in committed: tok[k] = MASK
    lines = ["".join(V.itos[int(tok[idx[(i, j)]])] for j in range(lens[i]) if int(tok[idx[(i, j)]]) >= 4) for i in range(L)]
    return "\n".join(lines)


def rows_preview(path):
    return [(lambda r: r.get("text") or r.get("prompt"))(json.loads(l)) for l in open(path, encoding="utf-8") if l.strip()]


def refill(a, model, V, seg, clean, fam_prior, topic_freq, eff_boost):
    """Re-draw ONE วรรค of an already-selected poem, keeping the other three byte for byte.

    Why: on the 220 judged benchmark poems 20% fail on cleanness alone with exactly one bad วรรค.
    Selection throws such a poem away and hopes another candidate is wholly clean; this keeps the three
    good วรรค and offers the selector K alternatives for each line. The rhyme family of every link
    touching a kept line is inherited; the refilled line still chooses its own rhyme word inside it.

    --prompts is a best-pick file (id, prompt, poem). Writes OUT/r{line}_{k}.jsonl, one row per id, so
    the files drop straight into select_by_typhoon.py next to the original. A poem whose lines do not
    round-trip through the segmenter and vocabulary is copied unchanged, and counted.
    """
    assert a.commit_n >= 1, "refill schedules one commit per step; pass --commit-n 1"
    rows = [json.loads(l) for l in open(a.prompts, encoding="utf-8") if l.strip()]
    if a.limit: rows = rows[:a.limit]
    todo = [int(x) - 1 for x in a.refill_lines.split(",")]
    line_map = None
    if a.refill_line_map:
        # one line per poem, chosen by a locator (e.g. the canvas score's weakest วรรค), 1-based
        line_map = {k: int(v) - 1 for k, v in json.load(open(a.refill_line_map, encoding="utf-8")).items()}
        todo = [0, 1, 2, 3]
    tf = topic_freq if a.refill_boost_mode == "rarity" else None
    os.makedirs(a.out, exist_ok=True)
    outs = {(i, k): open(os.path.join(a.out, f"r{i + 1}_{k}.jsonl"), "w", encoding="utf-8")
            for i in todo for k in range(a.refill_k)}
    n_skip = n_pass = n_tot = 0
    for ri, r in enumerate(rows, 1):
        text = r.get("prompt") or r.get("text")
        orig_lines = r["poem"].split("\n")
        orig = [[V.enc(sy) for sy in seg(ln, C)] for ln in orig_lines]
        ok = len(orig_lines) == L and all(UNK not in o and "".join(V.itos[t] for t in o) == ln
                                          for o, ln in zip(orig, orig_lines))
        title = seg(re.split(r"\s+|[,;.!?]", text)[0], TL)
        for i in todo:
            if line_map is not None and line_map.get(r["id"]) != i:
                continue
            for k in range(a.refill_k):
                if not ok:
                    outs[(i, k)].write(json.dumps({**r, "refill_line": None, "sample": k}, ensure_ascii=False) + "\n")
                    continue
                _rs = (a.seed * 100003 + int(hashlib.sha256(f"{r['id']}|{i}|{k}".encode()).hexdigest()[:8], 16))
                torch.manual_seed(_rs % (2**31)); rng_len = random.Random(_rs)
                want = pick_lens(a.wak_lens, rng_len)[i]
                cells_i, best = want, None
                prefill = {l: orig[l] for l in range(L) if l != i}
                for attempt in range(a.retries):
                    lens = [len(o) for o in orig]; lens[i] = cells_i
                    steps = max(a.steps, cells_i // a.commit_n + 1)
                    poem = decode_one(model, V, title, T=steps, temp=a.temp, lens=lens, fam_temp=a.fam_temp,
                                      topic=clauses_of(text) if a.topic_boost else None, topic_boost=eff_boost(text), inner_boost=a.inner_boost, clean=clean, fam_distinct=a.fam_distinct, fam_size_alpha=a.fam_size_alpha, fam_prior=fam_prior, fam_prior_w=a.fam_prior_w, ban_junk=a.ban_junk, seg=seg, rhyme_first=not a.no_rhyme_first, commit_n=a.commit_n, rhyme_pos=a.rhyme_pos, fam_soft=a.fam_soft, fam_delay=a.fam_delay, rhyme_last=a.rhyme_last, seam_repair=a.seam_repair, anaphora_penalty=a.anaphora_penalty, poetic=a.poetic_vec, word_rhyme=a.word_rhyme, word_mlm_w=a.word_mlm_w, topic_zero=a.topic_zero, topic_freq=tf,
                                      prefill=prefill, orig=orig)
                    new = poem.split("\n")
                    assert all(new[l] == orig_lines[l] for l in prefill), f"{r['id']}: a kept วรรค changed"
                    s_ = P.score(poem, n_bot_expected=N_BOT)
                    got = syl_counts(poem)
                    key = (got[i] == want, s_["all_pass"], s_["ext_rhyme"] + s_["tone"] + 0.5 * s_["internal"])
                    if best is None or key > best[0]: best = (key, poem, s_, got)
                    if got[i] == want and s_["all_pass"]: break
                    cells_i = max(5, min(C, cells_i - (got[i] - want)))
                _, poem, s_, got = best
                n_tot += 1; n_pass += bool(s_["all_pass"])
                outs[(i, k)].write(json.dumps({"id": r["id"], "prompt": text, "sample": k, "poem": poem,
                    "refill_line": i + 1, "orig_poem": r["poem"], "syl": got, "topic_boost": eff_boost(text), "poetic_boost": a.poetic_boost,
                    "refill_boost_mode": a.refill_boost_mode, "all_pass": s_["all_pass"], "structure": s_["structure"],
                    "ext_rhyme": s_["ext_rhyme"], "tone": s_["tone"], "internal": s_["internal"]}, ensure_ascii=False) + "\n")
        n_skip += not ok
        print(f"{ri}/{len(rows)} {'refilled' if ok else 'COPIED (no round-trip)'}", file=sys.stderr, flush=True)
    for f in outs.values(): f.close()
    print(f"refilled {n_tot} variants, form {100 * n_pass / max(n_tot, 1):.1f}%; "
          f"{n_skip}/{len(rows)} poems copied unchanged (segmenter round-trip failed)", file=sys.stderr, flush=True)
    P.save_cache()


def decode(a):
    model, V = load_model(a.model)
    if getattr(model, "diction", False) and getattr(a, "diction_level", None) is None:
        a.diction_level = 0                      # as trained: level 0 = "no level"
    if getattr(a, "diction_level", None) is not None:
        assert model.diction, f"--diction-level needs a model trained with --diction ({a.model} was not)"
        assert 0 <= a.diction_level <= DICTION_LEVELS, f"--diction-level must be 0..{DICTION_LEVELS}"
        model.diction_level = a.diction_level
        print(f"  diction level {a.diction_level} (0 = none, 1..{DICTION_LEVELS} = within-board quartile of poetic rhyme-zone cells)", file=sys.stderr, flush=True)
    if a.word_rhyme:
        assert getattr(model, "word_lex", None), f"--word-rhyme needs a model trained with --word-loss-w > 0 ({a.model} was not)"
        print(f"  word-rhyme: line endings drawn as words from {len(model.word_lex)} (canvas weight {a.word_mlm_w})", file=sys.stderr, flush=True)
    seg = (lambda t, lim=None: syl_align.syls_aligned(t, lim)) if getattr(V, "seg", "grid") == "aligned" else syls_of

    def eff_boost(text):
        """A short prompt gets --short-topic-boost, not --topic-boost (sweep 2026-09-20: 3.0 repeats
        the topic words, 1.5 keeps form). That is deliberate, but it means a job script saying
        '--topic-boost 3.0' applies 1.5 to 98% of our topics, and runs have been reported under the
        wrong number. The effective value is printed and recorded per row."""
        return a.short_topic_boost if len(seg(text)) <= 8 else a.topic_boost
    if getattr(a, "rhyme_last", False) and (a.fam_delay or a.no_rhyme_first is False):
        print("  NOTE: --rhyme-last overrides --fam-delay and rhyme-first ordering", file=sys.stderr, flush=True)
    print(f"decoding {a.model}: topic segmented with '{getattr(V, 'seg', 'grid')}' (as trained)", file=sys.stderr, flush=True)
    _n_short = sum(1 for r in rows_preview(a.prompts) if len(seg(r)) <= 8)
    _n_all = len(rows_preview(a.prompts))
    print(f"  topic boost: {a.short_topic_boost} on {_n_short}/{_n_all} prompts (<=8 syllables), "
          f"{a.topic_boost} on the other {_n_all - _n_short}", file=sys.stderr, flush=True)
    clean = CleanPrior(V, a.ngrams, a.rhyme_lex, a.bigram_w) if (a.rhyme_lex or a.bigram_w) else None
    fam_prior = None
    if a.fam_prior_w:
        d = json.load(open(a.ngrams, encoding="utf-8"))["fam_final"]; tot = sum(d.values())
        rime_key = {}
        for i, s in enumerate(V.itos):
            if i < 4: continue
            r, _ = phon(s)
            if r and V.rime[i]: rime_key[V.rime[i]] = "|".join(r)
        fam_prior = {fid: math.log(d.get(k, 0.5) / tot) for fid, k in rime_key.items()}   # log human line-final family frequency
    a.poetic_vec = None
    if a.poetic_boost:
        lex = json.load(open(a.poetic_lex, encoding="utf-8"))
        a.poetic_vec = torch.zeros(len(V.itos))
        for sy, lo in lex.items():
            if V.enc(sy) != UNK: a.poetic_vec[V.enc(sy)] = a.poetic_boost * min(max((lo - 1.5) / 3.0, 0.0), 1.0)
        print(f"  poetic boost {a.poetic_boost}: {int((a.poetic_vec > 0).sum())} syllables from {a.poetic_lex}, "
              f"mean bonus {float(a.poetic_vec[a.poetic_vec > 0].mean()):.2f}", file=sys.stderr, flush=True)
    topic_freq = None
    if a.boost_mode == "rarity":
        topic_freq = json.load(open(a.syl_freq, encoding="utf-8"))
        assert topic_freq.get("total", 0) > 0 and topic_freq.get("count"), f"{a.syl_freq} is not a frequency table"
        print(f"  boost mode rarity: {topic_freq['total']:,} corpus syllable tokens, "
              f"{len(topic_freq['count']):,} types", file=sys.stderr, flush=True)
    if a.refill_k:
        return refill(a, model, V, seg, clean, fam_prior, topic_freq, eff_boost)
    rows = [json.loads(l) for l in open(a.prompts, encoding="utf-8") if l.strip()]
    if a.limit: rows = rows[:a.limit]
    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    torch.manual_seed(a.seed)
    rng_len = random.Random(a.seed)
    n_hit = [0, 0]
    with open(a.out, "w", encoding="utf-8") as f:
        for i, r in enumerate(rows, 1):
            # Seed PER ROW. Seeding once before the loop makes a row's randomness depend on how much
            # every earlier row consumed, so two arms decoded with "the same seed" drift apart from
            # row 1 onward even where their inputs are identical - 83 of 211 supposedly-differing
            # pairs in runs/RARITY64 differ only by that drift. The invariant this restores: a
            # one-row prompt file must produce the same poem as that row inside the full file.
            # Seed per ROW IDENTITY, not per position. Seeding once before the loop makes a row's
            # randomness depend on how much every earlier row consumed, so two arms decoded with
            # "the same seed" drift apart from the first treated row onward - 83 of 211 supposedly
            # differing pairs in runs/RARITY64 differ only by that drift. Seeding by position would
            # fix the arms but still make a row's poem depend on where it sits in the file; keying on
            # the id gives the invariant that actually matters: one row decodes the same alone as it
            # does inside the full set.
            _rs = a.seed * 100003 + int(hashlib.sha256(str(r["id"]).encode()).hexdigest()[:8], 16)
            torch.manual_seed(_rs % (2**31))
            rng_len = random.Random(_rs)
            title = seg(re.split(r"\s+|[,;.!?]", r["text"])[0], TL)
            # A grid cell is an approximation of a syllable: the segmenter that fills the cells reads
            # ถวิล as one where it is ถ|วิล. So the target is a PRONOUNCED count, the poem is counted
            # with g2p_line after decoding, and the cell budget is corrected by the error and retried.
            want = pick_lens(a.wak_lens, rng_len)
            steps = a.steps if not a.commit_n else max(a.steps, sum(want) // a.commit_n + 1)
            cells, best = list(want), None
            for attempt in range(a.retries):
                if a.rhyme_last and a.commit_n:
                    # retries can grow the cells past sum(want); with rhyme-last the cells left over would
                    # be exactly the held rhyme cells, so the budget follows the ACTUAL cell count. Gated
                    # on the flag so ordinary decoding stays byte-identical to what was measured.
                    steps = max(a.steps, sum(cells) // a.commit_n + 1 + (16 if a.seam_repair else 0))
                _cl = {} if a.log_commit_steps else None
                poem = decode_one(model, V, title, T=steps, temp=a.temp, lens=cells, fam_temp=a.fam_temp,
                                  topic=clauses_of(r["text"]) if a.topic_boost else None, topic_boost=eff_boost(r["text"]), inner_boost=a.inner_boost, clean=clean, fam_distinct=a.fam_distinct, fam_size_alpha=a.fam_size_alpha, fam_prior=fam_prior, fam_prior_w=a.fam_prior_w, ban_junk=a.ban_junk, seg=seg, rhyme_first=not a.no_rhyme_first, commit_n=a.commit_n, rhyme_pos=a.rhyme_pos, fam_soft=a.fam_soft, fam_delay=a.fam_delay, rhyme_last=a.rhyme_last, seam_repair=a.seam_repair, anaphora_penalty=a.anaphora_penalty, poetic=a.poetic_vec, word_rhyme=a.word_rhyme, word_mlm_w=a.word_mlm_w, commit_log=_cl, topic_zero=a.topic_zero, topic_freq=topic_freq)
                s = P.score(poem, n_bot_expected=N_BOT)
                got = syl_counts(poem)
                hit = sum(1 for g, w in zip(got, want) if g == w)
                key = (hit, s["all_pass"], s["ext_rhyme"] + s["tone"] + 0.5 * s["internal"])
                if best is None or key > best[0]:
                    best = (key, poem, s, list(cells), got)
                    if _cl is not None:
                        # refinement step at which each line's final cell and first tail cell (last 3) committed
                        best_steps = {"end_step": [_cl.get(TL + i * C + cells[i] - 1) for i in range(L)],
                                      "tail_first_step": [min((_cl[k] for k in range(TL + i * C + max(0, cells[i] - WORD_MAX_K), TL + i * C + cells[i]) if k in _cl), default=None) for i in range(L)],
                                      "n_steps": steps}
                if hit == L and s["all_pass"]: break
                cells = [max(5, min(C, c - (g - w))) for c, g, w in zip(cells, got, want)]
            _, poem, s, cells, got = best
            eb = eff_boost(r["text"])
            n_hit[0] += sum(1 for g, w in zip(got, want) if g == w); n_hit[1] += L
            f.write(json.dumps({"id": r["id"], "prompt": r["text"], "sample": 0, "poem": poem,
                                "syl_target": want, "syl": got, "cells": cells, "topic_boost": eb, "poetic_boost": a.poetic_boost, "poetic_lex": a.poetic_lex if a.poetic_boost else None, **({"word_rhyme": True, "word_mlm_w": a.word_mlm_w} if a.word_rhyme else {}), **(best_steps if a.log_commit_steps else {}), "topic_zero": a.topic_zero, "boost_mode": a.boost_mode, "all_pass": s["all_pass"],
                                "structure": s["structure"], "ext_rhyme": s["ext_rhyme"], "tone": s["tone"], "internal": s["internal"]},
                               ensure_ascii=False) + "\n"); f.flush()
            print(f"{i}/{len(rows)} all_pass={s['all_pass']} syl={got} want={want}", file=sys.stderr, flush=True)
    print(f"วรรค hitting their pronounced syllable target: {n_hit[0]}/{n_hit[1]} ({100*n_hit[0]/max(n_hit[1],1):.1f}%)", file=sys.stderr, flush=True)
    P.save_cache()


if __name__ == "__main__":
    ap = argparse.ArgumentParser(); sub = ap.add_subparsers(dest="cmd")
    t = sub.add_parser("train"); t.add_argument("--data", default="data/modern_klon8.jsonl"); t.add_argument("--out", default="runs/CANVAS_WC")
    t.add_argument("--epochs", type=int, default=3); t.add_argument("--batch", type=int, default=32); t.add_argument("--lr", type=float, default=5e-5)
    t.add_argument("--expert", choices=["none", "asym", "sym"], default="none", help="extra transformer block inserted mid-stack (after body layer EXPERT_AFTER): asym = gated to the rhyme zone only (user's asymmetry idea), sym = every cell (capacity-matched control)")
    t.add_argument("--expert-lr", type=float, default=1e-4); t.add_argument("--gate-lr", type=float, default=3e-3)
    t.add_argument("--diction", action="store_true", help="add the gated diction-level input channel (user idea 2026-09-27; levels from runs/LOC/poetic_syl.json); the expert/gate learning rates apply to it")
    t.add_argument("--zone-weight", type=float, default=1.0, help="loss weight on rhyme-zone cells (1 = unweighted)")
    t.add_argument("--word-loss-w", type=float, default=0.0, help="C1: weight of the auxiliary ending-word loss; > 0 adds the word-rhyme head (0 = off, the old trainer)")
    t.add_argument("--word-lex", default="runs/LOC/rhyme_words.json", help="C1 lexicon built by scripts/build_rhyme_lex.py (copied into --out)")
    t.add_argument("--det-bank", default=None, help="B: in-context corruption bank from scripts/det_prep.py (title, grid, lab, line, split)")
    t.add_argument("--det-weight", type=float, default=0.0, help="B: weight of the replaced-cell detection loss (0 = off, no head: the old trainer)")
    t.add_argument("--det-batch", type=int, default=0, help="B: corrupted poems per step for the detection term (0 = --batch)")
    t.add_argument("--ul-data", default=None, help="A': flagged generated poems from scripts/ul_prep.py flags (title, grid, pairs)")
    t.add_argument("--ul-weight", type=float, default=0.0, help="A': weight of the unlikelihood term on flagged word pairs of the model's own poems (0 = off, the old trainer)")
    t.add_argument("--ul-heldout", default=None, help="ULKL: title-held-out flags for the in-training tripwire (off when absent)")
    t.add_argument("--kl-beta", type=float, default=0.0, help="ULKL: weight of KL(frozen init || model) on all non-PAD cells of the MLM and unlikelihood batches (0 = off)")
    t.add_argument("--ul-batch", type=int, default=0, help="A': generated poems per step for the unlikelihood term (0 = --batch)")
    t.add_argument("--stutter-ul", type=float, default=0.0, help="C5: weight of the unlikelihood loss on predicting a neighbour's syllable (0 = off)")
    t.add_argument("--patience", type=int, default=3); t.add_argument("--min-delta", type=float, default=0.01)
    t.add_argument("--init", default=None, help="warm-start state dict, e.g. runs/CANVAS_WC8/model.pt")
    t.add_argument("--seg", choices=["aligned", "grid"], default="aligned", help="how a วรรค becomes cells: aligned = one cell per PRONOUNCED syllable (scripts/syl_align.py), grid = the old character-level segmenter")
    t.add_argument("--init-vocab", default=None, help="the warm-start checkpoint's vocab.json; rows for syllables present in both vocabularies are carried over when --seg changes the segmentation")
    t.add_argument("--vocab", default=None, help="freeze the vocabulary from this vocab.json instead of rebuilding it from --data; REQUIRED with --init when --data is not the original corpus")
    t.add_argument("--valid", type=int, default=500, help="validation poems held out from the front (capped at 20%% of the data)")
    d = sub.add_parser("decode"); d.add_argument("--model", default="runs/CANVAS_WC"); d.add_argument("--prompts", required=True)
    d.add_argument("--diction-level", type=int, default=None, help="diction level for a model trained with --diction: 1..4 = within-board quartile of poetic rhyme-zone cells, 0 = no level (the default for such a model)")
    d.add_argument("--out", required=True); d.add_argument("--limit", type=int, default=0); d.add_argument("--steps", type=int, default=12)
    d.add_argument("--temp", type=float, default=0.8); d.add_argument("--seed", type=int, default=0); d.add_argument("--fam-temp", type=float, default=0.5)
    d.add_argument("--retries", type=int, default=4); d.add_argument("--topic-boost", type=float, default=0.0)
    d.add_argument("--short-topic-boost", type=float, default=1.5, help="topic bonus used instead of --topic-boost when the prompt is a short topic (<= 8 syllables); 3.0 repeats the topic words, 1.5 keeps form (sweep 2026-09-20)")
    d.add_argument("--fam-distinct", action="store_true", help="each of the 4 rhyme groups in a poem takes a different rime family (humans: 3.87/4)")
    d.add_argument("--fam-size-alpha", type=float, default=0.0, help="subtract alpha*log(family size) from the family score (1 = mean log-prob per member)")
    d.add_argument("--fam-prior-w", type=float, default=0.0, help="add w*log(human line-final family frequency) to the family score")
    d.add_argument("--rhyme-lex", type=int, default=0, help="rhyme cells limited to syllables seen >= N times at that position in the corpus (0 = off)")
    d.add_argument("--bigram-w", type=float, default=0.0, help="weight of the corpus syllable-bigram log-probability toward committed neighbours (0 = off)")
    d.add_argument("--wak-lens", default="8", help="target PRONOUNCED syllables per วรรค, verified with g2p_line after decoding and corrected by retry: 8, sample (drawn from the human distribution: 47%% of human บท are all-8, the usual alternative is 9), or a pattern like 8,8,9,8")
    d.add_argument("--boost-mode", choices=["flat", "rarity"], default="flat", help="flat: every topic syllable gets the same bonus, whatever its corpus prior (the original). rarity: the bonus is multiplied by clip(log(1e-3/p), 1, 3), so a syllable the corpus almost never uses is pushed hard enough to compete. Measured 2026-09-24: the topic gate is 4%% when the rarest content syllable sits below 1e-5 and 83%% above 1e-3, while cleanness is flat across that split")
    d.add_argument("--syl-freq", default="data/syl_freq.json", help="corpus syllable frequencies for --boost-mode rarity, segmented as trained")
    d.add_argument("--topic-zero", choices=["global", "row"], default="row", help="when a topic syllable is committed, stop boosting it everywhere (global, the original) or only in the วรรค that used it (row, the default since 2026-09-24). Global kills the topic channel after ~3 of ~33 commits: judged at N=8 on 50 paired topics it costs 26%% against row's 50%%, with the topic gate at 40%% against 74%% (McNemar p = 0.008)")
    d.add_argument("--fam-delay", type=int, default=0, help="postpone the rhyme-family choice until this refinement step, so content commits first and the rhyme is chosen against a line that exists (0 = the old behaviour, chosen at step 0 off a blank canvas)")
    d.add_argument("--fam-soft", type=float, default=0.0, help="replace the hard rhyme-family mask with a finite penalty of this size (e.g. 6). 0 = the old wall. Lets a much better word outrank the family; the checker and --retries decide whether the poem still scans")
    d.add_argument("--rhyme-pos", choices=["3", "auto"], default="3", help="which syllable receives a rhyme link: 3 (the only one the decoder ever tried) or auto (3 or 5, whichever the model prefers; humans use 5 in 10-12%% of links and the checker accepts both)")
    d.add_argument("--no-rhyme-first", action="store_true", help="do NOT commit the rhyme cells before everything else; let confidence order them. The model ranks the human's own rhyme word 1st when the บท is known and 3rd off a blank canvas")
    d.add_argument("--commit-n", type=int, default=0, help="cells committed per refinement step (0 = the old schedule, ~3). 1 removes the conditional-independence error of sampling several cells from one forward pass")
    d.add_argument("--ban-junk", action="store_true", help="forbid the 725 vocabulary entries containing . * or ๐ (crawl padding that became vocabulary)")
    d.add_argument("--ngrams", default="data/syl_ngrams.json")
    d.add_argument("--rhyme-last", action="store_true", help="commit each rhyme cell only after every other cell of its line, and choose a group's family only then (overrides rhyme-first and fam-delay). 2026-09-26: 56%% of judged bad spans run up to the line-final rhyme word against 23%% by chance")
    d.add_argument("--poetic-boost", type=float, default=0.0, help="bonus at uncommitted rhyme-zone cells toward poetic-diction syllables, scaled by clip((log-odds - 1.5)/3, 0, 1); 0 = off (plan B)")
    d.add_argument("--poetic-lex", default="runs/LOC/poetic_syl.json", help="syllable -> poem-vs-prose log-odds, built by scripts/poetic_syllables.py")
    d.add_argument("--anaphora-penalty", type=float, default=0.0, help="penalty on a line-opening syllable (position 1 or 2) that two other lines already open with; 0 = off. 32%% of our poems vs 1%% of human poems open 3+ lines alike")
    d.add_argument("--word-rhyme", action="store_true", help="C1: draw each line ending as a whole WORD from the model's ending-word head (inside the chosen rhyme family) and commit its syllables jointly; needs a model trained with --word-loss-w > 0")
    d.add_argument("--log-commit-steps", action="store_true", help="record per row the refinement step at which each line's final cell / first of its last 3 cells committed")
    d.add_argument("--word-mlm-w", type=float, default=1.0, help="C1: weight of the canvas's own syllable log-probs added to the head's word log-prob (0 = head only)")
    d.add_argument("--seam-repair", action="store_true", help="with --rhyme-last: when a rhyme cell commits, re-open its immediate neighbours once so they are re-chosen against the rhyme word")
    d.add_argument("--refill-k", type=int, default=0, help="REFILL mode: --prompts is a best-pick file (id, prompt, poem); re-draw each วรรค in --refill-lines K times with the other three kept, writing --out/r{line}_{k}.jsonl (--out is a directory)")
    d.add_argument("--refill-lines", default="1,2,3,4")
    d.add_argument("--refill-line-map", default=None, help="JSON {poem id: line 1-4}: refill only that line of each poem (rows without an entry are skipped); files r{line}_{k} then hold only the poems located on that line")
    d.add_argument("--refill-boost-mode", choices=["flat", "rarity"], default="flat", help="topic boost on the refilled line: flat (default) because the rarity-boosted syllable is in 80%% of the bad วรรค it produced")
    d.add_argument("--inner-boost", type=float, default=0.0, help="สัมผัสใน steering bonus toward a committed template partner's rime")
    a = ap.parse_args()
    train(a) if a.cmd == "train" else decode(a)