Train a GPT From Scratch on Your Mac

Build and train a 10M parameter GPT model on Shakespeare text from scratch.

Introduction

30 Second Summary

Every time you use ChatGPT or Claude, a transformer model is generating text one token at a time. But what's actually happening inside that model remains a black box to most people who use it.

In this project, you will build a GPT language model entirely from scratch in PyTorch, train it on Shakespeare text using your Mac's GPU, and generate new text from the trained model.

What You'll Build

A complete GPT training pipeline consisting of three Python files: a character-level tokenizer, a full transformer model with multi-head self-attention, and a training loop with text generation.

By the end of this project, you'll have:

  • A working ~10 million parameter GPT model you wrote entirely from scratch in PyTorch.
  • Trained it on Shakespeare text using your Mac's MPS GPU acceleration and observed the loss curve converge.
  • Generated coherent Shakespeare-like text from your trained model using temperature and top-k sampling.
  • Secret Mission: Implement nucleus (top-p) sampling and compare its output quality against your top-k implementation.

Are there any prerequisites?

You should be comfortable reading Python code, but no machine learning experience is required. A strong programming background in any language is enough to follow along.

You'll need a MacBook with Apple Silicon (M1 or later). This project uses the MPS backend for GPU-accelerated training.

Set Up Python and PyTorch

Before you can build a GPT model, you need the right tools on your machine. PyTorch is the deep learning framework you will use to define and train your model, and it needs Python 3.12 to run.

In this step, you will install Python via Homebrew, set up an isolated project environment, install PyTorch with GPU acceleration for your Mac's MPS (Metal Performance Shaders) backend, and download the Shakespeare text your model will learn from.

In this step, get ready to:
  • Install Python 3.12 and create a virtual environment.
  • Install PyTorch and confirm MPS GPU acceleration is available.
  • Create the project structure and download the Shakespeare training data.
Install Python 3.12 and create a virtual environment

Your Mac ships with an older system Python that is not suitable for this project. You need Python 3.12 installed through Homebrew so you have a modern version with full library support.

  • Check if Python 3.12 is already installed by running this command in your Terminal:
python3.12 --version

✔️ I see version 3.12 or higher

Python 3.12 is ready. Skip ahead to creating your project directory below.

ⓧ I see an older version

Your system is resolving a different Python version. Install Python 3.12 through Homebrew by running:

brew install python@3.12

After the install completes, confirm it worked by running:

python3.12 --version

You should now see Python 3.12.13 or similar. If your terminal still resolves an older version, use the full path /opt/homebrew/bin/python3.12 instead.

ⓧ Command not found

Python 3.12 is not installed yet. Install it through Homebrew by running:

brew install python@3.12

This downloads and installs the Python 3.12 formula. After it finishes, confirm the installation by running:

python3.12 --version

You should see Python 3.12.13 or similar printed in your terminal.

Still seeing command not found?

Homebrew installs Python at /opt/homebrew/bin/python3.12 on Apple Silicon Macs. Try using the full path directly. If Homebrew itself is not working, confirm it is installed by running brew --version.

help me troubleshoot my Python 3.12 installation on macOS with Homebrew

Now create your project directory and move into it.

  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir gpt-from-scratch
cd gpt-from-scratch
  • Create a virtual environment by running:
python3.12 -m venv venv

What is a virtual environment?

A virtual environment is an isolated Python installation that keeps your project's packages separate from the rest of your system. This prevents version conflicts between different projects.

  • Activate the virtual environment by running:
source venv/bin/activate

You should see (venv) appear at the start of your terminal prompt. This confirms you are now working inside the isolated environment.

Don't see (venv) in your prompt?

Make sure you ran the source command (not just the path). Also confirm you are inside the gpt-from-scratch folder by running pwd. You should see /Users/<your-username>/Desktop/gpt-from-scratch.

help me activate my Python virtual environment on macOS

Install PyTorch and verify MPS is available

PyTorch is the library that provides tensors, automatic differentiation, and neural network building blocks. On Apple Silicon Macs, PyTorch can use the MPS backend to run computations on your GPU instead of your CPU.

  • Install PyTorch by running:
pip3 install torch torchvision

Why not specify a GPU version?

On Linux, you need to pick between CPU and CUDA versions of PyTorch. On macOS, the standard pip3 install torch package already includes MPS support. No special index URL is needed.

The install may take a minute or two. Once it finishes, verify that PyTorch can see your Mac's GPU.

  • Run a quick MPS test by entering this command:
python3 -c "import torch; print('PyTorch version:', torch.__version__); print('MPS available:', torch.backends.mps.is_available())"

✔️ I see MPS available: True

Your Mac's GPU is ready for training. PyTorch will use the Metal Performance Shaders backend to accelerate tensor operations.

ⓧ I see MPS available: False

This can happen if your macOS version is too old. MPS requires macOS 14.0 or later. Check your macOS version by clicking the Apple menu and selecting About This Mac.

If you are on macOS 14.0+ and still see False, try restarting your terminal and re-activating the virtual environment. The project will still work on CPU if MPS is unavailable, but training will be slower.

help me troubleshoot MPS availability in PyTorch on my M4 Mac

ⓧ I see an import error

If you see ModuleNotFoundError: No module named 'torch', PyTorch did not install correctly. Make sure your virtual environment is active (you see (venv) in your prompt) and re-run:

pip3 install torch torchvision

help me fix a PyTorch installation error in my virtual environment

Create the project structure and download training data

Your model needs text to learn from. You will use the Tiny Shakespeare dataset, a ~1MB file containing works of Shakespeare that is the classic training corpus for character-level language models.

  • Create the data directory by running:
mkdir data
  • Download the Shakespeare text file by running:
