src/scripts/twolayer_canvas.py

"""Two-layer canvas (user idea 2026-09-29): SOUND layer = one cell per pronounced syllable, WORD layer = whole written
words spanning 1+ sound cells. One model, one loop.

Grid: title 12 cells + 4 วรรค x 10 cells, as the current canvas. Each cell holds
  word-slot token (w, k): word w of the lexicon, its k-th pronounced syllable (w has n_w syllables)
  sound channels: onset / rime family / tone of that syllable (exact, in context: the checker's own g2p path)
Embedding of (w, k)   = word_emb[w] + kn_emb[(k, n_w)]            (word rows initialised from WangchanBERTa pieces)
Logit of (w, k) at h  = z.word_emb[w] + z.kn_emb[(k, n_w)] + bias[w]    with z = pretrained MLM head transform(h)
So the output is a distribution over whole words placed at a position inside the word; the pretrained collocation
knowledge lives in word_emb, which for 86% of word tokens is exactly one native WangchanBERTa piece.

Training: whole-word masking (all cells of a word masked together, ratio U(0.15, 1.0) so the all-masked start of
decoding is seen), CE on masked cells; optional: rime reveal at masked rhyme cells (the sound layer is decided first
at decode), sound head (aux CE over pronounced syllables).

Decode: iterative; each step commits ONE WHOLE WORD across its cells. Rhyme family and tone masks act on the sound
of the cell exactly; a line-final cell only admits the LAST syllable of a word; a line-initial cell only the first.
Once a group's family is chosen its rime is written into the sound layer of its open cells (input channel), then
the words are chosen to fit it.

  twolayer_canvas.py train --data DIR --out RUN [--epochs E] [--limit N]      (DIR from twolayer_data.py build)
  twolayer_canvas.py decode --model RUN --prompts P.jsonl --out O.jsonl
  twolayer_canvas.py test --data DIR                                           (shapes, labels, overfit, decode smoke)
"""
import os, sys, json, math, random, argparse, collections, time
sys.path[:0] = [".", "scripts"]
import torch, torch.nn as nn, torch.nn.functional as F
from klon8 import prosody as P

WC = "airesearch/wangchanberta-base-att-spm-uncased"
DEV = torch.device(os.environ.get("KLON8_DEVICE") or ("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"))
C, TL, L = 10, 12, 4
T = TL + L * C
PAD, MASK, UNK, SEP = 0, 1, 2, 3                         # word ids AND token ids (specials have k=0, n=1)
KMAX = 8                                                 # longest word in the lexicon (8 syllables)
EVEN, CHECKED = (0, 3), (1, 2, 4)
TONE_AT = {1: CHECKED, 2: EVEN, 3: EVEN}                 # วรรค index (0-based) -> allowed final tones
# rhyme groups, (line 0-based, col) with col -1 = line-final; the receiver (col 2) may move to col 4 (auto)
GROUPS = [[(0, -1), (1, 2)], [(1, -1), (2, -1), (3, 2)]]
ROLE = [1, 2, 3, 4]


# ------------------------------------------------------------------ lexicon / tables

class Lex:
    """Words, word-slot tokens and sound tables."""
    def __init__(self, lex_json):
        d = json.load(open(lex_json, encoding="utf-8")) if isinstance(lex_json, str) else lex_json
        self.words = ["<pad>", "<mask>", "<unk>", "<sep>"] + list(d["words"])
        self.sounds = d["sounds"]; self.count = d.get("count", {})
        self.word_seen = torch.tensor([i >= 4 and self.count.get(w, 0) > 0 for i, w in enumerate(["<pad>", "<mask>", "<unk>", "<sep>"] + list(d["words"]))])
        self.wid = {w: i for i, w in enumerate(self.words)}
        # sound inventory: every syllable of a lexicon word (data cells outside it -> id 0 "unknown")
        phons = sorted({p for w in d["words"] for p in self.sounds[w]})
        self.phons = ["?"] + phons; self.pid = {p: i for i, p in enumerate(self.phons)}
        on, ri = {"?": 0}, {"?": 0}
        self.p_on, self.p_ri, self.p_to = [0], [0], [5]
        for p in phons:
            s = P._parse_syl("", p)
            self.p_on.append(on.setdefault(s.onset, len(on)))
            self.p_ri.append(ri.setdefault(P._MERGED_RIME.get(s.rime, s.rime), len(ri)))
            self.p_to.append(s.tone if 0 <= s.tone <= 4 else 5)   # g2p can emit a stray tone digit ('sip8'); 5 = unknown
        self.rimes = ri; self.n_on, self.n_ri = len(on), len(ri)
        assert set(self.p_to) <= set(range(6)), 'tone ids must index to_emb (6 rows)'
        # tokens: specials 0..3, then for each word its n_w slots
        tw, tk, tn, ts = [0, 1, 2, 3], [0] * 4, [1] * 4, [0] * 4
        self.first = [0, 1, 2, 3]
        for w in self.words[4:]:
            self.first.append(len(tw)); sy = self.sounds[w]
            for k, p in enumerate(sy):
                tw.append(self.wid[w]); tk.append(k); tn.append(len(sy)); ts.append(self.pid[p])
        self.tok_word, self.tok_k, self.tok_n, self.tok_sound = map(torch.tensor, (tw, tk, tn, ts))
        self.kn_of = lambda k, n: n * (n - 1) // 2 + k           # (k, n) pairs, n <= KMAX
        self.tok_kn = torch.tensor([self.kn_of(k, n) for k, n in zip(tk, tn)])
        self.n_kn = self.kn_of(KMAX - 1, KMAX) + 1
        self.tok_ri = torch.tensor([self.p_ri[s] for s in ts]); self.tok_to = torch.tensor([self.p_to[s] for s in ts])
        self.tok_on = torch.tensor([self.p_on[s] for s in ts])
        self.tok_last = self.tok_k == self.tok_n - 1
        self.NT = len(tw)

    def tok(self, w, k):
        i = self.wid.get(w, UNK)
        return i if i < 4 else self.first[i] + k


def encode(rec, X):
    """poem record -> (tok[T], sound[T], word_span[T]) ; word_span = index of the word occurrence (for whole-word masking)."""
    tok = torch.full((T,), PAD, dtype=torch.long); snd = torch.zeros(T, dtype=torch.long)
    span = torch.full((T,), -1, dtype=torch.long); occ = -1
    def put(base, cells):
        nonlocal occ
        for j, (w, k, n, ph) in enumerate(cells):
            if k == 0: occ += 1
            tok[base + j] = X.tok(w, k); snd[base + j] = X.pid.get(ph, 0); span[base + j] = occ
    put(0, rec["title"])
    for i, g in enumerate(rec["grid"]): put(TL + i * C, g)
    return tok, snd, span


def rhyme_cells(rec):
    """cells whose rime is part of a rhyme link in the DATA (line-final of every วรรค; receiver 3rd/5th if it rhymes)."""
    # line-final cells of วรรค 1-3 (the ส่ง line's final is in no rhyme group of a one-บท poem: never revealed at decode)
    out = [TL + i * C + len(g) - 1 for i, g in enumerate(rec["grid"][:3])]
    def rk(ph): s = P._parse_syl("", ph); return P._MERGED_RIME.get(s.rime, s.rime)
    for snd_line, rcv_line in ((0, 1), (1, 3)):
        src = rk(rec["grid"][snd_line][-1][3])
        for col in (2, 4):
            g = rec["grid"][rcv_line]
            if col < len(g) - 1 and rk(g[col][3]) == src: out.append(TL + rcv_line * C + col); break
    return out


# ------------------------------------------------------------------ model

_PIECES = {}


def piece_ids(X, wc_tok):
    """WangchanBERTa piece ids of every lexicon word (the leading "▁" dropped), computed once per process"""
    key = (len(X.words), X.words[4] if len(X.words) > 4 else "", X.words[-1])
    if key not in _PIECES:
        sp = wc_tok.convert_tokens_to_ids("▁")
        out = [[]] * 4
        for w in X.words[4:]:
            ids = [t for t in wc_tok(w, add_special_tokens=False)["input_ids"] if t != sp]
            out.append(ids or [wc_tok.unk_token_id])
        _PIECES[key] = out
    return _PIECES[key]


