"""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)