curl -o data/shakespeare.txt https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt

What is this dataset?

This is a ~1.1 million character file containing a subset of Shakespeare's plays. It has 65 unique characters (letters, punctuation, spaces) which becomes your model's vocabulary. The small vocabulary is what makes character-level tokenization work well at this scale.

  • Confirm the download succeeded by running:
wc -c data/shakespeare.txt

You should see approximately 1115394 bytes (about 1.1MB).

  • Preview the first few lines to confirm the text looks right:
head -5 data/shakespeare.txt

You should see the opening lines of a Shakespeare play starting with First Citizen:.

Seeing a 404 or empty file?

Make sure you are connected to the internet and that you typed the URL exactly as shown. If curl is not available, you can also try wget instead. Run wc -c data/shakespeare.txt to check the file is not empty (0 bytes means the download failed).

help me download the Shakespeare training data for my GPT project

Finally, open the project folder in your editor so you are ready to write code in the next step.

  • Open the gpt-from-scratch folder in VS Code or Cursor by running:
code .

Using Cursor instead?

If you use Cursor as your editor, run cursor . instead. If the code command is not recognized, open VS Code manually, click File then Open Folder, and select the gpt-from-scratch folder on your Desktop.

  • Confirm your project structure looks correct by running:
ls

You should see data and venv listed in the output.

Your development environment is fully set up. Next up, you will build the tokenizer that converts Shakespeare text into numbers your model can learn from.

Build the Tokenizer and Data Pipeline

Your environment is ready and your Shakespeare text is sitting in the data/ folder. Now comes the first real question of any language model project: how do you turn raw text into numbers a neural network can process?

In this step, you will build a character-level tokenizer that converts text into sequences of integers, then create a data pipeline that serves random training batches to your model on the MPS GPU.

In this step, get ready to:
  • Build a character-level tokenizer with encode and decode functions.
  • Create a train/val data split and a batch generator that loads tensors onto MPS.
  • Verify the pipeline by printing batch shapes and decoding a sample back to readable text.
Build a character-level tokenizer

Neural networks operate on numbers, not letters. A tokenizer is the bridge between human-readable text and the integer sequences your model will learn from.

Why character-level instead of BPE?

Modern systems like GPT-4 and Claude use Byte Pair Encoding (BPE) tokenizers with vocabularies of 50,000+ tokens. BPE needs massive datasets (100MB+) to learn useful token patterns.

Shakespeare is only ~1MB. At that scale, a character-level tokenizer with just 65 unique characters works perfectly. Every character in the text gets its own integer ID, and the model learns letter-by-letter patterns from there.

  • Create a new file called train.py in your gpt-from-scratch project folder.
  • Add the imports and training hyperparameters by pasting this code:
import torch
import math
from model import GPT, block_size, n_embd, vocab_size

# Training hyperparameters
batch_size = 64
max_iters = 3000
eval_interval = 200
eval_iters = 100
learning_rate = 3e-4
min_lr = 3e-5
warmup_iters = 100
device = 'mps' if torch.backends.mps.is_available() else 'cpu'

print(f"Using device: {device}")

What does this code do?

The import from model references a file you will create in the next step. It will not run yet, but it sets up the dependency structure correctly.

  • batch_size = 64 means the model will process 64 text chunks in parallel during training.
  • device selects your M4 GPU (MPS) for hardware acceleration, falling back to CPU if needed.
  • The remaining hyperparameters (max_iters, learning_rate, etc.) control the training loop you will build in Step 4.
  • Below the hyperparameters, add the code that loads the Shakespeare text and builds the tokenizer:
# Load data
with open('data/shakespeare.txt', 'r', encoding='utf-8') as f:
    text = f.read()

# Character-level tokenizer
chars = sorted(list(set(text)))
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])

print(f"Vocabulary size: {len(chars)}")
print(f"Dataset size: {len(text):,} characters")

What does this code do?

  • chars extracts every unique character from the text and sorts them. This becomes your vocabulary (65 characters: letters, digits, punctuation, newlines).
  • stoi (string-to-integer) maps each character to a unique number. itos (integer-to-string) maps numbers back to characters.
  • encode converts a string into a list of integers. decode reverses the process, turning integers back into readable text.
  • Save train.py.

You cannot run this file yet because model.py does not exist. You will verify the tokenizer works at the end of this step using a quick test command.

Seeing a red underline on the model import?

Your editor may flag from model import GPT, block_size, n_embd, vocab_size as an error because model.py does not exist yet. This is expected. The warning will disappear after Step 3.

Why is my editor showing import errors for model.py?

Create the train/val split and batch generator

With the tokenizer ready, you need to split the data into training and validation sets, then write a function that grabs random chunks of text as tensors for the model to learn from.

  • Below the print statements in train.py, add the train/val split:
# Train/val split
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]

What does this code do?

  • encode(text) turns the entire Shakespeare corpus into a list of integers using the tokenizer you just built.
  • torch.tensor(..., dtype=torch.long) converts those integers into a PyTorch tensor. The long dtype is required for embedding lookups.
  • The first 90% of characters become training data. The last 10% is held out for validation, so you can check whether the model is memorizing or genuinely learning patterns.
  • Below the split, add the batch generator function:

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])
    x, y = x.to(device), y.to(device)
    return x, y

What does this code do?

get_batch is the function your training loop will call repeatedly to fetch random chunks of text.

  • ix = torch.randint(...) picks 64 random starting positions in the dataset.
  • x is a batch of input sequences (each 256 characters long). y is the same sequences shifted one position forward. The model learns to predict each next character.
  • .to(device) moves the tensors onto your M4 GPU for fast computation.
  • Save train.py.