class TwoLayer(nn.Module):
    def __init__(self, X, wc_tok=None, wc=None, lm_head=None, peer=False, d=768, body=None, ln_in=None, abs_pos=True,
                 bias_mode="unigram", peer_in=True, peer_tied=False):
        super().__init__()
        self.peer_in, self.peer_tied = peer_in, peer_tied
        self.use_abs_pos = abs_pos
        self.X = X
        if wc is not None:
            body, ln_in, d = wc.encoder, wc.embeddings.LayerNorm, wc.config.hidden_size
        self.body, self.ln_in = body, ln_in
        W = len(X.words)
        self.word_emb = nn.Embedding(W, d)
        pre_bias = torch.zeros(W)
        if wc is not None:
            # every row (poem words AND unseen dictionary / WangchanBERTa words) from the pretrained pieces: the input row is
            # the mean of the pieces' embeddings (tied output), the output bias the mean of the pieces' pretrained MLM bias
            base = wc.embeddings.word_embeddings.weight.detach(); init = torch.zeros(W, d)
            pb = lm_head["bias"] if lm_head is not None else torch.zeros(base.shape[0])
            sp = [wc_tok.pad_token_id, wc_tok.mask_token_id, wc_tok.unk_token_id, wc_tok.sep_token_id]
            pieces = piece_ids(X, wc_tok)                                   # cached once per process (62k words)
            for i, w in enumerate(X.words):
                if i < 4: init[i] = base[sp[i]]; continue
                ids = pieces[i]
                init[i] = base[ids].mean(0); pre_bias[i] = pb[ids].mean()
            self.word_emb.weight.data.copy_(init)
        self.kn_emb = nn.Embedding(X.n_kn, d); nn.init.normal_(self.kn_emb.weight, std=0.02)
        # the body's own PRETRAINED absolute positions (RoBERTa offset 2) + token-type row: without them masked cells in a
        # run get near-identical inputs (one-batch overfit stalled at 0.53 with every cell of a masked run predicting
        # one word; the syllable canvas has the same omission - its fully-masked loss is 6.47 vs unigram 7.01)
        # and a leading <s> cell, always attended: the pretrained body uses <s> as its attention sink; with PAD masked
        # out and no <s>, poems with short titles collapsed (every masked cell the same hidden state, cos 1.0)
        self.abs_pos = nn.Parameter(torch.zeros(T, d)); self.bos = nn.Parameter(torch.zeros(d))
        if wc is not None:
            pe, te = wc.embeddings.position_embeddings.weight.detach(), wc.embeddings.token_type_embeddings.weight.detach()[0]
            self.abs_pos.data.copy_(pe[3:3 + T] + te)
            self.bos.data.copy_(wc.embeddings.word_embeddings.weight.detach()[wc_tok.bos_token_id] + pe[2] + te)
        self.on_emb = nn.Embedding(X.n_on, d); self.ri_emb = nn.Embedding(X.n_ri, d); self.to_emb = nn.Embedding(6, d)
        self.line = nn.Embedding(L + 1, d); self.pos = nn.Embedding(max(TL, C), d); self.role = nn.Embedding(5, d)
        for e in (self.on_emb, self.ri_emb, self.to_emb, self.line, self.pos, self.role): nn.init.normal_(e.weight, std=0.02)
        self.head_dense = nn.Linear(d, d); self.head_ln = nn.LayerNorm(d, eps=1e-5)
        # output bias: "pretrained" (default) = WangchanBERTa's own MLM bias for the word's pieces, so a word never seen in
        # the poems keeps the prior the pretrained LM gives it; "unigram" = centred log poem frequency (FIX 1's choice for
        # the syllable canvas, where every row is a poem unit). Chosen by step-0 masked loss (twolayer_test.py --bias-check).
        # "hybrid": poem words get the unigram bias; unseen words get the pretrained bias mapped onto the unigram scale
        # by a least-squares line fitted on the poem words (so their prior follows WangchanBERTa, not a flat floor)
        cnt = torch.tensor([float(X.count.get(w, 0)) for w in X.words])
        uni = torch.log((cnt + 1) / (cnt.sum() + W))
        if bias_mode == "unigram": bias = uni
        elif bias_mode == "pretrained": bias = pre_bias.clone()
        else:
            seen = X.word_seen
            A = torch.stack([pre_bias[seen], torch.ones(int(seen.sum()))], 1)
            coef = torch.linalg.lstsq(A, uni[seen][:, None]).solution[:, 0]
            bias = torch.where(seen, uni, coef[0] * pre_bias + coef[1])
            self.bias_fit = (float(coef[0]), float(coef[1]))
        bias = bias - bias[4:].mean()
        bias[:4] = -30.0                     # specials are never targets: keep them out of the partition (review: +12 nats)
        self.word_bias = nn.Parameter(bias)
        # never-in-poems rows stay frozen (train --unseen freeze); ONE learnable scalar shifts all of them together, so the
        # model can set how much mass unseen words get without pushing each row down (review 2026-09-29)
        self.unseen_off = nn.Parameter(torch.zeros(()))
        unseen = (~X.word_seen).float(); unseen[:4] = 0.0
        self.register_buffer("unseen_ind", unseen, persistent=False)
        if lm_head is not None:
            self.head_dense.weight.data.copy_(lm_head["dense.weight"]); self.head_dense.bias.data.copy_(lm_head["dense.bias"])
            self.head_ln.weight.data.copy_(lm_head["layer_norm.weight"]); self.head_ln.bias.data.copy_(lm_head["layer_norm.bias"])
        # SAME-LEVEL variant (user 2026-09-29 "maybe they both has to be at the same level"): the sound layer is a peer
        # stream with its own identity embedding (input) and its own head (output) over the pronounced syllables; rows
        # start from the mean WangchanBERTa embedding of the one-syllable lexicon words with that sound. The head is
        # UNTIED from the input embedding (review 2026-09-29: tied, the one-batch overfit stuck at ~3.1 nats on one
        # poem whose masked cells all collapsed to one hidden state, cos 1.0; untied it reached 0.009 in 150 steps)
        self.peer = bool(peer)
        if self.peer:
            self.snd_emb = nn.Embedding(len(X.phons), d)
            if peer_tied: self.snd_bias = nn.Parameter(torch.zeros(len(X.phons)))   # only the (unused) tied variant has it
            with torch.no_grad():
                self.snd_emb.weight.normal_(0, 0.02, generator=torch.Generator().manual_seed(21)); self.snd_emb.weight[0] = 0
                acc = torch.zeros(len(X.phons), d); n_ = torch.zeros(len(X.phons))
                for i, w in enumerate(X.words):
                    if i >= 4 and len(X.sounds[w]) == 1:
                        acc[X.pid[X.sounds[w][0]]] += self.word_emb.weight[i]; n_[X.pid[X.sounds[w][0]]] += 1
                has = n_ > 0; has[0] = False
                self.snd_emb.weight[has] = acc[has] / n_[has, None]
            self.register_buffer("tok_sound", X.tok_sound.clone(), persistent=False)
            if not peer_tied:
                self.snd_out = nn.Linear(d, len(X.phons))
                with torch.no_grad(): self.snd_out.weight.copy_(self.snd_emb.weight); self.snd_out.bias.zero_()
                # the INPUT identity channel starts at ZERO and grows (software-testing: gate new channels at zero): added
                # at full word-embedding scale it collapsed whole poems in the one-batch overfit (every masked cell one
                # hidden state, cos 1.000, on 2 of 8 poems; without the input channel: no collapse)
                with torch.no_grad(): self.snd_emb.weight.zero_()
        for n_, t_ in (("tok_word", X.tok_word), ("tok_kn", X.tok_kn), ("tok_on", X.tok_on), ("tok_ri", X.tok_ri),
                       ("tok_to", X.tok_to), ("p_on", torch.tensor(X.p_on)), ("p_ri", torch.tensor(X.p_ri)), ("p_to", torch.tensor(X.p_to))):
            self.register_buffer(n_, t_, persistent=False)
        self.register_buffer("line_id", torch.tensor([0] * TL + [1 + i for i in range(L) for _ in range(C)]), persistent=False)
        self.register_buffer("pos_id", torch.tensor(list(range(TL)) + [j for _ in range(L) for j in range(C)]), persistent=False)
        self.register_buffer("role_id", torch.tensor([0] * TL + [ROLE[i] for i in range(L) for _ in range(C)]), persistent=False)

    def embed(self, tok, snd, ri_in):
        """tok: word-slot ids; snd: sound id of KNOWN cells (0 = unknown); ri_in: revealed rime at masked cells (0 = none)."""
        known = (tok != MASK) & (snd > 0)
        on = torch.where(known, self.p_on[snd], torch.zeros_like(snd))
        ri = torch.where(known, self.p_ri[snd], ri_in)
        to = torch.where(known, self.p_to[snd], torch.full_like(snd, 5))
        x = (self.word_emb(self.tok_word[tok]) + self.kn_emb(self.tok_kn[tok]) + self.on_emb(on) + self.ri_emb(ri)
             + (self.snd_emb(torch.where(known, snd, torch.zeros_like(snd))) if (self.peer and self.peer_in) else 0)
             + self.to_emb(to) + self.line(self.line_id)[None] + self.pos(self.pos_id)[None] + self.role(self.role_id)[None]
             + (self.abs_pos[None] if self.use_abs_pos else 0))
        return self.ln_in(x)

    def hidden(self, tok, snd, ri_in):
        B = tok.shape[0]
        keep = torch.cat([torch.ones(B, 1, dtype=torch.bool, device=tok.device), tok != PAD], 1)
        ext = (~keep)[:, None, None, :].float() * -1e4                         # PAD is never attended (audit: PAD sink)
        x = torch.cat([self.ln_in(self.bos)[None, None].expand(B, 1, -1), self.embed(tok, snd, ri_in)], 1)
        h = self.body(x, attention_mask=ext).last_hidden_state[:, 1:]
        return self.head_ln(F.gelu(self.head_dense(h)))

    def eff_bias(self):
        return self.word_bias + self.unseen_off * self.unseen_ind

    def token_logits(self, z):
        """z [..., d] -> logits over ALL word-slot tokens [..., NT]"""
        wl = z @ self.word_emb.weight.T + self.eff_bias()                      # [..., W]
        kl = z @ self.kn_emb.weight.T                                          # [..., KN]
        return wl[..., self.tok_word] + kl[..., self.tok_kn]

    def token_nll(self, z, y):
        """exact cross-entropy over all word-slot tokens without building [N, NT]: logsumexp over tokens =
        logsumexp over words of (word logit + logsumexp of the slot logits k < n_w) - the slot term depends only on n_w"""
        wl = z @ self.word_emb.weight.T + self.eff_bias()                      # [N, W]
        kl = z @ self.kn_emb.weight.T                                          # [N, KN]
        if not hasattr(self, "_n_of_word") or self._n_of_word.device != z.device:
            nw = torch.zeros(len(self.X.words), dtype=torch.long).index_reduce_(0, self.X.tok_word, self.X.tok_n, "amax", include_self=False)
            kn_lists = [[self.X.kn_of(k, n) for k in range(n)] for n in range(1, KMAX + 1)]
            self._n_of_word = (nw - 1).to(z.device)
            self._kn_pad = torch.tensor([l + [l[0]] * (KMAX - len(l)) for l in kn_lists], device=z.device)       # [KMAX, KMAX]
            self._kn_valid = torch.tensor([[q < len(l) for q in range(KMAX)] for l in kn_lists], device=z.device)
        slot = kl[:, self._kn_pad].masked_fill(~self._kn_valid[None], float("-inf")).logsumexp(-1)   # [N, KMAX] per word length
        logZ = (wl + slot[:, self._n_of_word]).logsumexp(-1)
        return (logZ - wl.gather(1, self.tok_word[y][:, None])[:, 0] - kl.gather(1, self.tok_kn[y][:, None])[:, 0]).mean()

    def sound_logits(self, z):
        if not self.peer_tied: return self.snd_out(z)
        return z @ self.snd_emb.weight.T + self.snd_bias

    def forward(self, tok, snd, ri_in):
        z = self.hidden(tok, snd, ri_in)
        lg = self.token_logits(z)
        if self.peer:
            # joint score of a consistent (sound, word-slot) pair, a product of experts: log p_word(w,k) + log p_sound(s(w,k)).
            # The sound is NOT counted twice: the two heads are separate experts over the same pair, and Decoder.lp
            # renormalises the sum over all consistent pairs (log_softmax), i.e. p(w,k) ∝ p_word(w,k) p_sound(s(w,k)).
            lg = F.log_softmax(lg, -1) + F.log_softmax(self.sound_logits(z), -1)[..., self.tok_sound]
        return lg, z


