HAKO upload: hako/nlp/nlg_head.py

hako/nlp/nlg_head.py

@@ -0,0 +1,91 @@
+"""NLG head: tiny causal transformer decoder over HAKO's byte-BPE tokens,
+conditioned on the HAKO fused latent z_final (the "routed thinking state").
+
+Theorem T-LM (well-posedness). The decoder is a 2-layer pre-norm transformer
+with GELU and dropout; all parameters are Lipschitz on any bounded set, the
+softmax output is a probability vector, and the cross-entropy loss over the
+finite vocab V is bounded below by 0 and L-smooth in the logits; SGD with
+clipped gradients and Robbins-Monro rates therefore has a.s. convergent loss
+values (bounded-below supermartingale argument). Conditioning: the HAKO
+latent enters as a prefix KV-free additive bias on the first position
+(ContextFormer-style), keeping the causal mask exact for LM training.
+QED.
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Block(nn.Module):
+    def __init__(self, d: int, heads: int, dropout: float) -> None:
+        super().__init__()
+        self.ln1 = nn.LayerNorm(d)
+        self.attn = nn.MultiheadAttention(d, heads, dropout=dropout,
+                                          batch_first=True)
+        self.ln2 = nn.LayerNorm(d)
+        self.mlp = nn.Sequential(nn.Linear(d, 2 * d), nn.GELU(),
+                                 nn.Dropout(dropout), nn.Linear(2 * d, d))
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, causal_mask: torch.Tensor) -> torch.Tensor:
+        h = self.ln1(x)
+        a, _ = self.attn(h, h, h, attn_mask=causal_mask, need_weights=False)
+        x = x + self.drop(a)
+        x = x + self.drop(self.mlp(self.ln2(x)))
+        return x
+
+
+class NLGHead(nn.Module):
+    def __init__(self, vocab: int, N_dim: int = 256, d: int = 256,
+                 layers: int = 2, heads: int = 4, seq_len: int = 96,
+                 dropout: float = 0.1, seed: int = 7) -> None:
+        super().__init__()
+        self.d = d
+        self.seq_len = seq_len
+        self.tok = nn.Embedding(vocab, d)
+        self.pos = nn.Embedding(seq_len, d)
+        self.blocks = nn.ModuleList([Block(d, heads, dropout)
+                                     for _ in range(layers)])
+        self.ln_f = nn.LayerNorm(d)
+        self.head = nn.Linear(d, vocab, bias=False)
+        self.cond = nn.Linear(N_dim, d)     # HAKO latent -> context bias
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, ids: torch.Tensor,
+                z_cond: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, L = ids.shape
+        mask = torch.triu(torch.full((L, L), float("-inf")), diagonal=1)
+        x = self.drop(self.tok(ids) + self.pos(torch.arange(L, device=ids.device)))
+        if z_cond is not None:
+            x = x + self.cond(z_cond).view(B, 1, self.d)
+        for blk in self.blocks:
+            x = blk(x, mask)
+        return self.head(self.ln_f(x))      # (B, L, V) logits
+
+    def loss(self, ids: torch.Tensor,
+             z_cond: Optional[torch.Tensor] = None) -> torch.Tensor:
+        logits = self.forward(ids[:, :-1], z_cond)
+        return F.cross_entropy(logits.reshape(-1, logits.size(-1)),
+                               ids[:, 1:].reshape(-1), ignore_index=-100)
+
+    @torch.no_grad()
+    def generate(self, prompt_ids: list, z_cond: Optional[torch.Tensor],
+                 max_new: int = 40, temperature: float = 0.9,
+                 top_k: int = 50) -> list:
+        ids = torch.tensor([prompt_ids[-self.seq_len:]], dtype=torch.long)
+        for _ in range(max_new):
+            logits = self.forward(ids[:, -self.seq_len:], z_cond)[0, -1]
+            logits = logits / max(temperature, 1e-3)
+            if top_k > 0:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[-1]] = float("-inf")
+            probs = torch.softmax(logits, dim=-1)
+            nxt = int(torch.multinomial(probs, 1))
+            ids = torch.cat([ids, torch.tensor([[nxt]], dtype=torch.long)], 1)
+            if nxt == 0:
+                break
+        return ids[0].tolist()