Index Ai Academy
Project tutorials · 6 steps

Build a GPT from scratch in Python

Write a tiny GPT in PyTorch, train it on a product catalog, and watch it write new descriptions. The whole program is about 70 lines.

What you'll build

A GPT is smaller than it sounds. This tutorial writes the whole thing, attention included, and trains it on 300 short product descriptions. It's a real transformer with about 38,000 parameters, where the big models have hundreds of billions. The design is the same.

What you'll learn

  • Prepare a corpus for a language model: tokenize, split, and batch with shifted targets
  • Train a small GPT from scratch and watch its loss fall
  • Generate text from it and judge what it has and hasn't learned
  • Connect what you see to how large language models behave

What you need

Python and basic PyTorch, and it helps to know what attention is. In the lesson it runs in your browser, on an in-browser subset of PyTorch, and takes a while. On your own machine, install PyTorch and change the corpus address to https://www.indexaiacademy.com/mock/v1/corpus. Real PyTorch gives somewhat different numbers.

Step 1: Tokenize the text

Cut the text into tokens and give each one an ID. Here a token is a word, a number, a punctuation mark or a line break, which gives a vocabulary small enough to train quickly. Real models use BPE instead, which splits rare words into pieces. The last 10% is held back to show how the model does on text it hasn't seen.

import math
import re
import requests
import torch
import torch.nn as nn

# 1. tokenize: words, numbers, punctuation and line breaks
text = requests.get("https://api.gradient.test/v1/corpus").json()["text"]
tokens = re.findall(r"[A-Za-z]+|[0-9]+|[^\sA-Za-z0-9]|\n", text)
vocab = sorted(set(tokens))
stoi = {t: i for i, t in enumerate(vocab)}
data = torch.tensor([stoi[t] for t in tokens])
split = int(0.9 * len(data))
train_data, val_data = data[:split], data[split:]
print(f"{len(tokens):,} tokens, vocabulary of {len(vocab)}")
What it prints
8,758 tokens, vocabulary of 183

Step 2: Cut it into training windows

A training example is a window of text. The target y is the same window moved along by one token, so at every position the model learns to predict what comes next. One window of 24 tokens is really 24 small lessons.

# 2. batches of windows; y is x shifted by one token
def get_batch(d, block=24, batch=16):
    starts = torch.randint(0, len(d) - block - 1, (batch,)).tolist()
    x = torch.stack([d[i:i + block] for i in starts])
    y = torch.stack([d[i + 1:i + block + 1] for i in starts])
    return x, y

Rather write it yourself?

In the lesson you write this from a starter, in your browser with nothing to install, and checks tell you what's wrong as you go. A free account opens it and all 244 lessons.

Create a free account

Step 3: Write a transformer block

This is the heart of it. Each token produces a query, a key and a value. A query compared with every key says how much attention to pay to each token. split and view divide the work among the heads, so each can look for something different. The tril mask hides the future, because a model that could see the next token would learn nothing. The result goes through a small network, with a shortcut around each stage that keeps training stable.