SCHEMES = ("word", "cell", "wak", "all")
MIX = (0.4, 0.2, 0.25, 0.15)          # one model for every decode unit: each masking scheme is a state some unit meets


def mask_words(tok, span, gen, lo=0.15, hi=1.0, mix=MIX):
    """Mixed masking of the poem cells, one scheme per poem (so one model serves all four decode units):
      word: whole words masked together, ratio ~ U(lo, hi)        (word unit; refinement states of wak/bot)
      cell: cells masked independently, ratio ~ U(lo, hi)          (syllable unit: part of a word committed)
      wak:  วรรค i..4 fully masked for a random i, earlier วรรค visible  (wak unit, line by line)
      all:  every poem cell masked                                  (first step of every unit; bot unit)"""
    B = tok.shape[0]
    ratio = torch.rand(B, 1, generator=gen) * (hi - lo) + lo
    u_w = torch.rand(B, 64, generator=gen)                                      # one draw per word occurrence
    u_c = torch.rand(B, T, generator=gen)
    sch = torch.multinomial(torch.tensor(mix, dtype=torch.float), B, replacement=True, generator=gen)
    first = torch.randint(0, L, (B, 1), generator=gen)
    line_of = torch.tensor([-1] * TL + [i for i in range(L) for _ in range(C)])[None, :]
    real = (torch.arange(T)[None, :] >= TL) & (tok != PAD)
    mw = torch.gather(u_w, 1, span.clamp(min=0)) < ratio
    mc = u_c < ratio
    mk = line_of >= first
    m = torch.where((sch == 0)[:, None], mw, torch.where((sch == 1)[:, None], mc, torch.where((sch == 2)[:, None], mk, torch.ones_like(mw))))
    return real & m


def batch_loss(model, tok, snd, span, rcells, gen, reveal=0.0, sound_w=1.0, m=None):
    """random draws on the CPU generator (reproducible on any device); every tensor is moved to the model's device
    BEFORE it meets a model buffer (review 2026-09-29: CPU masks x MPS buffers crashed at step 0)"""
    dev = model.word_bias.device
    if m is None: m = mask_words(tok, span, gen)
    rv = rcells & m & (torch.rand(tok.shape, generator=gen) < reveal) if reveal else torch.zeros_like(m)
    tok, snd, m, rv = tok.to(dev), snd.to(dev), m.to(dev), rv.to(dev)
    x = torch.where(m, torch.full_like(tok, MASK), tok)
    s_in = torch.where(m, torch.zeros_like(snd), snd)
    ri_in = torch.where(rv, model.p_ri[snd], torch.zeros_like(snd))
    z = model.hidden(x, s_in, ri_in)
    tgt = m & (tok >= 4)                                                        # UNK / specials not trained as targets
    loss = model.token_nll(z[tgt], tok[tgt])
    out = {"mlm": float(loss)}
    if model.peer:        # the sound stream's own loss at the same masked cells (both streams of a cell are masked together)
        sm = m & (snd > 0)
        ls = F.cross_entropy(model.sound_logits(z[sm]), snd[sm]); out["sound"] = float(ls)
        loss = loss + sound_w * ls
    return loss, out


