File size: 8,295 Bytes
d7d850a | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 | """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 <fix> 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")
|