| """Step 2a: the transformer, written by hand. Symbols in, next-symbol probabilities out. | |
| Run on its own (`python model.py`) to print every weight and check that an untrained | |
| model's loss is near ln 28 = 3.33: with random weights it knows nothing yet, so it | |
| should guess every symbol about evenly. | |
| How to read this file: start at `TinyTowns.forward` near the bottom. It is the whole | |
| model in ~10 lines. Then go up into `Layer`, `Attention` and `Knowledge` for the parts. | |
| A note on shapes. Comments like [B, T, C] give the size of a grid of numbers: | |
| B = batch how many names we process at once (e.g. 64) | |
| T = positions how many symbols per name (at most 21) | |
| C = note size how many numbers in each note (16 nano, 64 small, 128 tinytowns) | |
| So [B, T, C] means: for every name, for every position, a note of C numbers. | |
| """ | |
| import json | |
| import math | |
| import random | |
| from dataclasses import dataclass | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| # --------------------------------------------------------------------------- | |
| # Letters <-> numbers (same as step 1) | |
| # --------------------------------------------------------------------------- | |
| # The model only understands numbers, so every symbol gets an index 0..27. | |
| # Index 0 '.' means "start of name" when it goes in, and "end of name" when it comes out. | |
| ALPHABET = ". abcdefghijklmnopqrstuvwxyz" | |
| stoi = {ch: i for i, ch in enumerate(ALPHABET)} # "string to int": 'a' -> 2 | |
| itos = {i: ch for i, ch in enumerate(ALPHABET)} # "int to string": 2 -> 'a' | |
| CONTEXT = 21 # longest input the model sees: one '.' plus at most 20 letters | |
| # --------------------------------------------------------------------------- | |
| # Config: the only thing that differs between the two sizes | |
| # --------------------------------------------------------------------------- | |
| # The code below is identical for both models; only these numbers change. | |
| class Config: | |
| note: int # numbers per note (C). Bigger = each symbol can carry more information. | |
| layers: int # how many (attention + knowledge step) layers are stacked | |
| heads: int # attention heads per layer; each head looks for its own pattern | |
| mlp: int # width of the knowledge step's middle (16 -> 64 -> 16 for nano) | |
| vocab: int = len(ALPHABET) # 28 symbols | |
| context: int = CONTEXT # 21 positions | |
| CONFIGS = { | |
| "tinytowns-nano": Config(note=16, layers=1, heads=1, mlp=64), # 4,572 weights | |
| "tinytowns-small": Config(note=64, layers=2, heads=4, mlp=256), # 105,052 weights | |
| "tinytowns": Config(note=128, layers=5, heads=4, mlp=512), # 1,001,500 weights | |
| } | |
| # --------------------------------------------------------------------------- | |
| # Turning names into model input | |
| # --------------------------------------------------------------------------- | |
| def encode(names): | |
| """Names -> (inputs x, targets y), each a [batch, 21] grid of numbers. | |
| Example, 'paris': | |
| position: 0 1 2 3 4 5 6 ... | |
| x (input): . p a r i s . . . <- '.' + name, padded with '.' | |
| y (answer):p a r i s . -1 -1 -1 <- name + '.', padded with -1 | |
| Read it column by column: at position t, the model sees x[0..t] and must predict y[t]. | |
| So one name gives one prediction per letter, plus one for "end of name". | |
| Why padding: a batch must be a neat rectangle, but names have different lengths. | |
| The padded inputs are harmless: the mask (in Attention) stops real positions from | |
| looking ahead at them. The padded targets are -1, which the loss is told to skip. | |
| """ | |
| x = torch.zeros((len(names), CONTEXT), dtype=torch.long) # all '.' to start | |
| y = torch.full((len(names), CONTEXT), -1, dtype=torch.long) # all "skip" to start | |
| for row, name in enumerate(names): | |
| ids = [stoi[ch] for ch in name] | |
| x[row, 1 : len(ids) + 1] = torch.tensor(ids, dtype=torch.long) # after the first '.' | |
| y[row, : len(ids) + 1] = torch.tensor(ids + [0], dtype=torch.long) # name, then '.' | |
| return x, y | |
| # --------------------------------------------------------------------------- | |
| # Attention: each note gathers information from the notes before it | |
| # --------------------------------------------------------------------------- | |
| class Attention(nn.Module): | |
| def __init__(self, cfg): | |
| super().__init__() | |
| self.heads = cfg.heads | |
| # The Q, K, V grids. nn.Linear(16, 16) is a 16x16 grid of weights plus 16 biases | |
| # (one number added to each output). note x grid + bias = q (or k, or v). | |
| # With several heads, all heads' grids sit side by side in one big grid; | |
| # we split the result into per-head slices in forward(). | |
| self.q = nn.Linear(cfg.note, cfg.note) # Q grid: "what am I looking for?" | |
| self.k = nn.Linear(cfg.note, cfg.note) # K grid: "what do I contain?" | |
| self.v = nn.Linear(cfg.note, cfg.note) # V grid: "what do I hand over if chosen?" | |
| # After mixing, one more grid turns the heads' results back into one note-shaped | |
| # update. With one head it is still useful: it decides how the mix is added in. | |
| self.out = nn.Linear(cfg.note, cfg.note) | |
| # The mask: future[i, j] is True when position j comes after position i. | |
| # torch.triu(..., diagonal=1) keeps only the part above the diagonal, e.g. for 4: | |
| # . T T T row = "me", column = "the one I might look at" | |
| # . . T T True = in my future = forbidden | |
| # . . . T | |
| # . . . . | |
| future = torch.triu(torch.ones(cfg.context, cfg.context, dtype=torch.bool), diagonal=1) | |
| # register_buffer: keep it with the model (moves to GPU with it) but it is a fixed | |
| # table, not a weight, so training never changes it. persistent=False: don't write it | |
| # into the weights file; it is rebuilt from the config, so the file holds only weights. | |
| self.register_buffer("future", future, persistent=False) | |
| def forward(self, notes): | |
| B, T, C = notes.shape # batch, positions, numbers per note | |
| hs = C // self.heads # numbers per head ("head size"): 16 nano, 64/4 = 16 small, 128/4 = 32 tinytowns | |
| # 1) q, k, v for every position at once: [B, T, C] x grid -> [B, T, C]. | |
| # .view splits the last C numbers into (heads, hs); .transpose moves "heads" forward | |
| # so each head becomes its own little [T, hs] problem: shape [B, heads, T, hs]. | |
| # With 1 head these two calls change nothing but the shape. | |
| q = self.q(notes).view(B, T, self.heads, hs).transpose(1, 2) | |
| k = self.k(notes).view(B, T, self.heads, hs).transpose(1, 2) | |
| v = self.v(notes).view(B, T, self.heads, hs).transpose(1, 2) | |
| # 2) Score: every q against every k, all in one matrix multiply. | |
| # score[i, j] = q_i . k_j = "how relevant is position j to position i?" -> [B, heads, T, T] | |
| # Divide by sqrt(hs): a dot product of hs numbers tends to grow with hs, and big | |
| # scores make softmax put ~100% on one position, which makes learning slow. | |
| score = q @ k.transpose(-2, -1) / math.sqrt(hs) | |
| # 3) Mask: set every "future" score to minus infinity. After softmax, e^-inf = 0, | |
| # so no position can take anything from a later one. This is what lets the model | |
| # train on all positions at once without cheating by looking at the answer. | |
| score = score.masked_fill(self.future[:T, :T], float("-inf")) | |
| # 4) Shares: softmax along each row turns scores into percentages that sum to 1. | |
| shares = F.softmax(score, dim=-1) | |
| # 5) Mix: each position's new information = its shares x everyone's v. | |
| # [B, heads, T, T] @ [B, heads, T, hs] -> [B, heads, T, hs] | |
| mix = shares @ v | |
| # 6) Put the heads back side by side ([B, T, C]) and pass through the out grid. | |
| mix = mix.transpose(1, 2).reshape(B, T, C) | |
| return self.out(mix) | |
| # --------------------------------------------------------------------------- | |
| # Knowledge step: each note is processed on its own, no looking at neighbours | |
| # --------------------------------------------------------------------------- | |
| class Knowledge(nn.Module): | |
| def __init__(self, cfg): | |
| super().__init__() | |
| self.up = nn.Linear(cfg.note, cfg.mlp) # 16 -> 64: spread out into more room | |
| self.down = nn.Linear(cfg.mlp, cfg.note) # 64 -> 16: squeeze back to note size | |
| def forward(self, notes): | |
| # GELU is a bend: roughly keeps positive numbers and squashes negative ones to ~0. | |
| # Without it, "x grid, then x grid" is the same as one combined grid, so the extra | |
| # 64-wide middle would add no new ability. The bend is what lets it learn rules like | |
| # "IF the note says X AND Y, THEN push towards letter Z". | |
| return self.down(F.gelu(self.up(notes))) | |
| # --------------------------------------------------------------------------- | |
| # One layer = attention, then knowledge step, each added onto the note | |
| # --------------------------------------------------------------------------- | |
| class Layer(nn.Module): | |
| def __init__(self, cfg): | |
| super().__init__() | |
| # LayerNorm rescales each note to average 0 and spread 1, then applies a learned | |
| # scale and shift (16 + 16 weights). It keeps numbers in a steady range so that | |
| # training doesn't blow up or fade out as notes pass through the layers. | |
| self.norm1 = nn.LayerNorm(cfg.note) | |
| self.attention = Attention(cfg) | |
| self.norm2 = nn.LayerNorm(cfg.note) | |
| self.knowledge = Knowledge(cfg) | |
| def forward(self, notes): | |
| # "notes + ..." not "notes = ...": each part ADDS an update to the note instead of | |
| # replacing it. The original information always flows through untouched, and each | |
| # part only has to learn a correction. (Standard name: residual connection.) | |
| notes = notes + self.attention(self.norm1(notes)) | |
| notes = notes + self.knowledge(self.norm2(notes)) | |
| return notes | |
| # --------------------------------------------------------------------------- | |
| # The whole model | |
| # --------------------------------------------------------------------------- | |
| class TinyTowns(nn.Module): | |
| def __init__(self, cfg): | |
| super().__init__() | |
| self.cfg = cfg | |
| # nn.Embedding(28, 16) is a 28-row table: give it symbol 4, get row 4 (16 numbers). | |
| self.lookup = nn.Embedding(cfg.vocab, cfg.note) # lookup table: symbol -> note | |
| self.position = nn.Embedding(cfg.context, cfg.note) # position table: position -> note | |
| self.layers = nn.ModuleList(Layer(cfg) for _ in range(cfg.layers)) | |
| self.norm = nn.LayerNorm(cfg.note) # one last tidy-up before output | |
| self.output = nn.Linear(cfg.note, cfg.vocab) # output table: note -> 28 scores | |
| self.apply(self._small_start) # run _small_start on every part above | |
| def _small_start(m): | |
| # Every weight starts as a small random number (around +-0.02), every bias at 0. | |
| # Small => all 28 output scores start near 0 => softmax gives ~1/28 each => | |
| # the untrained loss starts near ln 28 = 3.33. Random (not all equal) => each weight | |
| # starts slightly different, so each one gets a different nudge and they can learn | |
| # different jobs. If they all started equal they would all stay equal forever. | |
| if isinstance(m, (nn.Linear, nn.Embedding)): | |
| nn.init.normal_(m.weight, mean=0.0, std=0.02) | |
| if isinstance(m, nn.Linear): | |
| nn.init.zeros_(m.bias) | |
| def forward(self, x, y=None): | |
| """x: [B, T] symbol numbers. y: [B, T] answers, or None when we only want predictions. | |
| Returns (scores, loss). scores is [B, T, 28]: for every name and every position, | |
| 28 raw scores for the next symbol (softmax of them = the probabilities). | |
| """ | |
| B, T = x.shape | |
| # 1) Lookup + position: each symbol becomes a note, plus "where am I" information. | |
| # torch.arange(T) = [0, 1, ..., T-1], the position numbers. -> [B, T, C] | |
| notes = self.lookup(x) + self.position(torch.arange(T, device=x.device)) | |
| # 2) The layers: attention (look back) and knowledge step (think), repeated. | |
| for layer in self.layers: | |
| notes = layer(notes) | |
| # 3) Output table: note -> 28 scores, one per possible next symbol. -> [B, T, 28] | |
| # These raw scores are called "logits"; softmax turns them into probabilities. | |
| scores = self.output(self.norm(notes)) | |
| if y is None: | |
| return scores, None | |
| # 4) Loss, exactly as in step 1: average of -ln(probability of the real next symbol). | |
| # F.cross_entropy does softmax + -ln p + average in one go. .view(-1, 28) flattens | |
| # [B, T, 28] into a plain list of predictions [B*T, 28], and y into [B*T] answers. | |
| # ignore_index=-1 skips the padding positions. | |
| loss = F.cross_entropy(scores.view(-1, self.cfg.vocab), y.view(-1), ignore_index=-1) | |
| return scores, loss | |
| # --------------------------------------------------------------------------- | |
| # Weights file: safetensors (step 6) | |
| # --------------------------------------------------------------------------- | |
| # A .safetensors file is: 8 bytes (header length) + a JSON header (each grid's name, number type, | |
| # shape and byte position) + all the numbers back to back. Nothing in it can run as code. | |
| # The header may also carry "__metadata__": text only, so we store the config as JSON text. | |
| # read_weights.py opens this format by hand, without the library. | |
| def save_weights(model, path, **notes): | |
| from safetensors.torch import save_file | |
| meta = {"config": json.dumps(model.cfg.__dict__)} | |
| meta.update({k: json.dumps(v) for k, v in notes.items()}) # e.g. best={"step": 5500, ...} | |
| save_file(model.state_dict(), path, metadata=meta) | |
| def load_weights(path, device="cpu"): | |
| """Rebuild a model from a .safetensors file. Returns (model, notes).""" | |
| from safetensors import safe_open | |
| from safetensors.torch import load_file | |
| with safe_open(path, framework="pt") as f: | |
| meta = {k: json.loads(v) for k, v in f.metadata().items()} | |
| model = TinyTowns(Config(**meta.pop("config"))) | |
| model.load_state_dict(load_file(path)) # every grid, matched by name; errors if any is missing | |
| return model.to(device).eval(), meta | |
| # --------------------------------------------------------------------------- | |
| # Inventing names: predict, pick, feed back in, repeat (as in step 1) | |
| # --------------------------------------------------------------------------- | |
| # we're only predicting, not learning, so skip the gradient bookkeeping | |
| def sample(model, n, temperature=1.0, generator=None): | |
| """Invent n names. temperature < 1 = safer, more typical; > 1 = wilder.""" | |
| device = next(model.parameters()).device | |
| names = [] | |
| for _ in range(n): | |
| ids = [0] # start with '.' | |
| while len(ids) < CONTEXT: | |
| scores, _ = model(torch.tensor([ids], device=device)) | |
| # Only the last position matters: it's the prediction for the next symbol. | |
| probs = F.softmax(scores[0, -1] / temperature, dim=-1) | |
| nxt = torch.multinomial(probs.cpu(), 1, generator=generator).item() | |
| if nxt == 0: # '.' = end of name | |
| break | |
| ids.append(nxt) | |
| names.append("".join(itos[i] for i in ids[1:])) | |
| return names | |
| # --------------------------------------------------------------------------- | |
| # Self-check: `python model.py` | |
| # --------------------------------------------------------------------------- | |
| if __name__ == "__main__": | |
| torch.manual_seed(42) # same random starting weights every run | |
| with open("data/cities.txt", encoding="utf-8") as f: | |
| names = [line.strip() for line in f if line.strip()] | |
| batch = random.Random(42).sample(names, 256) | |
| x, y = encode(batch) | |
| print(f"example: '{batch[0]}'\n input {x[0].tolist()}\n target {y[0].tolist()}") | |
| for size, cfg in CONFIGS.items(): | |
| model = TinyTowns(cfg) | |
| print(f"\n{size}: {cfg}") | |
| if size == "tinytowns-nano": | |
| # named_parameters() lists every weight grid the model owns, by name. | |
| for name, w in model.named_parameters(): | |
| print(f" {name:<32} {str(list(w.shape)):<10} {w.numel():>6,}") | |
| total = sum(w.numel() for w in model.parameters()) | |
| _, loss = model(x, y) | |
| print(f" total weights: {total:,}") | |
| print(f" untrained loss on 256 names: {loss.item():.4f} (even guess: {math.log(28):.4f})") | |