Unsure about the blank line before get_batch?

Python style uses two blank lines before top-level function definitions. Make sure there is one blank line between val_data = data[n:] and def get_batch(split):. Your code will work either way, but matching the structure helps with readability.

Help me check my file structure

✔️ Awesome, I've got everything!

Great. Make sure you have saved train.py with all the code above.

ⓧ I'd like to double check the full code

Here is the complete train.py as it should look at the end of this step:

import torch
import math
from model import GPT, block_size, n_embd, vocab_size

# Training hyperparameters
batch_size = 64
max_iters = 3000
eval_interval = 200
eval_iters = 100
learning_rate = 3e-4
min_lr = 3e-5
warmup_iters = 100
device = 'mps' if torch.backends.mps.is_available() else 'cpu'

print(f"Using device: {device}")

# Load data
with open('data/shakespeare.txt', 'r', encoding='utf-8') as f:
    text = f.read()

# Character-level tokenizer
chars = sorted(list(set(text)))
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])

print(f"Vocabulary size: {len(chars)}")
print(f"Dataset size: {len(text):,} characters")

# Train/val split
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]


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])
    x, y = x.to(device), y.to(device)
    return x, y
Verify the tokenizer and batch pipeline

Since train.py cannot run directly yet (it imports from model.py), you will test the tokenizer and batching logic with a quick inline Python command.

  • Verify the tokenizer and batch shapes by running this command in your terminal:
python3.12 -c "
import torch
with open('data/shakespeare.txt', 'r', encoding='utf-8') as f:
    text = f.read()
chars = sorted(list(set(text)))
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])
print(f'Vocabulary size: {len(chars)}')
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
block_size = 256
batch_size = 64
device = 'mps' if torch.backends.mps.is_available() else 'cpu'
ix = torch.randint(len(train_data) - block_size, (batch_size,))
x = torch.stack([train_data[i:i + block_size] for i in ix]).to(device)
y = torch.stack([train_data[i + 1:i + block_size + 1] for i in ix]).to(device)
print(f'Batch x shape: {x.shape}')
print(f'Batch y shape: {y.shape}')
print(f'Decoded sample: {decode(x[0,:80].tolist())}')
"
  • You should see output similar to:
Vocabulary size: 65
Batch x shape: torch.Size([64, 256])
Batch y shape: torch.Size([64, 256])
Decoded sample: <a chunk of readable Shakespeare text>

What should I see?

  • Vocabulary size: 65 confirms your tokenizer found exactly 65 unique characters in the Shakespeare text (letters, digits, punctuation, spaces, and newlines).
  • torch.Size([64, 256]) means you have a batch of 64 sequences, each 256 tokens long. This is the shape your model will expect.
  • The decoded sample should be a recognizable chunk of Shakespeare. If you see readable English text (character names, dialogue), your encode/decode round-trip is working correctly.

Seeing garbled text or a KeyError?

If the decoded sample is garbled, check that your data/shakespeare.txt file downloaded correctly in Step 1. Re-download it if the file size is significantly smaller than ~1MB.

A KeyError means a character in the text is not in your vocabulary mapping. Make sure chars is built from the same text variable you are encoding.

Help me debug my tokenizer test

Your data pipeline is complete. Every time the training loop needs a fresh batch, get_batch will grab 64 random 256-character sequences and load them onto your GPU. Next up, you will build the transformer model that learns patterns from these batches.

Write the Transformer Model

Your data pipeline is ready. You have a tokenizer that converts Shakespeare into numbers, a train/val split, and a batch generator feeding sequences to your GPU.

Now comes the core of this project: writing the transformer model itself. This is the architecture that powers systems like ChatGPT and Claude. You will build every component from scratch in a single file called model.py.

In this step, get ready to:
  • Implement the building blocks: token and positional embeddings, a single self-attention head, and multi-head attention.
  • Build a full transformer block and stack six of them into a complete GPT model with a language model head.
  • Instantiate the model, print the parameter count (~10M), and verify a forward pass returns the expected output shape.
Implement the building blocks

Self-attention is the mechanism that lets the model look at other positions in a sequence when encoding each token. You will start by defining the model's hyperparameters, then build a single attention head, and finally combine multiple heads in parallel.

  • Create a new file called model.py in the root of your gpt-from-scratch folder.
  • Add the imports and hyperparameters by pasting this code:
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# Hyperparameters
block_size = 256
n_embd = 384
n_head = 6
n_layer = 6
dropout = 0.2
vocab_size = 65  # character-level: printable ASCII chars in Shakespeare

What do these hyperparameters control?

These six values define the entire model's shape. block_size is the maximum context length (256 characters the model can look back). n_embd is the width of each hidden representation (384 dimensions).

The model uses 6 attention heads (n_head) across 6 stacked layers (n_layer). dropout randomly zeros 20% of values during training to prevent overfitting. vocab_size matches the 65 unique characters found in the Shakespeare text.

  • Below the hyperparameters, add the Head class for a single self-attention head:
class Head(nn.Module):
    """Single 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)
        # Compute 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)
        # Weighted aggregation of values
        v = self.value(x)
        out = wei @ v
        return out

What does this code do?

Each token produces three vectors: a Query ("what am I looking for?"), a Key ("what do I contain?"), and a Value ("what do I communicate?"). The attention score between two tokens is the dot product of Query and Key, scaled by the square root of the head dimension.

The tril buffer is a lower-triangular matrix that acts as a causal mask. It prevents each token from attending to future positions. This is what makes the model autoregressive.

  • Below the Head class, add the MultiHeadAttention class that runs multiple heads in parallel:
class MultiHeadAttention(nn.Module):
    """Multiple heads of self-attention in parallel."""

    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)
        out = self.dropout(self.proj(out))
        return out

Why multiple heads?

A single attention head with dimension 384 can only learn one type of relationship at a time. By splitting into 6 heads of 64 dimensions each, the model learns 6 different relationship patterns simultaneously (grammar, meaning, rhythm, etc.). The outputs are concatenated and projected back to the full embedding dimension.

  • Save model.py.
  • Verify the file has no syntax errors by running this command:
python3.12 -c "from model import Head, MultiHeadAttention; print('Attention classes OK')"
  • You should see Attention classes OK printed with no errors.

Seeing a SyntaxError or IndentationError?

Check that the Head class uses consistent indentation (4 spaces per level). Mixing tabs and spaces causes Python to reject the file.

Make sure the register_buffer line is inside __init__, not inside forward.

Help me debug my model.py syntax error

Build the transformer block and full GPT model

With attention implemented, you now need two more pieces: a feed-forward network that processes each position independently, and a transformer block that combines attention with the feed-forward network using residual connections and layer normalization.

  • In model.py, below the MultiHeadAttention class, add the FeedForward class:
class FeedForward(nn.Module):
    """Feed-forward network with GELU activation."""

    def __init__(self):
        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)

What does this code do?

The feed-forward network expands each token's representation to 4x its size (384 to 1536), applies GELU activation (a smooth, non-linear function), then projects back down. This gives the model capacity to learn complex transformations at each position independently.

  • Below FeedForward, add the Block class that combines everything into one transformer layer:
class Block(nn.Module):
    """Transformer block: layernorm -> attention -> residual -> layernorm -> ffn -> residual."""

    def __init__(self):
        super().__init__()
        head_size = n_embd // n_head
        self.ln1 = nn.LayerNorm(n_embd)
        self.sa = MultiHeadAttention(n_head, head_size)
        self.ln2 = nn.LayerNorm(n_embd)
        self.ffwd = FeedForward()

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x

What are residual connections?

The x = x + self.sa(self.ln1(x)) pattern adds the attention output back to the original input. These "skip connections" let gradients flow directly through the network during training, which is essential for training deep models (6 layers deep in our case).

Layer normalization (ln1, ln2) stabilizes training by normalizing each token's representation to have zero mean and unit variance before passing it to attention or the feed-forward network.

  • Now add the full GPT class below the Block class. This is the top-level model that ties everything together:
class GPT(nn.Module):
    """GPT language model."""

    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() for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    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=idx.device))
        x = tok_emb + pos_emb
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits_flat = logits.view(B * T, C)
            targets_flat = targets.view(B * T)
            loss = F.cross_entropy(logits_flat, targets_flat)

        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

What does this code do?

The GPT class has four major parts:

  • Embeddings: token_embedding converts each character ID into a 384-dimensional vector. position_embedding adds positional information so the model knows where each token sits in the sequence.
  • Forward pass: tokens flow through 6 stacked transformer blocks, a final layer norm, and a linear layer (lm_head) that produces logits (scores) for each of the 65 characters at every position.
  • Loss: when targets are provided, the model computes cross-entropy loss by comparing its predictions against the actual next characters.
  • Generation: the generate method autoregressively produces new tokens by sampling from the model's predicted distribution with optional temperature scaling and top-k filtering.
  • Save model.py.

You won't see a visible change yet. The next substep verifies everything works with a forward pass.

Seeing a NameError when saving?

Make sure the classes appear in this exact order in your file: Head, MultiHeadAttention, FeedForward, Block, GPT. Each class depends on the ones above it. If Block appears before FeedForward, Python will raise a NameError.

Help me fix a NameError in my model.py file

Instantiate the model and verify the forward pass

Your architecture is complete. Before moving on, you need to confirm the model instantiates correctly, has roughly 10 million parameters, and produces output tensors of the expected shape.

  • Run this verification script in your terminal:
python3.12 -c "
import torch
from model import GPT, block_size, vocab_size

device = 'mps' if torch.backends.mps.is_available() else 'cpu'
model = GPT().to(device)

# Print parameter count
params = sum(p.numel() for p in model.parameters())
print(f'Model parameters: {params:,}')

# Test forward pass with a dummy batch
batch_size = 4
x = torch.randint(0, vocab_size, (batch_size, block_size), device=device)
logits, loss = model(x)
print(f'Logits shape: {logits.shape}')
print(f'Expected:     torch.Size([{batch_size}, {block_size}, {vocab_size}])')
print('Forward pass OK!' if logits.shape == (batch_size, block_size, vocab_size) else 'SHAPE MISMATCH')
"
  • You should see output showing approximately 10 million parameters and a logits shape of torch.Size([4, 256, 65]) with the message Forward pass OK!.

What do these numbers mean?

The ~10M parameters are all the learnable weights across 6 transformer blocks, 2 embedding tables, and the output head. The logits shape (4, 256, 65) means: for each of the 4 sequences in the batch, at each of the 256 positions, the model outputs 65 scores (one per character in the vocabulary).

The highest-scoring character at each position is the model's best guess for what comes next. Right now those guesses are random because the model is untrained.

Seeing an error during the forward pass?

If you see a shape mismatch error mentioning mat1 and mat2, check that n_embd is set to 384 and n_head is set to 6 at the top of model.py. The head size must divide evenly (384 / 6 = 64).

If you get an import error saying the module model cannot be found, make sure you are running the command from inside the gpt-from-scratch folder where model.py lives.

Help me debug my GPT model forward pass

✔️ Awesome, I've got everything!

Great. Double check you have saved model.py with all five classes in the correct order.

ⓧ I'd like to double check the full code

Here is the complete model.py file for comparison:

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# Hyperparameters
block_size = 256
n_embd = 384
n_head = 6
n_layer = 6
dropout = 0.2
vocab_size = 65  # character-level: printable ASCII chars in Shakespeare

class Head(nn.Module):
    """Single 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)
        # Compute 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)
        # Weighted aggregation of values
        v = self.value(x)
        out = wei @ v
        return out