def load_recs(data_dir, X, limit=0):
    recs = [json.loads(l) for l in open(os.path.join(data_dir, "poems.jsonl"), encoding="utf-8")]
    if limit: recs = recs[:limit]
    enc = [encode(r, X) for r in recs]
    TOK = torch.stack([e[0] for e in enc]); SND = torch.stack([e[1] for e in enc]); SPAN = torch.stack([e[2] for e in enc])
    RC = torch.zeros(len(recs), T, dtype=torch.bool)
    for n, r in enumerate(recs): RC[n, rhyme_cells(r)] = True
    return recs, TOK, SND, SPAN, RC


def build_model(X, peer=False, abs_pos=True, bias_mode="unigram", **kw):
    from transformers import AutoTokenizer, AutoModel, AutoModelForMaskedLM
    wc_tok = AutoTokenizer.from_pretrained(WC); wc = AutoModel.from_pretrained(WC)
    full = AutoModelForMaskedLM.from_pretrained(WC).state_dict()
    lmh = {k[len("lm_head."):]: v.detach().clone() for k, v in full.items() if k.startswith("lm_head.") and not k.startswith("lm_head.decoder")}
    del full
    return TwoLayer(X, wc_tok, wc, lm_head=lmh, peer=peer, abs_pos=abs_pos, bias_mode=bias_mode, **kw)


@torch.no_grad()
def eval_metrics(model, X, TOK, SND, SPAN, RC, idx, batch=32):
    """validation: mixed-masking loss; probability mass on words never seen in the poems (must not collapse to ~0);
    rime-use check: loss at masked rhyme cells of all-masked poems with the rime revealed vs hidden (the sound layer is
    set first at decode; if revealing it does not lower the loss, the channel is unused)"""
    model.eval(); vg = torch.Generator().manual_seed(1); dev = model.word_bias.device
    unseen_tok = (~X.word_seen[X.tok_word]).to(dev); unseen_tok[:4] = False
    tot = n = 0; mass = []; rv = {0: [], 1: []}
    for s_ in range(0, len(idx), batch):
        b = idx[s_:s_ + batch]; tok, snd, span, rc = TOK[b], SND[b], SPAN[b], RC[b]
        l, _ = batch_loss(model, tok, snd, span, rc, vg); tot += float(l) * len(b); n += len(b)
        m = mask_words(tok, span, vg, mix=(0, 0, 0, 1))                          # all masked
        tok, snd, m, rc = tok.to(dev), snd.to(dev), m.to(dev), rc.to(dev)
        x = torch.where(m, torch.full_like(tok, MASK), tok); s_in = torch.where(m, torch.zeros_like(snd), snd)
        tgt = m & (tok >= 4) & rc
        for r in (0, 1):
            ri = torch.where(tgt & (r == 1), model.p_ri[snd], torch.zeros_like(snd))
            z = model.hidden(x, s_in, ri)
            lg = model.token_logits(z[tgt]).float()
            rv[r].append(F.cross_entropy(lg, tok[tgt], reduction="none").cpu())
            if r == 0: mass.append(F.softmax(lg, -1)[:, unseen_tok].sum(-1).cpu())
    rv0, rv1 = torch.cat(rv[0]).mean(), torch.cat(rv[1]).mean()
    return {"valid": tot / max(n, 1), "unseen_mass": float(torch.cat(mass).mean()) if mass else 0.0,
            "rhyme_loss_hidden": float(rv0), "rhyme_loss_revealed": float(rv1)}


def make_opt(model, lr, snd_lr_mult=20.0):
    """AdamW groups: body/embeddings (wd 0.01); biases/positions/offset (no decay); the peer model's sound INPUT channel
    in its own group at snd_lr_mult x lr, no decay (re-review 2026-09-29: zero-initialised at the body's lr it grew to ~1% of
    the word rows' norm in a one-batch run - nearly dead; the user's idea needs sound used as an input peer)"""
    no_decay = {"word_bias", "bos", "abs_pos", "unseen_off"}; own = {"snd_emb.weight"}
    P_ = dict(model.named_parameters())
    groups = [{"params": [p for n_, p in P_.items() if n_ not in no_decay | own], "weight_decay": 0.01, "lr": lr},
              {"params": [p for n_, p in P_.items() if n_ in no_decay], "weight_decay": 0.0, "lr": lr}]
    if "snd_emb.weight" in P_: groups.append({"params": [P_["snd_emb.weight"]], "weight_decay": 0.0, "lr": lr * snd_lr_mult})
    return torch.optim.AdamW(groups, lr=lr)


@torch.no_grad()
def sound_input_use(model, X, TOK, SND, SPAN, RC, idx):
    """peer model: sound-input row norm / word row norm (rows in use) and the valid-loss change when it is zeroed"""
    used = torch.unique(SND[idx][SND[idx] > 0]); dev = model.word_bias.device
    ratio = float(model.snd_emb.weight[used.to(dev)].norm(dim=-1).mean() / model.word_emb.weight[4:].norm(dim=-1).mean())
    full = eval_metrics(model, X, TOK, SND, SPAN, RC, idx)["valid"]
    keep = model.snd_emb.weight.clone(); model.snd_emb.weight.zero_()
    zero = eval_metrics(model, X, TOK, SND, SPAN, RC, idx)["valid"]
    model.snd_emb.weight.copy_(keep)
    return ratio, zero - full


