| |
| |
| """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: |
| 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) |
| 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 ``<s>``).""" |
| 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) |
| 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 ``<s> t1 .. tn-1``, output ``t1 .. tn``) padded with ``<pad>``.""" |
| 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: |
| """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) |
| ce = F.cross_entropy(logits.transpose(1, 2), out, ignore_index=PAD_ID, |
| label_smoothing=smoothing, reduction="none") |
| 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: |
| """``temp``-squared scaled KL(teacher || student) at the ``valid`` positions ``[B, L]``; |
| pieces the teacher cannot emit (``<s>``, ``<pad>``) 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: |
| """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 |
|
|