class MultiHeadAttention(nn.Module):
    """Multiple heads of self-attention in parallel."""

    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)
        out = self.dropout(self.proj(out))
        return out


class FeedForward(nn.Module):
    """Feed-forward network with GELU activation."""

    def __init__(self):
        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: layernorm -> attention -> residual -> layernorm -> ffn -> residual."""

    def __init__(self):
        super().__init__()
        head_size = n_embd // n_head
        self.ln1 = nn.LayerNorm(n_embd)
        self.sa = MultiHeadAttention(n_head, head_size)
        self.ln2 = nn.LayerNorm(n_embd)
        self.ffwd = FeedForward()

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x


class GPT(nn.Module):
    """GPT language model."""

    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() for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    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=idx.device))
        x = tok_emb + pos_emb
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits_flat = logits.view(B * T, C)
            targets_flat = targets.view(B * T)
            loss = F.cross_entropy(logits_flat, targets_flat)

        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

Your transformer architecture is complete. You have a full GPT model ready to learn. Next up, you will write the training loop that teaches this model to predict Shakespeare one character at a time.

Train the Model

In the last step, you built a complete transformer architecture with ~10M parameters. The model exists, but right now its weights are random noise. It has no understanding of English, Shakespeare, or even what a word is.

Training is the process of showing the model millions of examples from Shakespeare and adjusting those random weights until it learns the patterns. In this step, you will write the training loop that makes your model learn and then watch the loss curve converge over ~3000 iterations.

In this step, get ready to:
  • Set up the training infrastructure: loss estimation, learning rate schedule, and optimizer.
  • Run training for 3000 iterations and save a checkpoint.
  • Understand what the loss numbers mean and how they connect to model quality.
Set up the training infrastructure

Before the training loop itself, you need three supporting pieces: a function to reliably measure how well the model is performing, a learning rate schedule that adjusts the step size over time, and an optimizer that knows how to update the weights.

  • In train.py, find the get_batch function you wrote in the last step. Add the following estimate_loss function below it:
@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            _, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

What does this code do?

Training loss on a single batch is noisy because each batch is a random sample. This function averages the loss over 100 batches to get a stable reading.

  • @torch.no_grad() disables gradient tracking during evaluation, saving memory and compute.
  • model.eval() turns off dropout so the evaluation is deterministic.
  • model.train() at the end re-enables dropout for the next training iteration.
  • Add the learning rate schedule function below estimate_loss:
def get_lr(it):
    """Cosine learning rate schedule with linear warmup."""
    if it < warmup_iters:
        return learning_rate * (it + 1) / warmup_iters
    if it >= max_iters:
        return min_lr
    decay_ratio = (it - warmup_iters) / (max_iters - warmup_iters)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (learning_rate - min_lr)

Why use a cosine learning rate schedule?

A fixed learning rate either starts too slow (wasting early iterations) or ends too fast (overshooting once the model is close to a good solution). The cosine schedule solves both problems.

During warmup (first 100 steps), the rate ramps up linearly from near-zero to the full value. This prevents unstable gradients early on when weights are random. After warmup, the rate smoothly decays following a cosine curve, making smaller and smaller adjustments as the model converges.

  • Add the model initialization and optimizer setup below the get_lr function:
# Initialize model
model = GPT()
model = model.to(device)
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

# Optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.1)

What does this code do?

  • GPT() creates your model with all the layers you defined in model.py.
  • model.to(device) moves all weights to your M4 GPU via MPS.
  • AdamW is an optimizer that combines adaptive learning rates per parameter with weight decay. The weight_decay=0.1 gently pushes weights toward zero, preventing the model from memorizing the training data.
  • Save train.py.

Seeing a red underline on model in estimate_loss?

Your editor may warn that model is undefined inside estimate_loss(). This is normal. The function references model as a module-level variable that gets created later in the file. Python resolves it at runtime, not at definition time.

help me understand why estimate_loss references model before it's defined

Write the training loop and run it

The training loop is where learning happens. Each iteration, the model sees a batch of Shakespeare, predicts the next character, measures how wrong it was, and adjusts its weights to be less wrong next time.

  • In train.py, add the training loop below the optimizer line:
# Training loop
print("\nStarting training...")
for iter in range(max_iters):
    # Update learning rate
    lr = get_lr(iter)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    # Evaluate periodically
    if iter % eval_interval == 0 or iter == max_iters - 1:
        losses = estimate_loss()
        print(f"Step {iter:4d} | train loss {losses['train']:.4f} | val loss {losses['val']:.4f} | lr {lr:.2e}")

    # Forward pass
    xb, yb = get_batch('train')
    logits, loss = model(xb, yb)

    # Backward pass
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()

# Save checkpoint
print("\nSaving checkpoint...")
torch.save(model.state_dict(), 'checkpoint.pt')
print("Training complete! Checkpoint saved to checkpoint.pt")

# Generate a sample
print("\n--- Sample generation ---")
context = torch.zeros((1, 1), dtype=torch.long, device=device)
generated = model.generate(context, max_new_tokens=500, temperature=0.8, top_k=40)
print(decode(generated[0].tolist()))

What does this code do?

Each iteration of the loop performs one complete learning cycle:

  • get_lr(iter) computes the current learning rate from the cosine schedule and applies it to the optimizer.
  • Every 200 steps, estimate_loss() runs to print a stable train/val loss reading.
  • The forward pass computes predictions and cross-entropy loss (how wrong the model is).
  • optimizer.zero_grad(set_to_none=True) clears old gradients before computing new ones.
  • clip_grad_norm_(model.parameters(), 1.0) caps gradient magnitude to prevent exploding gradients from destabilizing training.
  • After the loop, the trained weights are saved to checkpoint.pt and a 500-character sample is generated to preview the model's output.
  • Save train.py.

✔️ Awesome, I've got everything!

Great. Double check you have saved train.py before running the training script.

ⓧ I'd like to double check the full code

import torch
import math
from model import GPT, block_size, n_embd, vocab_size

# Training hyperparameters
batch_size = 64
max_iters = 3000
eval_interval = 200
eval_iters = 100
learning_rate = 3e-4
min_lr = 3e-5
warmup_iters = 100
device = 'mps' if torch.backends.mps.is_available() else 'cpu'

print(f"Using device: {device}")

# Load data
with open('data/shakespeare.txt', 'r', encoding='utf-8') as f:
    text = f.read()

# Character-level tokenizer
chars = sorted(list(set(text)))
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])

print(f"Vocabulary size: {len(chars)}")
print(f"Dataset size: {len(text):,} characters")

# Train/val split
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]


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])
    x, y = x.to(device), y.to(device)
    return x, y


@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            _, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out


def get_lr(it):
    """Cosine learning rate schedule with linear warmup."""
    if it < warmup_iters:
        return learning_rate * (it + 1) / warmup_iters
    if it >= max_iters:
        return min_lr
    decay_ratio = (it - warmup_iters) / (max_iters - warmup_iters)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (learning_rate - min_lr)


# Initialize model
model = GPT()
model = model.to(device)
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

# Optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.1)

# Training loop
print("\nStarting training...")
for iter in range(max_iters):
    # Update learning rate
    lr = get_lr(iter)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    # Evaluate periodically
    if iter % eval_interval == 0 or iter == max_iters - 1:
        losses = estimate_loss()
        print(f"Step {iter:4d} | train loss {losses['train']:.4f} | val loss {losses['val']:.4f} | lr {lr:.2e}")

    # Forward pass
    xb, yb = get_batch('train')
    logits, loss = model(xb, yb)

    # Backward pass
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()

# Save checkpoint
print("\nSaving checkpoint...")
torch.save(model.state_dict(), 'checkpoint.pt')
print("Training complete! Checkpoint saved to checkpoint.pt")

# Generate a sample
print("\n--- Sample generation ---")
context = torch.zeros((1, 1), dtype=torch.long, device=device)
generated = model.generate(context, max_new_tokens=500, temperature=0.8, top_k=40)
print(decode(generated[0].tolist()))
  • Start training by running this command in your terminal:
python train.py

Training will take approximately 20-25 minutes on an M4 Mac. You will see loss printed every 200 steps.

What should I see?

The first few lines will print your device (mps), vocabulary size (65), dataset size, and parameter count (~10M). Then the training loop starts printing loss every 200 steps.

The initial loss will be around 4.1 (which is basically random guessing across 65 characters). Over 3000 steps, you should see it drop steadily toward 1.5 or lower. After training completes, a 500-character Shakespeare-like sample will print automatically.

Training seems stuck or very slow?

  • Check that the first line says Using device: mps. If it says cpu, the GPU is not being used and training will be much slower.
  • Close other GPU-heavy apps (browsers with many tabs, video editors) to free up Metal resources.
  • If you see a RuntimeError about MPS, make sure your macOS is up to date and that you are using the virtual environment where PyTorch is installed.

help me debug my GPT training script that isn't running correctly on MPS

Understand what the loss numbers mean

While your model trains (or after it finishes), take a moment to understand what the numbers scrolling by actually tell you about your model's intelligence.

What do the loss values mean?

The loss is cross-entropy loss, which measures how surprised the model is by the correct answer. Lower loss means less surprise, which means better predictions.

  • Loss ~4.17 (starting point): This equals ln(65), meaning the model is guessing uniformly across all 65 characters. No learning has occurred yet.
  • Loss ~2.5: The model has learned basic character frequencies. It knows spaces and 'e' are common, but cannot form words.
  • Loss ~1.5: The model generates recognizable English words, verse structure, and character names. This is your target.

You can convert loss to perplexity by computing e^loss. A loss of 1.5 means a perplexity of ~4.5, so on average the model is choosing between about 4-5 equally likely next characters at each position.

Why track both train and val loss?

Train loss measures performance on data the model has seen. Validation loss measures performance on the held-out 10% it has never seen. If train loss drops but val loss starts rising, the model is memorizing instead of learning patterns.

For this project, both losses should drop together and end up close to each other (within ~0.05). The dropout and weight decay you configured prevent overfitting.

  • Wait for training to complete. Your final output should show a validation loss below 1.6 and a generated Shakespeare-like text sample.
  • Confirm that checkpoint.pt was saved by running:
ls -lh checkpoint.pt

You should see a file of approximately 40MB.

Final loss above 1.6?

  • Make sure your max_iters is set to 3000. Fewer iterations means less learning time.
  • Verify that learning_rate is 3e-4 (not 3e-5). A rate that is too low means the model learns too slowly.
  • Check that model.py has n_layer = 6 and n_embd = 384. Smaller values produce a weaker model.

help me figure out why my GPT training loss isn't converging below 1.6

Your model just went from random noise to Shakespeare in 3000 steps. The checkpoint sitting on your disk contains everything the model learned about English structure, verse patterns, and character dialogue. Next up, you will build a dedicated generation script that loads this checkpoint and lets you control exactly how the model creates text.

Generate Text

Your model has been trained and the loss curve has converged. The weights sitting in checkpoint.pt encode everything the model learned about Shakespeare's writing patterns.

Now you need a way to actually talk to your model. In this step, you will build a standalone generation script that loads the trained checkpoint, accepts a text prompt, and produces new Shakespeare-like text on demand.

In this step, get ready to:
  • Write an autoregressive generation script that loads your trained model and produces text.
  • Understand how temperature and top-k sampling control output quality.
  • Generate Shakespeare-like text with different settings and compare the results.
Write the autoregressive generation loop

Autoregressive generation means producing text one token at a time. The model predicts the next character, you append it to the sequence, and feed the extended sequence back in. Your model.py already has a generate() method that does this loop internally.

The script you are about to write wraps that method with a CLI so you can experiment with prompts and settings from the terminal.

  • Create a new file called generate.py in your project folder.
  • Add the CLI argument parser by pasting this code:
import argparse
import torch
from model import GPT, block_size

# Parse arguments
parser = argparse.ArgumentParser(description='Generate text from trained GPT model')
parser.add_argument('--prompt', type=str, default='', help='Starting text prompt')
parser.add_argument('--tokens', type=int, default=500, help='Number of tokens to generate')
parser.add_argument('--temperature', type=float, default=0.8, help='Sampling temperature')
parser.add_argument('--top-k', type=int, default=40, help='Top-k filtering')
parser.add_argument('--checkpoint', type=str, default='checkpoint.pt', help='Path to model checkpoint')
args = parser.parse_args()

device = 'mps' if torch.backends.mps.is_available() else 'cpu'

What does this code do?

This sets up a command-line interface using Python's built-in argparse module. Each add_argument call defines a flag you can pass when running the script.

  • --prompt sets the starting text the model continues from. Defaults to empty (random start).
  • --tokens controls how many characters to generate. Defaults to 500.
  • --temperature and --top-k control the randomness of sampling (more on these in the next substep).
  • --checkpoint points to your saved model weights file.
  • Below the device line, add the tokenizer loading, model loading, and generation logic:
# Load data for tokenizer
with open('data/shakespeare.txt', 'r', encoding='utf-8') as f:
    text = f.read()

chars = sorted(list(set(text)))
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])

# Load model
model = GPT()
model.load_state_dict(torch.load(args.checkpoint, map_location=device, weights_only=True))
model = model.to(device)
model.eval()

# Encode prompt
if args.prompt:
    context = torch.tensor([encode(args.prompt)], dtype=torch.long, device=device)
else:
    context = torch.zeros((1, 1), dtype=torch.long, device=device)

# Generate
with torch.no_grad():
    generated = model.generate(context, max_new_tokens=args.tokens, temperature=args.temperature, top_k=args.top_k)

print(decode(generated[0].tolist()))

What does this code do?

This is the generation pipeline. It rebuilds the same character-level tokenizer from the training data, loads your trained model weights, and runs the autoregressive loop.

  • The tokenizer is rebuilt from data/shakespeare.txt to ensure the character-to-index mapping matches what the model was trained on.
  • model.eval() disables dropout so every forward pass is deterministic given the same random seed.
  • torch.no_grad() disables gradient tracking since you are only running inference, saving memory.
  • If no prompt is provided, the model starts from a zero token (essentially a blank canvas) and generates freely.
  • Save generate.py.
  • Test that the script runs by executing this in your terminal:
python generate.py --prompt "ROMEO:" --tokens 200

You should see approximately 200 characters of Shakespeare-like text printed to your terminal, starting from the prompt ROMEO:.

Seeing an error about the checkpoint?

Make sure you are running the command from your project folder where checkpoint.pt lives. If you moved the file, pass the correct path with --checkpoint.

If you see a KeyError about a character, make sure data/shakespeare.txt is the same file you trained on. The tokenizer must match exactly.

Help me debug my generate.py script

Understand temperature and top-k sampling

Your model's generate() method in model.py already implements both temperature scaling and top-k filtering. These two parameters give you direct control over how creative or conservative the generated text is.

How does temperature work?

After the model produces logits (raw scores for each character), temperature divides those scores before converting them to probabilities. This changes the shape of the probability distribution.

  • Temperature close to 0 (e.g. 0.1) makes the distribution very peaked. The model almost always picks the single most likely character. Output reads like a repetitive loop.
  • Temperature = 1.0 leaves the distribution unchanged. The model samples proportionally to what it learned.
  • Temperature greater than 1.0 flattens the distribution, making unlikely characters more probable. Output becomes more random and creative but also more incoherent.

The sweet spot for this model is around 0.8. It adds enough variety to avoid repetition while staying close to learned patterns.

How does top-k filtering work?

Top-k filtering zeroes out all but the k most likely tokens before sampling. This prevents the model from ever picking extremely unlikely characters, even when temperature is high.

With top_k=40, the model only considers the 40 most probable next characters at each step. The remaining 25 characters in the vocabulary are set to negative infinity (zero probability). This acts as a safety net that keeps output coherent.

  • Look at the generate() method in model.py. Find these lines that implement both controls:
logits = logits[:, -1, :] / temperature
if top_k is not None:
    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
    logits[logits < v[:, [-1]]] = float('-inf')
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)

What's happening in this sequence?

  • Line 1 divides logits by temperature, reshaping the distribution.
  • Lines 2-4 find the top-k values, then set everything below the k-th highest to negative infinity.
  • Line 5 converts the filtered logits to probabilities with softmax.
  • Line 6 randomly samples one token from the resulting distribution.
Generate text with different settings

Now you will see temperature and top-k in action by generating text with different configurations and comparing the results.

  • Generate near-greedy text (very low temperature) by running:
python generate.py --prompt "ROMEO:" --tokens 300 --temperature 0.1 --top-k 40

Notice how the output is repetitive. Low temperature forces the model to always pick the most likely next character, creating loops.

  • Generate creative text (temperature = 1.0) by running:
python generate.py --prompt "ROMEO:" --tokens 300 --temperature 1.0 --top-k 40

This output has more variety. You should see invented words and unexpected phrasing mixed with recognizable Shakespeare patterns.

  • Try reducing top-k to see how vocabulary restriction affects output:
python generate.py --prompt "ROMEO:" --tokens 300 --temperature 0.8 --top-k 10

With only 10 candidate characters at each step, the output stays tightly constrained to common letter sequences.

  • Finally, generate 500 characters with the balanced defaults to see your model's best output:
python generate.py --prompt "ROMEO:" --tokens 500 --temperature 0.8 --top-k 40

You should see recognizable Shakespeare-like structure. Look for verse form, character names in all caps followed by colons, and iambic rhythm in the lines.

What makes good output?

Your model will not produce perfect Shakespeare. It has learned character-level patterns from only 1MB of text. Good output at this scale means recognizable English words arranged in verse form with occasional character names and stage directions.

The fact that coherent words appear at all is remarkable. The model never saw "words" as a concept. It learned that certain character sequences are more likely than others, purely from statistical patterns.

Output looks like random characters?

If the output is gibberish, your training may not have converged. Check that your final validation loss was below 1.6. If it was higher, run training for more iterations.

Also verify you are loading the correct checkpoint. If you have multiple checkpoint.pt files from different training runs, pass the right one with --checkpoint.

Help me figure out why my generated text is not coherent

✔️ Awesome, I've got everything!

Great. Double check you have saved generate.py and can run it without errors.

ⓧ I'd like to double check the full code

import argparse
import torch
from model import GPT, block_size

# Parse arguments
parser = argparse.ArgumentParser(description='Generate text from trained GPT model')
parser.add_argument('--prompt', type=str, default='', help='Starting text prompt')
parser.add_argument('--tokens', type=int, default=500, help='Number of tokens to generate')
parser.add_argument('--temperature', type=float, default=0.8, help='Sampling temperature')
parser.add_argument('--top-k', type=int, default=40, help='Top-k filtering')
parser.add_argument('--checkpoint', type=str, default='checkpoint.pt', help='Path to model checkpoint')
args = parser.parse_args()

device = 'mps' if torch.backends.mps.is_available() else 'cpu'

# Load data for tokenizer
with open('data/shakespeare.txt', 'r', encoding='utf-8') as f:
    text = f.read()

chars = sorted(list(set(text)))
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])

# Load model
model = GPT()
model.load_state_dict(torch.load(args.checkpoint, map_location=device, weights_only=True))
model = model.to(device)
model.eval()

# Encode prompt
if args.prompt:
    context = torch.tensor([encode(args.prompt)], dtype=torch.long, device=device)
else:
    context = torch.zeros((1, 1), dtype=torch.long, device=device)

# Generate
with torch.no_grad():
    generated = model.generate(context, max_new_tokens=args.tokens, temperature=args.temperature, top_k=args.top_k)

print(decode(generated[0].tolist()))

You now have a complete text generation pipeline. Your model can take any prompt and continue it in the style of Shakespeare, with full control over creativity through temperature and top-k settings.

Secret mission

Implement Nucleus (Top-p) Sampling

Your top-k sampling always keeps exactly 40 tokens, whether the model is 99% sure of the next character or completely uncertain. Can you implement a smarter strategy that adapts to the model's confidence level?

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely locally, so there are no ongoing costs.

Resources you used:

  • Python virtual environment (local disk only).
  • Trained model checkpoint file checkpoint.pt (~40MB on disk).
  • Shakespeare text file data/shakespeare.txt (~1MB).
  • Project source files (model.py, train.py, generate.py).

Keep everything running

No action needed. Choose this if you want to keep experimenting with hyperparameters, different prompts, or generation settings.

  • Your project files, virtual environment, and trained checkpoint all stay in place.
  • There are no running processes or background services to manage.
  • Total disk usage is approximately 42MB (mostly the checkpoint).

Pause - I'll come back to this later

There are no running processes to shut down. Everything in this project is just files on disk.

  • Your trained checkpoint and source files will be ready whenever you come back.
  • To resume later, activate the virtual environment again and run generate.py to pick up where you left off.

Delete - I don't want to use this again

Remove all project files, the virtual environment, and the trained checkpoint.

  • If your virtual environment is currently active, deactivate it first by running:
deactivate
  • Delete the entire project directory (which includes the virtual environment, checkpoint, data, and source files) by running:
rm -rf ~/gpt-from-scratch
  • Confirm the directory is gone by running:
ls ~/gpt-from-scratch

You should see a "No such file or directory" message confirming the deletion.

That's a Wrap!

That's a Wrap!

Nice work! You just built a working GPT language model entirely from scratch and trained it to generate Shakespeare-like text on your own Mac.

You've learned how to:

  • Build a character-level tokenizer and understand why it suits small datasets, then write a complete transformer architecture from scratch: embeddings, multi-head self-attention with causal masking, feed-forward networks, layer normalization, and residual connections.
  • Train a ~10M parameter model on Shakespeare using MPS GPU acceleration with AdamW optimizer, cosine learning rate scheduling, and gradient clipping.
  • Generate Shakespeare-like text using temperature scaling and top-k sampling to control output creativity and coherence.
  • Secret Mission: Implemented nucleus (top-p) sampling and compared it against fixed-k truncation to see how dynamic vocabulary selection adapts to varying model confidence.

Ready to quiz yourself?