# SPDX-License-Identifier: Apache-2.0 # Copyright 2026 ayutaz """Plan A-prime: a tiny decoder that writes Action JSON pieces from frame features (docs/02 8.4). encoder: frames [B, T, L, 320] -> softmax layer mix -> LayerNorm -> Linear 320->128 -> causal separable Conv1d(k=5) x1 decoder: 2 layers, d=128, 4 heads, FFN 384, learned positions (max 24), causal self-attention and cross-attention over the valid steps; output over the 62 fixed pieces Decoding is grammar-constrained greedy (``pieces.ActionGrammar``); ``min_prob`` is the smallest unconstrained probability of a chosen piece, as in JTALM, so a forced piece lowers it; ``p_act`` is 1 - P(``[]``) at the first step (E04: the better rejection score). """ import torch from torch import nn from torch.nn import functional as F from koeact.action import parse_output from koeact.model.head_a import CausalSepConv from koeact.model.pieces import ( BOS_ID, EOS_ID, MAX_TOKENS, PAD_ID, VOCAB, ActionGrammar, decode_text, ) from koeact.model.slots import Decoded class DecoderA2(nn.Module): def __init__(self, n_layers: int = 2, d_in: int = 320, d: int = 128, n_heads: int = 4, ffn: int = 384, depth: int = 2, dropout: float = 0.1) -> None: # fmt: skip super().__init__() self.d_model = d self.mix = nn.Parameter(torch.zeros(n_layers)) self.norm = nn.LayerNorm(d_in) self.proj = nn.Linear(d_in, d) self.conv = CausalSepConv(d, dropout=dropout) self.embed = nn.Embedding(len(VOCAB), d, padding_idx=PAD_ID) self.pos = nn.Embedding(MAX_TOKENS + 1, d) layer = nn.TransformerDecoderLayer(d, n_heads, ffn, dropout, batch_first=True, norm_first=True) # fmt: skip self.decoder = nn.TransformerDecoder(layer, depth) self.out_norm = nn.LayerNorm(d) self.out = nn.Linear(d, len(VOCAB)) def encode(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: h = (x * torch.softmax(self.mix, 0)[None, None, :, None]).sum(2) h = self.proj(self.norm(h)) * mask[..., None] return self.conv(h) * mask[..., None] def step_logits(self, memory, mask, inp: torch.Tensor) -> torch.Tensor: """Logits ``[B, L, V]`` for the input pieces ``inp`` (starting with ````).""" length = inp.shape[1] y = self.embed(inp) + self.pos(torch.arange(length, device=inp.device))[None] causal = torch.ones(length, length, dtype=torch.bool, device=inp.device).triu(1) y = self.decoder(y, memory, tgt_mask=causal, tgt_is_causal=True, tgt_key_padding_mask=inp == PAD_ID, memory_key_padding_mask=~mask) # fmt: skip return self.out(self.out_norm(y)) def forward(self, x, mask, inp): return self.step_logits(self.encode(x, mask), mask, inp) def teacher_forcing(targets: list[list[int]]) -> tuple[torch.Tensor, torch.Tensor]: """(input `` t1 .. tn-1``, output ``t1 .. tn``) padded with ````.""" length = max(len(t) for t in targets) inp = torch.full((len(targets), length), PAD_ID, dtype=torch.long) out = torch.full((len(targets), length), PAD_ID, dtype=torch.long) for i, t in enumerate(targets): inp[i, : len(t)] = torch.tensor([BOS_ID, *t[:-1]]) out[i, : len(t)] = torch.tensor(t) return inp, out def loss_fn(logits: torch.Tensor, out: torch.Tensor, smoothing: float = 0.05, weights: torch.Tensor | None = None) -> torch.Tensor: # fmt: skip """Token cross-entropy; ``weights`` ``[B]`` scales each row's tokens (weighted mean).""" if weights is None: return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), out.reshape(-1), ignore_index=PAD_ID, label_smoothing=smoothing) # fmt: skip ce = F.cross_entropy(logits.transpose(1, 2), out, ignore_index=PAD_ID, label_smoothing=smoothing, reduction="none") # fmt: skip w = weights[:, None] * (out != PAD_ID) return (ce * w).sum() / w.sum() class Teacher: """JTALM-kana piece logits and pooled hidden states per D01 text (``scripts/kd_teacher.py``).""" def __init__(self, path) -> None: import numpy as np data = np.load(path) self.row = {str(i): k for k, i in enumerate(data["ids"])} self.logits = torch.from_numpy(data["logits"].astype(np.float32)) self.hidden = torch.from_numpy(data["hidden"].astype(np.float32)) self.lengths = data["lengths"] def kd_loss(logits: torch.Tensor, teacher: torch.Tensor, valid: torch.Tensor, temp: float) -> torch.Tensor: # fmt: skip """``temp``-squared scaled KL(teacher || student) at the ``valid`` positions ``[B, L]``; pieces the teacher cannot emit (````, ````) have -inf teacher logits.""" t = F.log_softmax(teacher[valid] / temp, -1) s = F.log_softmax(logits[valid] / temp, -1) kl = torch.where(t.isfinite(), t.exp() * (t - s), torch.zeros_like(s)).sum(-1) return kl.mean() * temp**2 def align_loss(memory: torch.Tensor, mask: torch.Tensor, proj: nn.Module, hidden: torch.Tensor) -> torch.Tensor: # fmt: skip """1 - cosine between the projected mean of the valid encoder steps and the teacher's pooled hidden state (the Phase 3 alignment of docs/02 8.5).""" pooled = (memory * mask[..., None]).sum(1) / mask.sum(1, keepdim=True).clamp(min=1) return (1 - F.cosine_similarity(proj(pooled), hidden, dim=-1)).mean() @torch.no_grad() def greedy(model: DecoderA2, x, mask, grammar: ActionGrammar) -> list[Decoded]: memory = model.encode(x, mask) batch = x.shape[0] inp = torch.full((batch, 1), BOS_ID, dtype=torch.long, device=x.device) done = torch.zeros(batch, dtype=torch.bool, device=x.device) min_prob = torch.ones(batch, device=x.device) p_act = None for _ in range(MAX_TOKENS): probs = model.step_logits(memory, mask, inp)[:, -1].float().softmax(-1) if p_act is None: p_act = 1 - probs[:, VOCAB.index("[]")] allowed = torch.full_like(probs, -1.0) for row in range(batch): if not done[row]: ids = grammar.allowed(inp[row, 1:].tolist()) allowed[row, ids] = probs[row, ids] nxt = allowed.argmax(-1) p = probs.gather(1, nxt[:, None]).squeeze(1) min_prob = torch.where(done, min_prob, torch.minimum(min_prob, p)) nxt = torch.where(done, torch.full_like(nxt, EOS_ID), nxt) inp = torch.cat([inp, nxt[:, None]], 1) done |= nxt == EOS_ID if bool(done.all()): break out = [] for row in range(batch): calls = parse_output(decode_text(inp[row, 1:].tolist())).calls or [] out.append(Decoded(calls, float(min_prob[row]), float(p_act[row]))) return out