| """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 |
|
|
| SIZES = { |
| "tiny": dict(d=256, heads=4, ff=1024, enc=5, dec=2), |
| "small": dict(d=384, heads=6, ff=1536, enc=6, dec=3), |
| "base": dict(d=512, heads=8, ff=2048, enc=8, dec=4), |
| "big": dict(d=640, heads=10, ff=2560, enc=10, dec=5), |
| } |
|
|
|
|
| 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): |
| 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: |
| 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: |
| 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) |
|
|
| 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 = [] |
| 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") |
|
|