def train(a):
    random.seed(0); torch.manual_seed(0)
    X = Lex(os.path.join(a.data, "lexicon.json"))
    recs, TOK, SND, SPAN, RC = load_recs(a.data, X, a.limit)
    perm = torch.randperm(len(recs), generator=torch.Generator().manual_seed(0))
    nv = min(a.valid, len(recs) // 5); vi, ti = perm[:nv], perm[nv:]
    model = build_model(X, peer=a.peer, bias_mode=a.bias).to(DEV)
    unseen = ~X.word_seen; unseen[:4] = False
    print(f"lexicon {len(X.words) - 4:,} words ({int(unseen.sum()):,} never in the poems), {X.NT:,} word-slot tokens; "
          f"poems {len(recs):,} (train {len(ti):,}, valid {nv}); bias {a.bias}; unseen rows {a.unseen}", flush=True)
    unseen_d = unseen.to(DEV)
    keep_e = model.word_emb.weight.detach()[unseen_d].clone(); keep_b = model.word_bias.detach()[unseen_d].clone()
    keep_e_d, keep_b_d = keep_e, keep_b
    opt = make_opt(model, a.lr, a.snd_lr_mult)
    steps = a.epochs * (len(ti) // a.batch)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=[g["lr"] for g in opt.param_groups], total_steps=max(steps, 1), pct_start=0.06)
    gen = torch.Generator().manual_seed(0); os.makedirs(a.out, exist_ok=True); t0 = time.time()
    json.dump({"lexicon": os.path.abspath(os.path.join(a.data, "lexicon.json")), "peer": a.peer, "peer_tied": False, "bias": a.bias,
               "unseen": a.unseen}, open(os.path.join(a.out, "config.json"), "w"))
    e0 = eval_metrics(model, X, TOK, SND, SPAN, RC, vi[:256]); print(f"== init {e0}", flush=True)
    for ep in range(a.epochs):
        model.train(); order = ti[torch.randperm(len(ti), generator=gen)]
        for it in range(len(ti) // a.batch):
            b = order[it * a.batch:(it + 1) * a.batch]
            loss, o = batch_loss(model, TOK[b], SND[b], SPAN[b], RC[b], gen, a.reveal, a.sound_w)
            if a.unseen == "anchor":
                reg = ((model.word_emb.weight[unseen_d] - keep_e.to(DEV)) ** 2).sum(-1).mean() + ((model.word_bias[unseen_d] - keep_b.to(DEV)) ** 2).mean()
                loss = loss + a.anchor_w * reg; o["anchor"] = float(reg)
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step(); sched.step()
            if a.unseen == "freeze":                  # rows of never-seen words stay at their pretrained init (no push-down)
                with torch.no_grad():
                    model.word_emb.weight[unseen_d] = keep_e_d; model.word_bias[unseen_d] = keep_b_d
            if it % 50 == 0:
                mem = f" mps {torch.mps.driver_allocated_memory() / 1e9:.1f}GB" if DEV.type == "mps" else ""
                print(f"ep {ep} it {it}/{len(ti) // a.batch} {o} {time.time() - t0:.0f}s{mem}", flush=True)
            if a.max_steps and ep * (len(ti) // a.batch) + it + 1 >= a.max_steps: break
        e = eval_metrics(model, X, TOK, SND, SPAN, RC, vi)
        off = float(model.unseen_off)
        alarm = ("HIGH" if e["unseen_mass"] > 0.05 else "LOW" if e["unseen_mass"] < 0.001 else "ok")   # human 0.0054
        bad = [k for k, v in e.items() if not math.isfinite(v)] + ([] if math.isfinite(off) else ["unseen_off"])
        print(f"== epoch {ep} valid masked-loss {e['valid']:.3f} unseen-mass {e['unseen_mass']:.4f} ({alarm}) unseen-offset "
              f"{off:+.3f} rhyme-loss hidden {e['rhyme_loss_hidden']:.3f} revealed {e['rhyme_loss_revealed']:.3f}", flush=True)
        if bad: print(f"STOPPING: non-finite {bad}", flush=True); sys.exit(1)
        if model.peer:
            ratio, delta = sound_input_use(model, X, TOK, SND, SPAN, RC, vi[:256])
            low = ep >= 3 and ratio < 0.10          # prereg: LOW if the sound input is < 10% of the word rows by epoch 4
            print(f"== sound-input epoch {ep}: row norm / word row norm {ratio:.3f}, zero-ablation valid-loss delta "
                  f"{delta:+.4f} nats{' (LOW: sound input < 10% of word norm)' if low else ''}", flush=True)
        torch.save(model.state_dict(), os.path.join(a.out, "model.pt")); print(f"   saved model.pt (epoch {ep})", flush=True)
        if a.max_steps and (ep + 1) * (len(ti) // a.batch) >= a.max_steps: break


# ------------------------------------------------------------------ decoding

# PRONOUNCED syllables per วรรค in human บท (the fixed base's --wak-lens sample table, canvas_bot.WAK_LEN_DIST)
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(rng):
    out = []
    for i in range(L):
        d = WAK_LEN_DIST[i]; r = rng.random() * sum(d.values())
        for k, v in sorted(d.items()):
            r -= v
            if r <= 0: out.append(k); break
        else: out.append(8)
    return out


def title_input(X, text):
    import twolayer_data as TD, title_clean
    cells = TD.line_cells(title_clean.clean_title(text)) or []
    tok = torch.full((T,), PAD, dtype=torch.long); snd = torch.zeros(T, dtype=torch.long); n = 0; words = []
    for (w, k, nw, s) in cells:
        if k == 0 and n + nw > TL: break
        tok[n] = X.tok(w, k); snd[n] = X.pid.get(s.phon, 0); n += 1
        if k == 0: words.append(w)
    return tok, snd, words


class Decoder:
    """Iterative decoding on the two-layer canvas with a selectable COMMIT UNIT (user 2026-09-29: "try 'em all"):
      syllable  one sound cell per step; the word layer still binds it: a cell committed as (w, k) forces the other
                cells of w's span to (w, k') - so a syllable is always part of a real word
      word      one whole word per step, across its sound cells (placement = mean log-prob over its cells)
      wak       one วรรค per step, all its cells jointly: an exact sample of a whole-word tiling of the line (forward-
                filtering backward-sampling over the per-cell log-probs), then Mask-Predict refinement inside the line
      bot       all 4 วรรค jointly in one step (a tiling per line from ONE forward pass), then refinement over the stanza
    Rhyme / tone / count act on the SOUND layer in every unit: count = number of cells; tone masks at line-final cells;
    a line-final cell admits only a word's LAST syllable, a line-initial cell only a FIRST; a rhyme cell admits only
    tokens whose sound has the group's rime family. How the family is set:
      syllable / word: family chosen (after fam_delay steps) from the model's family mass on the group's cells, as the
                current canvas does; rhyme cells are held until then; the chosen rime is written into the sound layer
                (ri_in) of the group's open cells before any word is chosen for them
      wak:      the SENDER's line-final word is sampled inside the line's joint tiling with no family mask (only its
                tone class); that word fixes the family; receivers are later lines, tiled under the family mask, the
                receiving position (3rd or 5th) chosen by the tiling's log-partition under each
      bot:      ONE forward pass from the all-masked canvas; the 4 lines are tiled in order from it, so each sender line
                fixes its group's family (as wak) before its receivers are tiled - no family off the blank canvas
    Refinement (wak, bot): R rounds; each re-masks the lowest-scoring fraction of words (score = mean log-prob of the word
    when it was drawn) and re-tiles the gaps; families stay fixed, so re-drawn rhyme words keep the rhyme."""

    def __init__(self, model, X, title_text, lens=None, unit="word", temp=0.8, fam_temp=0.5, fam_delay=4, topic_boost=1.5,
                 rep_pen=2.0, rhyme_pos="auto", refine=3, gen=None, trace=None):
        self.m, self.X, self.unit, self.temp, self.fam_temp, self.fam_delay = model, X, unit, temp, fam_temp, fam_delay
        self.rep_pen, self.rhyme_pos, self.refine, self.trace = rep_pen, rhyme_pos, refine, trace
        self.gen = gen or torch.Generator().manual_seed(0)
        self.lens = lens or [8] * L
        self.tok, self.snd, twords = title_input(X, title_text)
        self.ri_in = torch.zeros(T, dtype=torch.long)
        for i in range(L):
            for j in range(self.lens[i]): self.tok[self.idx(i, j)] = MASK
        self.committed, self.forced, self.fam, self.gpos = set(), {}, {}, {}
        self.wscore = {}                                     # (line, start) -> (word id, mean log-prob when drawn)
        NT = X.NT; self.NT = NT
        self.base = torch.zeros(NT); self.base[:4] = -1e9
        self.tone_ok = {i: torch.tensor([(t in allowed) for t in X.tok_to.tolist()]) for i, allowed in TONE_AT.items()}
        topic = torch.zeros(NT)
        if topic_boost:
            for w in twords:
                wi = X.wid.get(w, UNK)
                if wi >= 4 and len(w) > 1: topic[X.first[wi]: X.first[wi] + len(X.sounds[w])] = topic_boost
        self.topic_rows = [topic.clone() for _ in range(L)]
        res = lambda g, alt: [(l, self.lens[l] - 1 if c == -1 else (4 if alt and 4 < self.lens[l] - 1 else c)) for (l, c) in g]
        self.groups = [res(g, False) for g in GROUPS]; self.groups_alt = [res(g, True) for g in GROUPS]
        self.n_fwd = 0

    idx = staticmethod(lambda i, j: TL + i * C + j)

    def gcells(self, gi):
        if gi in self.gpos: return (self.groups_alt if self.gpos[gi] else self.groups)[gi]
        return list(dict.fromkeys(self.groups[gi] + (self.groups_alt[gi] if self.rhyme_pos == "auto" else [])))

    def fam_cells(self, gi):
        """cells the family mask / rime channel applies to now: an undecided receiver (3rd or 5th) is left out"""
        if gi in self.gpos or self.rhyme_pos != "auto": return self.gcells(gi)
        return [c for c in self.groups[gi][:-1]]

    def used_words(self, exclude=()):
        return {int(self.X.tok_word[self.tok[k]]) for k in self.committed if k not in exclude and int(self.tok[k]) >= 4}

    # ---------------------------------------------------------------- scores
    def lp(self, hold=True):
        """masked log-probs [T, NT] for the current state"""
        X, tk, tn, tw = self.X, self.X.tok_k, self.X.tok_n, self.X.tok_word
        self.n_fwd += 1
        dev = self.m.word_bias.device                     # the model's own device, so model.to(...) after loading works
        logits, _ = self.m(self.tok[None].to(dev), self.snd[None].to(dev), self.ri_in[None].to(dev))
        lg = logits[0].float().cpu() + self.base[None, :]
        for i in range(L):
            lg[self.idx(i, 0)] += torch.where(tk == 0, 0.0, -1e9)
            lg[self.idx(i, self.lens[i] - 1)] += torch.where(tk == tn - 1, 0.0, -1e9)
            if i in self.tone_ok: lg[self.idx(i, self.lens[i] - 1)] += torch.where(self.tone_ok[i], 0.0, -1e9)
            lg[TL + i * C: TL + i * C + self.lens[i]] += self.topic_rows[i]
        uw = self.used_words()
        if uw: lg[TL:] += torch.isin(tw, torch.tensor(sorted(uw))).float() * -self.rep_pen
        for k, t in self.forced.items():
            keep = lg[k, t].clone(); lg[k] = -1e9; lg[k, t] = keep
        lp = F.log_softmax(lg, -1)               # peer model: renormalises the product log p_word + log p_sound (see forward)
        self.apply_fams(lp)
        for i in range(L): self.ban_group_words(lp, i)
        return lp

    def topk_cut(self, lp, line=None, K=512):
        """the SAME candidate cut in every unit (review 2026-09-29): at each open poem cell only its K best tokens (after
        every mask) stay admissible; applied right before a commit or a tiling, never before a family choice"""
        rows = [self.idx(i, j) for i in range(L) for j in range(self.lens[i])
                if self.idx(i, j) not in self.committed and (line is None or i == line)]
        if not rows: return lp
        r = torch.tensor(rows); sub = lp[r]
        kth = torch.topk(sub, min(K, self.NT), dim=-1).values[:, -1:]
        lp[r] = torch.where(sub >= kth, sub, torch.full_like(sub, -1e9))
        return lp

    def apply_fams(self, lp, line=None):
        """family / early-receiver / same-word-in-group masks for every chosen family (idempotent), optionally one line"""
        tw = self.X.tok_word
        for gi in self.fam:
            cg = self.fam_cells(gi)
            for c in cg:
                k = self.idx(*c)
                if k in self.committed or (line is not None and c[0] != line): continue
                lp[k] += torch.where(self.X.tok_ri == self.fam[gi], 0.0, -1e9)
            if line is None or self.groups[gi][-1][0] == line: self.ban_early_receiver(lp, gi)
        return lp

    def choose_families(self, lp, groups=None):
        X = self.X
        for gi in (groups if groups is not None else range(len(GROUPS))):
            if gi in self.fam: continue
            done = [c for c in self.gcells(gi) if self.idx(*c) in self.committed]
            a, b = self.groups[gi][-1], self.groups_alt[gi][-1]
            if done:
                # a committed cell fixes the family (wak unit: the sender's line was tiled freely); the receiving position
                # is fixed only if a receiver cell is committed - otherwise tile_lines picks it by the partition
                self.fam[gi] = int(X.tok_ri[self.tok[self.idx(*done[0])]])
                if a in done: self.gpos[gi] = 0
                elif b in done and b != a: self.gpos[gi] = 1
            else:
                fams = [f for f in range(1, X.n_ri) if f not in self.fam.values()]
                ft = torch.tensor(fams); sc = torch.zeros(len(fams))
                for c in self.groups[gi]:
                    mass = torch.zeros(X.n_ri).index_add_(0, X.tok_ri, lp[self.idx(*c)].exp())
                    sc += torch.log(mass[ft] + 1e-30)
                pr = F.softmax(sc / len(self.groups[gi]) / max(self.fam_temp, 1e-6), 0)
                self.fam[gi] = fams[int(torch.multinomial(pr, 1, generator=self.gen))]
                if self.rhyme_pos == "auto" and gi not in self.gpos:
                    rm = lambda c: float(lp[self.idx(*c)][X.tok_ri == self.fam[gi]].exp().sum())
                    self.gpos[gi] = 1 if (a != b and rm(b) > rm(a)) else 0
            for c in self.fam_cells(gi):                                     # the sound layer is set first
                if self.idx(*c) not in self.committed: self.ri_in[self.idx(*c)] = self.fam[gi]

    def held(self):
        h = set()
        for gi in range(len(GROUPS)):
            if gi not in self.fam: h |= {self.idx(*c) for c in self.gcells(gi)}
        return h

    # ---------------------------------------------------------------- commits
    def put_word(self, i, s, t_first, score=None):
        """commit the whole word whose FIRST slot token is t_first at cells s.. of line i"""
        X = self.X; w = int(X.tok_word[t_first]); n = int(X.tok_n[t_first])
        for q in range(n):
            k = self.idx(i, s + q); self.tok[k] = t_first + q; self.snd[k] = int(X.tok_sound[t_first + q])
            self.committed.add(k); self.ri_in[k] = 0; self.forced.pop(k, None)
        self.topic_rows[i][X.tok_word == w] = 0.0
        self.wscore[(i, s)] = (w, score if score is not None else 0.0)
        if self.trace is not None: self.trace.append((self.unit, i, s, X.words[w]))

    def span_scores(self, lp, i, j, held, own_only=False):
        """for the tokens (w,k) still admissible at this cell (topk_cut left <= 512): admissible over the word's whole span?
        and the SUM of log-probs over its open cells (own cell only for the syllable unit). Returns full [NT] vectors
        (sum, number of open cells), -1e9 where inadmissible."""
        X = self.X
        a_ = j
        while a_ - 1 >= 0 and (self.idx(i, a_ - 1) not in self.committed or own_only): a_ -= 1
        b_ = j
        while b_ + 1 < self.lens[i] and (self.idx(i, b_ + 1) not in self.committed or own_only): b_ += 1
        c0 = self.idx(i, j)
        ids = (lp[c0] > -1e8).nonzero()[:, 0]
        out = torch.full((self.NT,), -1e9); nout = torch.ones(self.NT)
        if not len(ids): return out, nout
        tk, tn = X.tok_k[ids], X.tok_n[ids]
        s = j - tk; e = s + tn - 1
        ok = (s >= a_) & (e <= b_)
        tot = lp[c0][ids].clone(); cnt = torch.ones(len(ids))
        for d in range(-(KMAX - 1), KMAX):
            if d == 0: continue
            jj = j + d
            inw = ok & (tk + d >= 0) & (tk + d < tn)
            if not bool(inw.any()): continue
            if not (0 <= jj < self.lens[i]): ok = ok & ~inw; continue
            ck = self.idx(i, jj); tt = (ids + d).clamp(0, self.NT - 1)
            if ck in self.committed:                                   # syllable unit: must match what is there
                ok = ok & ~(inw & (tt != int(self.tok[ck])))
                continue
            v = lp[ck][tt]
            ok = ok & ~(inw & ((v < -1e8) | (ck in held)))
            tot = tot + torch.where(inw, v, torch.zeros_like(v)); cnt = cnt + inw.float()
        sc = lp[c0][ids] if own_only else tot
        out[ids] = torch.where(ok, sc, torch.full_like(sc, -1e9)); nout[ids] = 1.0 if own_only else cnt
        return out, nout

    def step_cell_or_word(self, t):
        lp = self.lp()
        if t >= self.fam_delay:
            self.choose_families(lp); self.apply_fams(lp)                  # masks for families chosen this step
            for i in range(L): self.ban_group_words(lp, i)
        self.topk_cut(lp)
        held = self.held()
        best = None
        for i in range(L):
            for j in range(self.lens[i]):
                k = self.idx(i, j)
                if k in self.committed or k in held: continue
                sc, n_open = self.span_scores(lp, i, j, held, own_only=(self.unit == "syllable"))
                # ORDER by the best per-cell mean (confidence, length-neutral); SAMPLE by the sum = the joint log-prob of
                # the cells committed, the same quantity the wak/bot tiling samples by (review: make them consistent)
                conf = float((sc / n_open).max())
                if conf > -1e8 and (best is None or conf > best[0]): best = (conf, i, j, sc, n_open)
        if best is None: return False
        _, i, j, sc, n_open = best
        t_ = int(torch.multinomial(F.softmax(sc / self.temp, 0), 1, generator=self.gen))
        X = self.X; k_, n_ = int(X.tok_k[t_]), int(X.tok_n[t_]); s = j - k_; first = t_ - k_
        if self.unit == "word":
            self.put_word(i, s, first, float(sc[t_] / n_open[t_]))
        else:
            kk = self.idx(i, j); self.tok[kk] = t_; self.snd[kk] = int(X.tok_sound[t_]); self.committed.add(kk)
            self.ri_in[kk] = 0; self.forced.pop(kk, None)
            for q in range(n_):
                c = self.idx(i, s + q)
                if c not in self.committed: self.forced[c] = first + q
            if all(self.idx(i, s + q) in self.committed for q in range(n_)):
                self.wscore[(i, s)] = (int(X.tok_word[t_]), float(sc[t_] / n_open[t_]))
            if self.trace is not None: self.trace.append(("syllable", i, j, X.words[int(X.tok_word[t_])], k_))
        return True

    # ---------------------------------------------------------------- joint tiling (wak / bot)
    def word_end_scores(self, lp, i, j, a_):
        """[W] sum of log-probs of each word ending at cell j of line i with its start >= a_ (else -inf), and n per word"""
        X = self.X; tk, tn = X.tok_k, X.tok_n
        cell = j - (tn - 1) + tk                                          # where each token would sit
        ok = (cell >= a_) & (cell <= j)
        v = torch.full((self.NT,), 0.0)
        cols = cell.clamp(a_, j)
        rows = torch.stack([lp[self.idx(i, c)] for c in range(a_, j + 1)])  # [j-a_+1, NT]
        v = rows[(cols - a_), torch.arange(self.NT)]
        v = torch.where(ok, v, torch.full_like(v, -1e9))
        W = len(X.words)
        ws = torch.zeros(W).index_add_(0, X.tok_word, v)                  # sum over the word's slots
        ws[:4] = -1e9
        nw = torch.zeros(W, dtype=torch.long).index_reduce_(0, X.tok_word, X.tok_n, "amax", include_self=False)
        # a word whose span does not fit gets a -1e9 from at least one slot; clip so sums stay "impossible"
        ws = torch.where(ws < -1e8, torch.full_like(ws, -1e9), ws)
        return ws, nw

    def tile_run(self, lp, i, a_, b_):
        """forward-filtering backward-sampling of a whole-word tiling of open cells a_..b_ of line i at temperature temp.
        returns (log partition, [(start, first_token, mean log-prob)])"""
        X = self.X
        ends = {}
        for j in range(a_, b_ + 1): ends[j] = self.word_end_scores(lp, i, j, a_)
        alpha = {a_ - 1: 0.0}
        for j in range(a_, b_ + 1):
            ws, nw = ends[j]
            terms = []
            for n in range(1, min(KMAX, j - a_ + 1) + 1):
                if (j - n) not in alpha or alpha[j - n] < -1e8: continue
                sel = ws[nw == n]
                if sel.numel(): terms.append(alpha[j - n] + torch.logsumexp(sel / self.temp, 0))
            alpha[j] = float(torch.logsumexp(torch.stack(terms), 0)) if terms else -1e9
        if alpha[b_] < -1e8: return -1e9, None
        out, j = [], b_
        while j >= a_:
            ws, nw = ends[j]
            cand = torch.full_like(ws, -1e9)
            for n in range(1, min(KMAX, j - a_ + 1) + 1):
                if alpha.get(j - n, -1e9) < -1e8: continue
                m_ = nw == n
                cand = torch.where(m_, alpha[j - n] + ws / self.temp, cand)
            w = int(torch.multinomial(F.softmax(cand, 0), 1, generator=self.gen)); n = int(nw[w])
            out.append((j - n + 1, X.first[w], float(ws[w]) / n)); j -= n
        return alpha[b_], out[::-1]

    def ban_early_receiver(self, lp, gi, alt=None):
        """receiver at the 5th syllable: the 3rd may not carry the family's rime, else the checker lands the rhyme on the
        3rd (ชิงสัมผัส, or ซ้ำคำสัมผัส when it is the sender's word) - smoke 2026-09-29, wak unit"""
        alt = self.gpos.get(gi) if alt is None else alt
        a, b = self.groups[gi][-1], self.groups_alt[gi][-1]
        if alt == 1 and a != b and self.idx(*a) not in self.committed:
            lp[self.idx(*a)] += torch.where(self.X.tok_ri == self.fam[gi], -1e9, 0.0)

    def ban_group_words(self, lp, i):
        """no word twice inside a rhyme group (ซ้ำคำสัมผัส): cells of line i in a group cannot take a word another cell of
        that group already holds - re-applied line by line, since bot tiles all lines from ONE forward pass"""
        tw = self.X.tok_word
        for gi in range(len(GROUPS)):
            cg = self.gcells(gi)
            words = {int(tw[self.tok[self.idx(*o)]]) for o in cg if self.idx(*o) in self.committed}
            if not words: continue
            bad = torch.isin(tw, torch.tensor(sorted(words)))
            for c in cg:
                if c[0] == i and self.idx(*c) not in self.committed: lp[self.idx(*c)] += torch.where(bad, -1e9, 0.0)

    def tile_lines(self, lines):
        """one forward pass, then a joint tiling of every open run in the given lines"""
        lp = self.lp(); placed = set()
        for i in lines:
            self.apply_fams(lp, i); self.ban_group_words(lp, i)
            rows = slice(TL + i * C, TL + i * C + self.lens[i])
            if placed:     # the repetition penalty lp() gave committed words, now also for words placed earlier in this pass
                lp[rows] += torch.isin(self.X.tok_word, torch.tensor(sorted(placed))).float() * -self.rep_pen
            for gi in range(len(GROUPS)):
                snd_c = self.groups[gi][0]
                if gi not in self.fam and snd_c[0] == i and self.fam:
                    # a new sender picks a family no other group uses (humans: 3.87 of 4 families distinct)
                    lp[self.idx(*snd_c)] += torch.isin(self.X.tok_ri, torch.tensor(sorted(set(self.fam.values())))).float() * -1e9
            runs, j = [], 0
            while j < self.lens[i]:
                if self.idx(i, j) in self.committed: j += 1; continue
                a_ = j
                while j + 1 < self.lens[i] and self.idx(i, j + 1) not in self.committed: j += 1
                runs.append((a_, j)); j += 1
            # receiving position (3rd or 5th) still open for a group with a family: pick by the tiling's partition
            for gi in list(self.fam):
                if gi in self.gpos or self.rhyme_pos != "auto": continue
                a, b = self.groups[gi][-1], self.groups_alt[gi][-1]
                if a[0] != i: continue
                best = None
                for alt, c in ((0, a), (1, b)):
                    lp2 = lp.clone(); k = self.idx(*c)
                    lp2[k] += torch.where(self.X.tok_ri == self.fam[gi], 0.0, -1e9)
                    self.ban_early_receiver(lp2, gi, alt); self.topk_cut(lp2, i)
                    Z = sum(self.tile_run(lp2, i, a_, b_)[0] for a_, b_ in runs)
                    if best is None or Z > best[0]: best = (Z, alt)
                self.gpos[gi] = best[1]
                k = self.idx(*self.gcells(gi)[-1]); lp[k] += torch.where(self.X.tok_ri == self.fam[gi], 0.0, -1e9)
                self.ban_early_receiver(lp, gi)
            self.topk_cut(lp, i)                                            # after every mask of this line
            for a_, b_ in runs:
                Z, words = self.tile_run(lp, i, a_, b_)
                if words is None: return False
                for s, first, sc in words: self.put_word(i, s, first, sc); placed.add(int(self.X.tok_word[first]))
            # a sender line just tiled fixes its group's family (wak unit)
            for gi in range(len(GROUPS)):
                if gi not in self.fam and any(self.idx(*c) in self.committed for c in self.gcells(gi)):
                    self.choose_families(lp, [gi])
        return True

    def remask(self, lines, frac):
        """re-open the lowest-scoring fraction of the words in these lines (whole words)"""
        cand = sorted((sc, i, s, w) for (i, s), (w, sc) in self.wscore.items() if i in lines)
        n = int(len(cand) * frac)
        for sc, i, s, w in cand[:n]:
            for q in range(len(self.X.sounds[self.X.words[w]])):
                k = self.idx(i, s + q); self.committed.discard(k); self.tok[k] = MASK; self.snd[k] = 0
                for gi in self.fam:                                      # a re-opened rhyme cell shows its family again
                    if (i, s + q) in self.fam_cells(gi): self.ri_in[k] = self.fam[gi]
            del self.wscore[(i, s)]
        return n

    def refine_lines(self, lines):
        for r in range(1, self.refine + 1):
            frac = 0.5 * (self.refine - r + 1) / (self.refine + 1)       # e.g. R=3: 0.375, 0.25, 0.125
            if self.remask(lines, frac) and not self.tile_lines(lines): return False
        return True

    # ---------------------------------------------------------------- run
    def run(self):
        ok = True
        if self.unit in ("syllable", "word"):
            for t in range(200):
                if not any(self.idx(i, j) not in self.committed for i in range(L) for j in range(self.lens[i])): break
                if not self.step_cell_or_word(t) and t >= self.fam_delay: ok = False; break
        elif self.unit == "wak":
            for i in range(L):
                ok = ok and self.tile_lines([i]) and self.refine_lines([i])
        elif self.unit == "bot":
            # ONE forward pass; the lines are tiled in order from it, so each sender line fixes its group's family before
            # the receiving lines are tiled (review: no family chosen off the all-masked canvas)
            ok = self.tile_lines(list(range(L))) and self.refine_lines(list(range(L)))
        return self.text(), ok

    def text(self):
        X = self.X; lines = []
        for i in range(L):
            ws, j = [], 0
            while j < self.lens[i]:
                t_ = int(self.tok[self.idx(i, j)])
                if t_ < 4: ws.append("?"); j += 1; continue
                ws.append(X.words[int(X.tok_word[t_])]); j += int(X.tok_n[t_]) - int(X.tok_k[t_])
            lines.append("".join(ws))
        return lines

    def placed(self):
        return sorted((i, s, self.X.words[w]) for (i, s), (w, _) in self.wscore.items())


@torch.no_grad()
def decode_one(model, X, title_text, unit="word", retries=4, lens=None, **kw):
    """decode with retries until klon8.prosody reads back exactly the sound cells (the checker re-cuts the joined words)"""
    for r in range(retries + 1):
        D = Decoder(model, X, title_text, lens=lens, unit=unit, **kw)
        lines, ok = D.run()
        want = [[X.phons[int(X.tok_sound[int(D.tok[D.idx(i, j)])])] for j in range(D.lens[i])] for i in range(L)]
        got = [[s.phon for s in P.g2p_line(l)] for l in lines]
        if ok and got == want: return lines, D, r
        kw["gen"] = D.gen                                                 # continue the same RNG stream
    return lines, D, r


def resolve_lexicon(path):
    """the model config stores the lexicon as an absolute path of the machine it was trained on. Used as is when it exists
    (every existing job); otherwise $KLON8_LEXICON, then the same path relative to this repository (the part after the last
    '/klon8/'), then runs/TWOLAYER/data/lexicon.json"""
    if os.path.exists(path): return path
    repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    cands = [os.environ.get("KLON8_LEXICON", "")]
    if "/klon8/" in path: cands.append(os.path.join(repo, path.rsplit("/klon8/", 1)[1]))
    if not os.path.isabs(path): cands.append(os.path.join(repo, path))
    cands.append(os.path.join(repo, "runs", "TWOLAYER", "data", "lexicon.json"))
    for c in cands:
        if c and os.path.exists(c): return c
    raise FileNotFoundError(f"lexicon not found: {path} (tried {[c for c in cands if c]}); set KLON8_LEXICON")


def load(run):
    cfg = json.load(open(os.path.join(run, "config.json")))
    X = Lex(resolve_lexicon(cfg["lexicon"])); model = build_model(X, peer=cfg.get("peer", False), bias_mode=cfg.get("bias", "unigram"),
                                            peer_tied=cfg.get("peer_tied", False))
    model.load_state_dict(torch.load(os.path.join(run, "model.pt"), map_location="cpu")); return model.to(DEV).eval(), X


if __name__ == "__main__":
    ap = argparse.ArgumentParser(); sub = ap.add_subparsers(dest="cmd")
    t = sub.add_parser("train"); t.add_argument("--data", required=True); t.add_argument("--out", required=True)
    t.add_argument("--epochs", type=int, default=12); t.add_argument("--batch", type=int, default=32); t.add_argument("--lr", type=float, default=3e-5)
    t.add_argument("--valid", type=int, default=1000); t.add_argument("--limit", type=int, default=0)
    t.add_argument("--reveal", type=float, default=0.5, help="prob. of showing the rime at a masked rhyme cell (0 = off)")
    t.add_argument("--snd-lr-mult", type=float, default=20.0, help="peer model: lr multiplier of the sound input channel's own group")
    t.add_argument("--peer", action="store_true", help="SAME-LEVEL model: sound stream with its own embedding and head, joint (sound, word) scoring")
    t.add_argument("--sound-w", type=float, default=1.0, help="weight of the sound-head loss (--peer only)")
    t.add_argument("--bias", choices=["hybrid", "pretrained", "unigram"], default="unigram", help="output bias init (see TwoLayer); unigram won the pre-committed step-0 loss rule (11.00 vs hybrid 11.73, pretrained 12.08)")
    t.add_argument("--unseen", choices=["train", "freeze", "anchor"], default="freeze", help="rows of words never seen in the poems: trained freely, frozen at the pretrained init, or L2-anchored to it")
    t.add_argument("--anchor-w", type=float, default=0.1); t.add_argument("--max-steps", type=int, default=0, help="stop after N steps (smoke tests)")
    d = sub.add_parser("decode"); d.add_argument("--model", required=True); d.add_argument("--prompts", required=True); d.add_argument("--out", required=True)
    d.add_argument("--seed", type=int, default=0); d.add_argument("--temp", type=float, default=0.8); d.add_argument("--limit", type=int, default=0)
    d.add_argument("--unit", choices=["syllable", "word", "wak", "bot"], default="word", help="commit unit per decoding step")
    d.add_argument("--refine", type=int, default=3, help="Mask-Predict refinement rounds for --unit wak/bot")
    d.add_argument("--topic-boost", type=float, default=1.5, help="bonus on the title's words (= the fixed base's effective short-topic boost)")
    a = ap.parse_args()
    if a.cmd == "train": train(a)
    elif a.cmd == "decode":
        model, X = load(a.model); rows = [json.loads(l) for l in open(a.prompts, encoding="utf-8")]
        if a.limit: rows = rows[:a.limit]
        rng_len = random.Random(a.seed)
        with open(a.out + ".tmp", "w", encoding="utf-8") as f:
            for n, r in enumerate(rows):
                text = r.get("text") or r.get("prompt") or r.get("title") or r.get("topic")   # dev49 prompts use "text"
                assert text, f"prompt row without text: {r}"
                want = pick_lens(rng_len)
                lines, D, tries = decode_one(model, X, text, unit=a.unit, temp=a.temp, lens=want, refine=a.refine,
                                             topic_boost=a.topic_boost, gen=torch.Generator().manual_seed(a.seed * 100003 + n))
                poem = "\n".join(lines); sc = P.score(poem, n_bot_expected=1)
                got = [len(P.g2p_line(x)) for x in lines]
                f.write(json.dumps({"id": r.get("id"), "prompt": text, "sample": 0, "poem": poem, "unit": a.unit, "syl_target": want,
                                    "syl": got, "retries": tries, "forward_passes": D.n_fwd, "all_pass": sc["all_pass"],
                                    "structure": sc["structure"], "ext_rhyme": sc["ext_rhyme"], "tone": sc["tone"],
                                    "internal": sc["internal"]}, ensure_ascii=False) + "\n"); f.flush()
                print(f"{n}/{len(rows)} all_pass={sc['all_pass']} syl={got} want={want} retries={tries}", file=sys.stderr, flush=True)
        os.replace(a.out + ".tmp", a.out)
        P.save_cache()