File size: 1,912 Bytes
f930dac | 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 | """Transcript-window baseline: "the context window is the mind".
A causal decoder-only transformer over the last W transcript tokens (the same
per-event token stream PNSR sees, concatenated), parameter-matched to PNSR,
with the IDENTICAL factorized head system and cache record bank. TXE is the
same class with W = one event (current-event-only trained null).
"""
from __future__ import annotations
from dataclasses import dataclass
import torch
import torch.nn as nn
from .modules import Block, HeadSystem, RecordEncoder
@dataclass
class TXConfig:
vocab: int = 8192
d: int = 384
heads: int = 8
layers: int = 10 # param-matched to PNSR (see PNSRConfig.enc_layers note)
window: int = 768
class TX(nn.Module):
def __init__(self, cfg: TXConfig):
super().__init__()
self.cfg = cfg
d = cfg.d
self.tok_emb = nn.Embedding(cfg.vocab, d)
self.pos_emb = nn.Parameter(torch.zeros(cfg.window, d))
self.blocks = nn.ModuleList(Block(d, cfg.heads) for _ in range(cfg.layers))
self.ln = nn.LayerNorm(d)
self.recenc = RecordEncoder(d, self.tok_emb)
self.heads = HeadSystem(d)
mask = torch.triu(torch.ones(cfg.window, cfg.window, dtype=torch.bool), 1)
self.register_buffer("causal", mask, persistent=False)
def forward(self, tok, rec_bank, live_mask):
"""tok: [B,W] left-padded with 0; final position is the query point."""
B, W = tok.shape
pad = tok == 0
x = self.tok_emb(tok.long()) + self.pos_emb[:W]
am = self.causal[:W, :W]
for blk in self.blocks:
x = blk(x, attn_mask=am, self_padding_mask=pad)
h = self.ln(x[:, -1]) # final token = query position
return self.heads(None, h, None, rec_bank, live_mask)
def param_report(self):
return {"total": sum(p.numel() for p in self.parameters())}
|