model.py

"""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.
@dataclass
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

    @staticmethod
    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)
# ---------------------------------------------------------------------------
@torch.no_grad()  # 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})")