File size: 6,920 Bytes
504fb6d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | # 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 ``<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) # 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 ``<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: # 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 (``<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: # 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
|