Building a Brain

In App 1 we built the engine. Now we build the car.

We are going to build a GPT (Generative Pre-trained Transformer) from scratch in PyTorch. Not use one. Build one.

By the end of this chapter, you will have a model that generates text. It will not be ChatGPT. It will be more like a drunk Shakespeare. But it will be yours.

The Brain

The Plan

Load a text dataset (Shakespeare)
Turn characters into numbers (tokenizer)
Build a Transformer with self-attention
Train it to predict the next character
Generate text

Step 1: The Data

import torch
import torch.nn as nn
import torch.nn.functional as F
 
# Download Shakespeare
import urllib.request
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
text = urllib.request.urlopen(url).read().decode('utf-8')
 
print(f"Dataset: {len(text)} characters")
print(text[:200])

About 1 million characters of Shakespeare. Plenty for a tiny GPT.

Step 2: The Tokenizer

We use the simplest possible tokenizer — one character = one token:

chars = sorted(list(set(text)))
vocab_size = len(chars)
print(f"Vocabulary: {vocab_size} characters")
print(''.join(chars))
 
# Character to integer, and back
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])
 
# Encode the whole dataset
data = torch.tensor(encode(text), dtype=torch.long)

Step 3: Train/Val Split

n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]

Step 4: Data Loading

The model sees blocks of text at a time:

block_size = 64   # context length
batch_size = 32
 
def get_batch(split):
    d = train_data if split == 'train' else val_data
    ix = torch.randint(len(d) - block_size, (batch_size,))
    x = torch.stack([d[i:i+block_size] for i in ix])
    y = torch.stack([d[i+1:i+block_size+1] for i in ix])
    return x, y

For every input character, the target is the next character. That is the whole game.

Step 5: The Model

Here is the full GPT. Read it top to bottom — every piece builds on the last:

n_embd = 64
n_head = 4
n_layer = 4
dropout = 0.1
device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
class Head(nn.Module):
    """One head of self-attention."""
    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)
        q = self.query(x)
        # Attention scores
        wei = q @ k.transpose(-2, -1) * C**-0.5
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1)
        wei = self.dropout(wei)
        v = self.value(x)
        return wei @ v
 
class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        return self.dropout(self.proj(out))
 
class FeedForward(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )
 
    def forward(self, x):
        return self.net(x)
 
class Block(nn.Module):
    """Transformer block: attention + feedforward."""
    def __init__(self, n_embd, n_head):
        super().__init__()
        head_size = n_embd // n_head
        self.sa = MultiHeadAttention(n_head, head_size)
        self.ffwd = FeedForward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)
 
    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x
 
class GPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, n_embd)
        self.position_embedding = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)
 
    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T, device=device))
        x = tok_emb + pos_emb
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
 
        loss = None
        if targets is not None:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)
        return logits, loss
 
    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

That is a complete GPT. Token embeddings, position embeddings, multi-head self-attention, feedforward layers, layer norm, and text generation. About 100 lines.

Step 6: Train It

model = GPT().to(device)
print(f"Parameters: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M")
 
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
 
for step in range(5000):
    xb, yb = get_batch('train')
    xb, yb = xb.to(device), yb.to(device)
    logits, loss = model(xb, yb)
 
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
 
    if step % 500 == 0:
        print(f"Step {step}, Loss: {loss.item():.4f}")
 
print(f"Final loss: {loss.item():.4f}")

On a CPU this takes a few minutes. On a GPU, under a minute.

Step 7: Generate Text

context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(model.generate(context, max_new_tokens=500)[0].tolist()))

The output will look something like:

ROMEO:
And therefore hath she so much of my heart,
That I do know the state of my lord the king.

JULIET:
What say you to the prince?

It is not real Shakespeare. But it sounds like Shakespeare. Your GPT learned English grammar, character names, iambic-ish rhythm — all from predicting the next character.

What You Built

Component	What It Does
Tokenizer	Characters to numbers
Embedding	Numbers to vectors
Self-Attention	"Which other characters should I pay attention to?"
FeedForward	"Now that I know the context, what comes next?"
Block	Attention + FeedForward, stacked 4 times
Generate	Sample one token at a time, autoregressively

This is the same architecture behind GPT-2, GPT-3, GPT-4, and ChatGPT. They just have more layers, more heads, and more data. The idea is identical.

Inspired by Andrej Karpathy's nanoGPT and build-nanogpt.

Next up: we build something practical — an image classifier.