| """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 |
| 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]) |
| return self.heads(None, h, None, rec_bank, live_mask) |
|
|
| def param_report(self): |
| return {"total": sum(p.numel() for p in self.parameters())} |
|
|
