Build a Tiny GPT From Scratch on a CPU: Attention, Transformer Blocks and Text Generation

"""The smallest real GPT: it reads the Sherlock Holmes stories one character at a time and learns to write more.
Every piece of GPT-2, Llama and Claude is here, just small enough to train on one CPU core in under a minute."""
import math

import torch
from torch import nn
from torch.nn import functional as F

torch.set_num_threads(1)   # the lab has one CPU; more threads only fight over it


# ---------- Step 1: characters to numbers ----------
class Chars:
    """The vocabulary: every distinct character in the text, sorted. encode() maps text to ids, decode() back."""

    def __init__(self, text):
        # TODO (Step 1): self.chars = sorted distinct characters; self.stoi = {char: index}.
        raise NotImplementedError("Step 1: write Chars.__init__()")

    def __len__(self):
        return len(self.chars)

    def encode(self, text):
        # TODO (Step 1): a list of ids, one per character.
        raise NotImplementedError("Step 1: write Chars.encode()")

    def decode(self, ids):
        # TODO (Step 1): join the characters back into a string.
        raise NotImplementedError("Step 1: write Chars.decode()")


def get_batch(data, batch_size, block_size, generator):
    """batch_size random windows of data: x = data[i:i+block_size], y = the same window one character later
    (data[i+1:i+block_size+1]). Starts i = torch.randint(len(data) - block_size, (batch_size,), generator=generator).
    Returns (x, y), both (batch_size, block_size) long tensors."""
    # TODO (Step 1): draw the starts, then stack the x windows and the y windows (one later).
    raise NotImplementedError("Step 1: write get_batch()")


# ---------- Step 2: a bigram model ----------
def lm_loss(logits, targets):
    """Cross-entropy between (B, T, V) scores and (B, T) next characters: flatten to (B*T, V) and (B*T,)."""
    raise NotImplementedError("lm_loss() arrives in Step 2")


class Bigram(nn.Module):
    """Scores for the next character from the current character alone: one row of a (V, V) table per character."""

    def __init__(self, vocab_size, block_size=64):
        super().__init__()
        self.block_size = block_size
        raise NotImplementedError("Bigram.__init__() arrives in Step 2")

    def forward(self, idx):
        raise NotImplementedError("Bigram.forward() arrives in Step 2")


@torch.no_grad()
def generate(model, idx, new_tokens, temperature=1.0, generator=None):
    """Extend idx (1, T) by new_tokens characters, one at a time: feed the last model.block_size ids, take the
    scores at the last position, divide by temperature, softmax, sample one id with torch.multinomial(probs, 1,
    generator=generator), append. Eval mode. Returns the (1, T + new_tokens) tensor."""
    raise NotImplementedError("generate() arrives in Step 2")


# ---------- Step 3: self-attention ----------
class Head(nn.Module):
    """One head of causal self-attention. Each position makes a query, a key and a value (Linear, no bias);
    weights = softmax(q @ k^T / sqrt(head_size)) with every future position masked to -inf; out = weights @ v.
    Keeps the last weights in self.weights, (B, T, T), so you can look at them."""

    def __init__(self, n_embd, head_size, block_size):
        super().__init__()
        raise NotImplementedError("Head.__init__() arrives in Step 3")
        self.weights = None

    def forward(self, x):
        raise NotImplementedError("Head.forward() arrives in Step 3")


# ---------- Step 4: the GPT ----------
class MultiHead(nn.Module):
    """n_head Heads of size n_embd // n_head side by side, outputs concatenated, then a Linear(n_embd, n_embd)."""

    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        raise NotImplementedError("MultiHead.__init__() arrives in Step 4")

    def forward(self, x):
        raise NotImplementedError("MultiHead.forward() arrives in Step 4")


class Block(nn.Module):
    """x = x + attention(LayerNorm(x)); x = x + mlp(LayerNorm(x)), the mlp being Linear(n_embd, 4 * n_embd), GELU,
    Linear(4 * n_embd, n_embd). Attention lets positions talk; the mlp thinks about what each one heard."""

    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        raise NotImplementedError("Block.__init__() arrives in Step 4")

    def forward(self, x):
        raise NotImplementedError("Block.forward() arrives in Step 4")


class GPT(nn.Module):
    """Token embedding + position embedding -> n_layer Blocks -> LayerNorm -> Linear to vocab_size scores."""

    def __init__(self, vocab_size, n_embd=64, n_head=4, n_layer=2, block_size=64):
        super().__init__()
        self.block_size = block_size
        raise NotImplementedError("GPT.__init__() arrives in Step 4")

    def forward(self, idx):
        raise NotImplementedError("GPT.forward() arrives in Step 4")


# ---------- Step 5: train it, and measure what context buys ----------
def train(model, data, steps, lr=3e-3, batch_size=32, seed=0):
    """AdamW(lr); `steps` batches from get_batch(data, batch_size, model.block_size, a generator seeded with seed).
    Returns the loss of every step."""
    raise NotImplementedError("train() arrives in Step 5")


@torch.no_grad()
def loss_by_position(model, data, batches=20, batch_size=32, seed=1):
    """The mean loss at each position 0 .. block_size - 1 of the window, over `batches` batches of data (eval mode,
    get_batch with a generator seeded with seed). Position 0 sees one character; the last sees block_size."""
    raise NotImplementedError("loss_by_position() arrives in Step 5")