"""WriteFix: a small text-correction network, written from scratch in PyTorch. Messy text -> [conv stem -> Transformer encoder] -> [Transformer decoder] -> clean text. No pretrained weights anywhere: every parameter starts random. """ import math import torch import torch.nn as nn import torch.nn.functional as F PAD, BOS, EOS, UNK, FIX = 0, 1, 2, 3, 4 N_SPECIAL = 40 # ids 4..39 are control tokens; only is trained in v1, the rest are reserved for tone modes SIZES = { "tiny": dict(d=256, heads=4, ff=1024, enc=5, dec=2), # ~10M parameters: trains in a few hours on a laptop "small": dict(d=384, heads=6, ff=1536, enc=6, dec=3), # ~24M parameters "base": dict(d=512, heads=8, ff=2048, enc=8, dec=4), # ~50M "big": dict(d=640, heads=10, ff=2560, enc=10, dec=5), # ~95M } class ConvBlock(nn.Module): """Residual 1-D convolution over neighbouring tokens. A misspelt word breaks into small pieces; this lets the encoder glue those pieces back together before attention runs.""" def __init__(self, d, k=5, drop=0.1): super().__init__() self.norm = nn.LayerNorm(d) self.depth = nn.Conv1d(d, d, k, padding=k // 2, groups=d) self.point = nn.Conv1d(d, d, 1) self.drop = nn.Dropout(drop) def forward(self, x, keep): # x: B,T,D keep: B,T,1 (1 = real token) h = (self.norm(x) * keep).transpose(1, 2) h = self.point(F.gelu(self.depth(h))).transpose(1, 2) return x + self.drop(h) class Attention(nn.Module): def __init__(self, d, heads, drop): super().__init__() self.h, self.drop = heads, drop self.q, self.k, self.v, self.o = (nn.Linear(d, d) for _ in range(4)) def split(self, x): B, T, D = x.shape return x.view(B, T, self.h, D // self.h).transpose(1, 2) def forward(self, x, mem=None, mask=None, cache=None): q = self.split(self.q(x)) if mem is None: # self-attention (optionally growing a key/value cache while decoding) k, v = self.split(self.k(x)), self.split(self.v(x)) if cache is not None: if "k" in cache: k, v = torch.cat([cache["k"], k], 2), torch.cat([cache["v"], v], 2) cache["k"], cache["v"] = k, v else: # cross-attention over the encoder output (keys/values computed once) if cache is not None and "k" in cache: k, v = cache["k"], cache["v"] else: k, v = self.split(self.k(mem)), self.split(self.v(mem)) if cache is not None: cache["k"], cache["v"] = k, v y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.drop if self.training else 0.0) B, _, T, _ = y.shape return self.o(y.transpose(1, 2).reshape(B, T, -1)) class FF(nn.Module): def __init__(self, d, ff, drop): super().__init__() self.a, self.b, self.drop = nn.Linear(d, ff), nn.Linear(ff, d), nn.Dropout(drop) def forward(self, x): return self.b(self.drop(F.gelu(self.a(x)))) class EncLayer(nn.Module): def __init__(self, d, heads, ff, drop): super().__init__() self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d) self.att, self.ff, self.drop = Attention(d, heads, drop), FF(d, ff, drop), nn.Dropout(drop) def forward(self, x, mask): x = x + self.drop(self.att(self.n1(x), mask=mask)) return x + self.drop(self.ff(self.n2(x))) class DecLayer(nn.Module): def __init__(self, d, heads, ff, drop): super().__init__() self.n1, self.n2, self.n3 = nn.LayerNorm(d), nn.LayerNorm(d), nn.LayerNorm(d) self.self_att, self.cross = Attention(d, heads, drop), Attention(d, heads, drop) self.ff, self.drop = FF(d, ff, drop), nn.Dropout(drop) def forward(self, x, mem, self_mask, mem_mask, cache=None): c1 = cache.setdefault("self", {}) if cache is not None else None c2 = cache.setdefault("cross", {}) if cache is not None else None x = x + self.drop(self.self_att(self.n1(x), mask=self_mask, cache=c1)) x = x + self.drop(self.cross(self.n2(x), mem=mem, mask=mem_mask, cache=c2)) return x + self.drop(self.ff(self.n3(x))) class WriteFix(nn.Module): def __init__(self, vocab, d=384, heads=6, ff=1536, enc=6, dec=3, drop=0.1, max_len=320): super().__init__() self.cfg = dict(vocab=vocab, d=d, heads=heads, ff=ff, enc=enc, dec=dec, drop=drop, max_len=max_len) self.d, self.max_len = d, max_len self.emb = nn.Embedding(vocab, d, padding_idx=PAD) self.pos_e, self.pos_d = nn.Embedding(max_len, d), nn.Embedding(max_len, d) self.stem = nn.ModuleList([ConvBlock(d, 5, drop), ConvBlock(d, 5, drop)]) self.enc = nn.ModuleList([EncLayer(d, heads, ff, drop) for _ in range(enc)]) self.dec = nn.ModuleList([DecLayer(d, heads, ff, drop) for _ in range(dec)]) self.ne, self.nd, self.drop = nn.LayerNorm(d), nn.LayerNorm(d), nn.Dropout(drop) self.out_bias = nn.Parameter(torch.zeros(vocab)) self.apply(self._init) def _init(self, m): if isinstance(m, (nn.Linear, nn.Conv1d)): nn.init.xavier_uniform_(m.weight) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, std=self.d ** -0.5) def encode(self, src): B, T = src.shape keep = src != PAD x = self.emb(src) * math.sqrt(self.d) + self.pos_e(torch.arange(T, device=src.device)) x = self.drop(x) k3 = keep.unsqueeze(-1).to(x.dtype) for conv in self.stem: x = conv(x, k3) mask = keep[:, None, None, :] for layer in self.enc: x = layer(x, mask) return self.ne(x), mask def logits(self, h): return F.linear(h, self.emb.weight, self.out_bias) # output layer shares the embedding matrix def forward(self, src, tgt_in): mem, mem_mask = self.encode(src) T = tgt_in.shape[1] x = self.emb(tgt_in) * math.sqrt(self.d) + self.pos_d(torch.arange(T, device=src.device)) x = self.drop(x) causal = torch.ones(T, T, dtype=torch.bool, device=src.device).tril() for layer in self.dec: x = layer(x, mem, causal, mem_mask) return self.logits(self.nd(x)) @torch.no_grad() def generate(self, src, max_new=None): """Greedy decoding with a key/value cache. Returns (token ids, mean log-prob, lowest token prob).""" self.eval() B = src.shape[0] dev = src.device mem, mem_mask = self.encode(src) if max_new is None: max_new = min(self.max_len - 2, int(src.shape[1] * 1.3) + 12) caches = [{} for _ in self.dec] tok = torch.full((B, 1), BOS, device=dev) outs = [] self.last_probs = [] # per-token confidence of the latest call, read by the safety gate in infer.py done = torch.zeros(B, dtype=torch.bool, device=dev) logp_sum = torch.zeros(B, device=dev) min_p = torch.ones(B, device=dev) n = torch.zeros(B, device=dev) for t in range(max_new): x = self.emb(tok) * math.sqrt(self.d) + self.pos_d(torch.tensor([t], device=dev)) for layer, c in zip(self.dec, caches): x = layer(x, mem, None, mem_mask, cache=c) lp = F.log_softmax(self.logits(self.nd(x[:, -1])).float(), -1) best, tok = lp.max(-1) live = ~done logp_sum += best * live n += live min_p = torch.where(live, torch.minimum(min_p, best.exp()), min_p) tok = torch.where(done, torch.full_like(tok, PAD), tok) outs.append(tok) self.last_probs.append(best.exp()) done |= tok == EOS if done.all(): break tok = tok.unsqueeze(1) return torch.stack(outs, 1), logp_sum / n.clamp(min=1), min_p def build(vocab, size="small", **kw): return WriteFix(vocab, **{**SIZES[size], **kw}) if __name__ == "__main__": for s in SIZES: m = build(16000, s) print(s, f"{sum(p.numel() for p in m.parameters()) / 1e6:.1f}M parameters")