# 3. the model from the last lesson
class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.h = h
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv, self.proj = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        B, T, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).split(d, dim=-1)
        q, k, v = (t.view(B, T, self.h, d // self.h).transpose(1, 2) for t in (q, k, v))
        s = (q @ k.transpose(-2, -1) / math.sqrt(d // self.h)).masked_fill(~torch.tril(torch.ones(T, T)).bool(), float("-inf"))
        x = x + self.proj((torch.softmax(s, dim=-1) @ v).transpose(1, 2).reshape(B, T, d))
        return x + self.mlp(self.ln2(x))

Step 4: Assemble the GPT

A GPT is two embeddings, a stack of blocks and a final layer. One embedding says which token this is, and the other says where it sits, because attention alone has no sense of order. The final layer turns each position's numbers into a score for every token in the vocabulary.

class GPT(nn.Module):
    def __init__(self, vocab, context, d, n_layers, n_heads):
        super().__init__()
        self.context = context
        self.tok, self.pos = nn.Embedding(vocab, d), nn.Embedding(context, d)
        self.blocks = nn.Sequential(*[Block(d, n_heads) for _ in range(n_layers)])
        self.ln_f, self.head = nn.LayerNorm(d), nn.Linear(d, vocab)

    def forward(self, idx):
        x = self.tok(idx) + self.pos(torch.arange(idx.shape[1]))
        return self.head(self.ln_f(self.blocks(x)))

Step 5: Train it

Before training, the model guesses evenly among 183 tokens, and the natural log of 183 is 5.209, so a starting loss near that says the setup is right. Then the loop is the usual one: run the model, measure the loss, backpropagate, step. Validation loss staying close to training loss means it's learning patterns, and not yet memorizing the text.

torch.manual_seed(0)
model = GPT(len(vocab), context=24, d=32, n_layers=2, n_heads=4)
print(f"{sum(p.numel() for p in model.parameters()):,} parameters")
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)

# 4. train
for step in range(301):
    x, y = get_batch(train_data)
    loss = nn.functional.cross_entropy(model(x).view(-1, len(vocab)), y.view(-1))
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 50 == 0:
        with torch.no_grad():
            xv, yv = get_batch(val_data, batch=64)
            val = nn.functional.cross_entropy(model(xv).view(-1, len(vocab)), yv.view(-1)).item()
        print(f"step {step:>3}  training loss {loss.item():.3f}  validation loss {val:.3f}")
What it prints
38,135 parameters
step   0  training loss 5.397  validation loss 5.326
step  50  training loss 2.636  validation loss 2.646
step 100  training loss 1.379  validation loss 1.561
step 150  training loss 1.173  validation loss 1.221
step 200  training loss 0.970  validation loss 1.067
step 250  training loss 0.908  validation loss 0.989
step 300  training loss 0.861  validation loss 0.908

Step 6: Generate text

To generate, run the model, take the scores for the last position, turn them into probabilities and sample one token. Add it to the sequence and repeat. Temperature below 1 makes the likely choices more likely.

# 5. generate
@torch.no_grad()
def generate(idx, n, temperature=0.8):
    for _ in range(n):
        logits = model(idx[:, -model.context:])[:, -1, :] / temperature
        idx = torch.cat([idx, torch.multinomial(torch.softmax(logits, dim=-1), 1)], dim=1)
    return idx

torch.manual_seed(1)
out = generate(torch.tensor([[stoi["\n"]]]), 70)[0].tolist()
words = " ".join(vocab[i] for i in out)
print("\n" + re.sub(r" ?\n ?", "\n", re.sub(r" ([.,])", r"\1", words)).strip())
What it prints
Our Vista camp mug is made from recycled ripstop nylon. It is packable, light and hikes. Made for day hikes. Packs into its own pocket.
Meet the River sleeping bag, made from a recycled polyester fleece. It is light, easy to layer and light. Made for travel. Packs into its own pocket.
Meet the Nomad rain jacket

This is only 300 training steps on a tiny dataset, and it has learned the catalog's sentence patterns and which words go where. It also has no idea what's true, which is why "light and hikes" comes out. A large language model is this design with every number turned up.

Where to take it next

  • Raise the width, the layers or the context, and see how the loss and the time respond.
  • Swap in your own text. The tokenizer will need to handle what's in it.
  • Replace the word tokenizer with BPE, which is what real models use.

Mistakes to watch for

  • Thinking a model that writes fluently must be accurate.
  • Assuming bigger models work differently rather than at larger scale.
  • Evaluating a language model only on its training text.
  • Believing the model stores sentences and retrieves them.
The complete program
import math
import re
import requests
import torch
import torch.nn as nn

# 1. tokenize: words, numbers, punctuation and line breaks
text = requests.get("https://api.gradient.test/v1/corpus").json()["text"]
tokens = re.findall(r"[A-Za-z]+|[0-9]+|[^\sA-Za-z0-9]|\n", text)
vocab = sorted(set(tokens))
stoi = {t: i for i, t in enumerate(vocab)}
data = torch.tensor([stoi[t] for t in tokens])
split = int(0.9 * len(data))
train_data, val_data = data[:split], data[split:]
print(f"{len(tokens):,} tokens, vocabulary of {len(vocab)}")

# 2. batches of windows; y is x shifted by one token
def get_batch(d, block=24, batch=16):
    starts = torch.randint(0, len(d) - block - 1, (batch,)).tolist()
    x = torch.stack([d[i:i + block] for i in starts])
    y = torch.stack([d[i + 1:i + block + 1] for i in starts])
    return x, y

# 3. the model from the last lesson
class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.h = h
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv, self.proj = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        B, T, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).split(d, dim=-1)
        q, k, v = (t.view(B, T, self.h, d // self.h).transpose(1, 2) for t in (q, k, v))
        s = (q @ k.transpose(-2, -1) / math.sqrt(d // self.h)).masked_fill(~torch.tril(torch.ones(T, T)).bool(), float("-inf"))
        x = x + self.proj((torch.softmax(s, dim=-1) @ v).transpose(1, 2).reshape(B, T, d))
        return x + self.mlp(self.ln2(x))

class GPT(nn.Module):
    def __init__(self, vocab, context, d, n_layers, n_heads):
        super().__init__()
        self.context = context
        self.tok, self.pos = nn.Embedding(vocab, d), nn.Embedding(context, d)
        self.blocks = nn.Sequential(*[Block(d, n_heads) for _ in range(n_layers)])
        self.ln_f, self.head = nn.LayerNorm(d), nn.Linear(d, vocab)

    def forward(self, idx):
        x = self.tok(idx) + self.pos(torch.arange(idx.shape[1]))
        return self.head(self.ln_f(self.blocks(x)))

torch.manual_seed(0)
model = GPT(len(vocab), context=24, d=32, n_layers=2, n_heads=4)
print(f"{sum(p.numel() for p in model.parameters()):,} parameters")
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)

# 4. train
for step in range(301):
    x, y = get_batch(train_data)
    loss = nn.functional.cross_entropy(model(x).view(-1, len(vocab)), y.view(-1))
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 50 == 0:
        with torch.no_grad():
            xv, yv = get_batch(val_data, batch=64)
            val = nn.functional.cross_entropy(model(xv).view(-1, len(vocab)), yv.view(-1)).item()
        print(f"step {step:>3}  training loss {loss.item():.3f}  validation loss {val:.3f}")

# 5. generate
@torch.no_grad()
def generate(idx, n, temperature=0.8):
    for _ in range(n):
        logits = model(idx[:, -model.context:])[:, -1, :] / temperature
        idx = torch.cat([idx, torch.multinomial(torch.softmax(logits, dim=-1), 1)], dim=1)
    return idx

torch.manual_seed(1)
out = generate(torch.tensor([[stoi["\n"]]]), 70)[0].tolist()
words = " ".join(vocab[i] for i in out)
print("\n" + re.sub(r" ?\n ?", "\n", re.sub(r" ([.,])", r"\1", words)).strip())

Write it yourself

In the lesson you write this from a starter, in your browser with nothing to install, and checks tell you what's wrong as you go. A free account opens it and all 244 lessons.

Create a free account