Build a GPT Model from Scratch
Implement a transformer-based GPT in PyTorch and generate Shakespeare-style text.
Introduction
30 Second Summary
Every time you ask ChatGPT or Claude a question, a language model is predicting the next word in the sequence. These models feel like magic, but they are built from a repeatable set of components that you can understand and implement yourself.
In this project, you will build a miniature GPT language model from scratch using PyTorch. You will implement every piece of the transformer architecture and train it to generate Shakespeare-style text.
What You'll Build
A character-level GPT language model implemented entirely from scratch in a single Python file, trained on Shakespeare's collected works.
By the end of this project, you'll have:
- A working transformer architecture with self-attention, feedforward layers, normalization, and residual connections.
- A trained language model that generates coherent Shakespeare-style text after a few minutes of training.
- Deep understanding of every component inside the models that power ChatGPT and Claude.
- Secret Mission: Modernize your transformer with RMSNorm and GELU activation to see how small architectural changes improve output quality.
Are there any prerequisites?
You should be comfortable writing basic Python (variables, functions, loops). Conceptual familiarity with neural networks is helpful but not required.
You do not need prior experience with PyTorch or transformers. This project teaches both from the ground up.
Set Up Your Project
Before you build the transformer architecture, you need training data and a way to convert text into numbers a neural network can process.
In this step, you will create the foundation of your GPT model file. You will download the Tiny Shakespeare dataset, build a character-level tokenizer, and split the data into training and validation sets.
In this step, get ready to:
- Create gpt.py and define the model's hyperparameters.
- Download and load the Tiny Shakespeare dataset.
- Build a character-level tokenizer and prepare training data splits.
Create your model file and define hyperparameters
Every configuration value your GPT needs lives at the top of a single Python file. These hyperparameters control the model's size, training duration, and learning speed.
- In your code editor, create a new file called gpt.py in your project folder.
- Add the following imports at the top of gpt.py:
import torch
import torch.nn as nn
from torch.nn import functional as F
import urllib.request
import osWhat do these imports do?
- torch is the core PyTorch library for tensor operations and automatic differentiation.
- torch.nn provides neural network building blocks (layers, loss functions, containers).
- torch.nn.functional gives you stateless functions like softmax and cross-entropy.
- urllib.request and os handle downloading the dataset and checking file paths.
- Save gpt.py.
You won't see any output yet. The next chunk adds the configuration these imports support.
- Below the imports in gpt.py, add the hyperparameters block:
# Hyperparameters
batch_size = 64
block_size = 256
max_iters = 5000
eval_interval = 500
learning_rate = 3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 384
n_head = 6
n_layer = 6
dropout = 0.2
torch.manual_seed(1337)What do these hyperparameters control?
- batch_size sets how many text sequences the model processes in parallel during each training step.
- block_size is the maximum context length (256 characters the model can "look back" at).
- n_embd, n_head, and n_layer define the model's size: 384-dimensional embeddings, 6 attention heads, and 6 transformer layers.
- device automatically selects GPU if available, otherwise falls back to CPU.
- torch.manual_seed(1337) makes results reproducible by fixing the random number generator.
- Save gpt.py.
You won't see output yet. The next substep downloads the training data that these settings will govern.
Download the Tiny Shakespeare dataset
Your model needs text to learn from. The Tiny Shakespeare dataset is about 1 MB of Shakespeare's collected works, containing 65 unique characters.
- Below the torch.manual_seed(1337) line in gpt.py, add the dataset download and load code:
# Download and load the dataset
data_path = 'input.txt'
if not os.path.exists(data_path):
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
urllib.request.urlretrieve(url, data_path)
with open(data_path, 'r', encoding='utf-8') as f:
text = f.read()What does this code do?
The script checks whether input.txt already exists in your project folder. If it does not, it downloads the full Tiny Shakespeare text from GitHub and saves it locally.
After the download check, the file is opened and its entire contents are stored in the text variable as one long string. This avoids re-downloading every time you run the script.
- Save gpt.py.
- Run the script to trigger the download by running this command in your terminal:
python gpt.py- Confirm the dataset downloaded by running:
lsYou should see input.txt alongside gpt.py in the output.
Don't see input.txt?
- Check that your terminal is in the same directory as gpt.py. Run pwd to confirm.
- If you see a ModuleNotFoundError, make sure your virtual environment is activated.
- If you see a network error, check your internet connection and try running the script again.
Help me debug why input.txt didn't download when I ran gpt.py.
Build the character-level tokenizer and prepare data splits
Neural networks only understand numbers, not letters. A tokenizer converts each character into a unique integer so the model can process text mathematically.
- Below the dataset loading code in gpt.py, add the character-level tokenizer:
# Character-level tokenizer
chars = sorted(list(set(text)))
vocab_size = len(chars)
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])What does this tokenizer do?
chars extracts every unique character from the text and sorts them alphabetically. For Tiny Shakespeare, this gives you 65 characters (letters, punctuation, newlines, and spaces).
stoi (string-to-integer) maps each character to a number. itos (integer-to-string) maps each number back to a character. The encode and decode functions use these mappings to convert entire strings to number lists and back.
- Save gpt.py.
You won't see output yet. The next chunk uses this tokenizer to create the training data.
- Below the tokenizer code in gpt.py, add the train/validation 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) converts the entire Shakespeare text into a list of integers using the tokenizer you just built. torch.tensor(...) wraps that list into a PyTorch tensor so the model can process it efficiently.
The first 90% of the data becomes train_data (what the model learns from). The remaining 10% becomes val_data (used to check if the model is memorizing or actually learning patterns).
- Save gpt.py.
- Verify everything works together by running:
python gpt.pyThe script should complete without any errors or output. A clean return to your terminal prompt means all the imports, downloads, tokenization, and data splitting ran successfully.
Seeing an error when you run the script?
- A NameError usually means a variable is referenced before it is defined. Check that your code blocks are in the correct order (imports, hyperparameters, dataset, tokenizer, split).
- A SyntaxError often points to a missing colon, bracket, or indentation issue. Compare your code to the full file below.
- If you see KeyError on the encode/decode lambdas, make sure you copied the tokenizer lines exactly as shown.
Help me debug the error I'm getting when running gpt.py with imports, hyperparameters, dataset download, tokenizer, and train/val split.
✔️ Awesome, I've got everything!
Great. Double check you have saved gpt.py and that it runs cleanly with no output.
ⓧ I'd like to double check the full code
import torch
import torch.nn as nn
from torch.nn import functional as F
import urllib.request
import os
# Hyperparameters
batch_size = 64
block_size = 256
max_iters = 5000
eval_interval = 500
learning_rate = 3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 384
n_head = 6
n_layer = 6
dropout = 0.2
torch.manual_seed(1337)
# Download and load the dataset
data_path = 'input.txt'
if not os.path.exists(data_path):
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
urllib.request.urlretrieve(url, data_path)
with open(data_path, 'r', encoding='utf-8') as f:
text = f.read()
# Character-level tokenizer
chars = sorted(list(set(text)))
vocab_size = len(chars)
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])
# 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:]Your data pipeline is ready. Next up, you will build the self-attention mechanism that lets your model learn which characters are most important to pay attention to.
Build the Self-Attention Mechanism
In Step 1, you set up your project environment and loaded the Tiny Shakespeare dataset into train and validation splits.
Now it's time to build the core innovation that makes transformers work: the self-attention mechanism. Self-attention is how a model decides which other tokens in a sequence are relevant to the current token. When predicting the next character after "thou art," attention lets the model look back and weigh the importance of every previous character.
In this step, get ready to:
- Implement a single causal self-attention head with query, key, and value matrices.
- Combine multiple attention heads into a multi-head attention module.
- Test the attention module with dummy data to verify correct output shapes.
Implement a single causal self-attention head
Before building the attention head itself, you need two utility functions that support data loading and loss evaluation. These will be used later during training, but their placement in the file comes before the model classes.
- In gpt.py, find the line that says val_data = data[n:] at the bottom of your file.
- Add the get_batch function below it, separated by two blank lines:
def get_batch(split):
# Select the correct data split
data_split = train_data if split == 'train' else val_data
# Pick random starting positions for our batch
ix = torch.randint(len(data_split) - block_size, (batch_size,))
# Stack sequences into input (x) and target (y) tensors
x = torch.stack([data_split[i:i+block_size] for i in ix])
y = torch.stack([data_split[i+1:i+block_size+1] for i in ix])
x, y = x.to(device), y.to(device)
return x, yWhat does this code do?
The get_batch function grabs random chunks of text from your dataset and formats them as input/target pairs.
- ix picks random starting indices. Each batch contains batch_size sequences of length block_size.
- x is the input (characters 0 through 255) and y is the target (characters 1 through 256). The model learns to predict the next character at every position.
You won't see output from this function yet. It will be used when we test the attention module in substep 3.
- Add the estimate_loss function directly below get_batch:
@torch.no_grad()
def estimate_loss():
out = {}
model.eval()
for split in ['train', 'val']:
# Average loss over eval_iters batches for stable measurement
losses = torch.zeros(eval_iters)
for k in range(eval_iters):
X, Y = get_batch(split)
logits, loss = model(X, Y)
losses[k] = loss.item()
out[split] = losses.mean()
model.train()
return outWhat does this code do?
The estimate_loss function evaluates the model by averaging loss over many batches. This gives a stable measurement instead of relying on a single noisy batch.
The @torch.no_grad() decorator disables gradient tracking during evaluation, saving memory. The function references a model variable that doesn't exist yet. Python is fine with this because the function is only defined here, not called.
Now for the main event. Self-attention works by having each token produce three vectors: a Query ("what am I looking for?"), a Key ("what do I contain?"), and a Value ("what information do I provide?").
Tokens with matching Query-Key pairs attend strongly to each other, and the output is a weighted combination of all Values.
- Add the Head class below estimate_loss:
class Head(nn.Module):
"""One head of self-attention."""
def __init__(self, head_size):
super().__init__()
# Query, Key, Value projections
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)
# Lower-triangular mask for causal (autoregressive) attention
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, scaled by sqrt(head_size)
wei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5
# Mask future positions so the model can only look backward
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 outWhat does this code do?
The Head class implements one head of causal self-attention. Here is how the forward pass works:
- Each input token gets projected into a query, key, and value vector via separate linear layers.
- Attention scores (wei) are computed by multiplying queries and keys, then scaling by 1/sqrt(head_size) to prevent the dot products from growing too large.
- The tril mask enforces causality. Positions that correspond to future tokens are filled with -inf before softmax, making their attention weights zero.
- The output is a weighted sum of value vectors. Each token's output is a blend of information from all preceding tokens, weighted by relevance.
What is causal masking?
In language generation, the model predicts the next token. It must never peek at future tokens during training. Causal masking blocks this by zeroing out attention scores for any position ahead of the current one.
The lower-triangular matrix (tril) is a simple way to achieve this. Position 3 can attend to positions 0, 1, 2, and 3, but not 4 or beyond.
- Save gpt.py.
You won't see visible output yet. The next substep adds multi-head attention, and then you will test both together.
Seeing a syntax error?
- Check that the Head class is indented consistently. Every method inside the class must be indented one level deeper than the class keyword.
- Make sure the docstring """One head of self-attention.""" uses triple double-quotes, not single-quotes.
Help me debug my Head class
Combine multiple heads into a MultiHeadAttention module
A single attention head can only focus on one type of relationship at a time. Multi-head attention runs several heads in parallel, each learning to attend to different patterns (e.g., one head might learn word adjacency while another learns longer-range grammar dependencies).
- Add the MultiHeadAttention class directly below the Head class in gpt.py:
class MultiHeadAttention(nn.Module):
"""Multiple heads of self-attention in parallel."""
def __init__(self, num_heads, head_size):
super().__init__()
# Create a list of independent attention heads
self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
# Projection layer to mix head outputs back to n_embd dimensions
self.proj = nn.Linear(n_embd, n_embd)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
# Run all heads in parallel, concatenate their outputs
out = torch.cat([h(x) for h in self.heads], dim=-1)
# Project concatenated output back to embedding dimension
out = self.dropout(self.proj(out))
return outWhat does this code do?
The MultiHeadAttention class creates num_heads independent Head instances and runs them all on the same input.
- Each head outputs a tensor of size head_size (which is n_embd // n_head = 64). Concatenating 6 heads gives back a 384-dimensional vector.
- The projection layer (self.proj) mixes information across heads. Without it, heads would be completely independent.
- Save gpt.py.
The next substep verifies that both modules produce the correct output shapes.
Test the attention module with dummy data
Before moving on, you should verify that your attention modules accept an input tensor and produce the correct output shape. This catches dimension mismatches early, before the full model is assembled.
- Add these test lines at the very bottom of gpt.py (below the MultiHeadAttention class):
# Temporary test code - remove after verifying
head_size = n_embd // n_head
x_test = torch.randn(batch_size, block_size, n_embd)
single_head = Head(head_size)
print(f"Input shape: {x_test.shape}")
print(f"Single head output shape: {single_head(x_test).shape}")
mha = MultiHeadAttention(n_head, head_size)
print(f"MultiHeadAttention output shape: {mha(x_test).shape}")- Run the script by executing this command in your terminal:
python3 gpt.pyYou should see output like this:
Input shape: torch.Size([64, 256, 384])
Single head output shape: torch.Size([64, 256, 64])
MultiHeadAttention output shape: torch.Size([64, 256, 384])What do these shapes mean?
- The input is (batch_size=64, block_size=256, n_embd=384). That's 64 sequences of 256 tokens, each represented as a 384-dimensional vector.
- A single head reduces the last dimension from 384 to 64 (the head_size). It focuses on a 64-dimensional subspace.
- Multi-head attention concatenates 6 heads of size 64 and projects back to 384. The output shape matches the input, which is essential for residual connections in the next step.
- Now remove the test lines you just added (everything below the MultiHeadAttention class) so your file is clean for the next step.
- Save gpt.py.
Seeing a shape mismatch or error?
- If you see RuntimeError mentioning matrix dimensions, check that n_embd is 384 and n_head is 6 in your hyperparameters. The head_size must be n_embd // n_head = 64.
- If you get NameError: name 'Head' is not defined, make sure the Head class appears before MultiHeadAttention in your file. Python reads top to bottom.
- If you get a memory error, your machine may not have enough RAM to handle batch_size=64 with block_size=256. Try changing batch_size to 32 temporarily for this test.
Help me fix my attention module shape error
✔️ Awesome, I've got everything!
Great. Make sure you have saved gpt.py with the test code removed.
ⓧ I'd like to double check the full code
Here is the complete gpt.py at this point:
import torch
import torch.nn as nn
from torch.nn import functional as F
import urllib.request
import os
# Hyperparameters
batch_size = 64
block_size = 256
max_iters = 5000
eval_interval = 500
learning_rate = 3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 384
n_head = 6
n_layer = 6
dropout = 0.2
torch.manual_seed(1337)
# Download and load the dataset
data_path = 'input.txt'
if not os.path.exists(data_path):
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
urllib.request.urlretrieve(url, data_path)
with open(data_path, 'r', encoding='utf-8') as f:
text = f.read()
# Character-level tokenizer
chars = sorted(list(set(text)))
vocab_size = len(chars)
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])
# 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):
data_split = train_data if split == 'train' else val_data
ix = torch.randint(len(data_split) - block_size, (batch_size,))
x = torch.stack([data_split[i:i+block_size] for i in ix])
y = torch.stack([data_split[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)
logits, loss = model(X, Y)
losses[k] = loss.item()
out[split] = losses.mean()
model.train()
return out
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)
wei = q @ k.transpose(-2, -1) * k.shape[-1]**-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)
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 outYour attention mechanism is working. Next up, you will assemble these components into a full transformer block with feedforward layers and residual connections.
Assemble the Transformer Block
Your attention mechanism can now look at all relevant tokens in a sequence. But attention alone only handles communication between tokens. Each token also needs to individually process the information it gathered.
In this step, you'll build the FeedForward network for individual token processing, compose a full TransformerBlock that combines attention with feedforward computation, and stack multiple blocks into a complete GPT model.
In this step, get ready to:
- Build the FeedForward network with a linear expansion, ReLU activation, and projection.
- Compose a full TransformerBlock with layer normalization and residual connections.
- Stack multiple blocks into the complete GPT model with embeddings and a language model head.
Build the FeedForward network
After tokens communicate through attention, each token needs to independently process what it learned. The FeedForward network expands the representation to 4x its size, applies a non-linearity, then projects back down.
- In gpt.py, find the MultiHeadAttention class and add the following FeedForward class below it:
class FeedForward(nn.Module):
"""A simple linear layer followed by a non-linearity."""
def __init__(self, n_embd):
super().__init__()
# Expand to 4x, apply ReLU, project back down
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(dropout),
)
def forward(self, x):
return self.net(x)What does this code do?
The FeedForward network gives each token a chance to "think" about the information it gathered from attention.
- nn.Linear(n_embd, 4 * n_embd) expands from 384 dimensions to 1536, giving the network more room to find patterns.
- nn.ReLU() introduces non-linearity so the network can learn complex relationships, not just linear transformations.
- nn.Linear(4 * n_embd, n_embd) projects back to the original size so it can be added to the residual stream.
Why expand to 4x the size?
The 4x expansion is a design choice from the original Transformer paper. A wider intermediate layer lets the network represent more complex functions before compressing back down. Think of it as spreading out on a big desk to do your work, then tidying everything back into a small notebook.
You won't see a change yet. The next chunk uses this FeedForward network inside a full TransformerBlock.
Compose a full TransformerBlock
A single transformer block pairs attention (communication) with a feedforward network (computation). Two critical additions make training stable: layer normalization before each sub-layer, and residual connections that add the input back to the output.
- Below the FeedForward class you just wrote, add the Block class:
class Block(nn.Module):
"""Transformer block: communication followed by computation."""
def __init__(self, n_embd, n_head):
super().__init__()
head_size = n_embd // n_head
# Communication: multi-head self-attention
self.sa = MultiHeadAttention(n_head, head_size)
# Computation: feedforward network
self.ffwd = FeedForward(n_embd)
# Layer normalization before each sub-layer
self.ln1 = nn.LayerNorm(n_embd)
self.ln2 = nn.LayerNorm(n_embd)
def forward(self, x):
# Residual connection: add input back to attention output
x = x + self.sa(self.ln1(x))
# Residual connection: add input back to feedforward output
x = x + self.ffwd(self.ln2(x))
return xWhat does this code do?
- self.ln1 and self.ln2 normalize the input before each sub-layer. This stabilizes training by keeping values in a consistent range.
- x = x + self.sa(self.ln1(x)) applies layer norm, runs attention, then adds the original input back. The addition is the residual connection.
- The same pattern repeats for the feedforward layer. Residual connections let gradients flow directly through the network during training, preventing the "vanishing gradient" problem in deep networks.
Why Pre-LN instead of Post-LN?
The original Transformer paper applied layer norm after each sub-layer (Post-LN). Modern implementations apply it before (Pre-LN) because it produces more stable training without needing learning rate warmup. GPT-2 and later models all use this Pre-LN pattern.
You won't see a change yet. The next chunk uses this Block inside the full GPT model.
Stack multiple blocks into the full GPT model
Now you have all the building blocks. The full GPT model stacks 6 transformer blocks on top of each other, preceded by token and positional embeddings, and followed by a final projection to vocabulary logits.
- Below the Block class, add the GPTLanguageModel class:
class GPTLanguageModel(nn.Module):
def __init__(self):
super().__init__()
# Each token gets an embedding vector
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
# Each position gets an embedding vector
self.position_embedding_table = nn.Embedding(block_size, n_embd)
# Stack of transformer blocks
self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])
# Final layer norm
self.ln_f = nn.LayerNorm(n_embd)
# Project from embedding space to vocabulary size
self.lm_head = nn.Linear(n_embd, vocab_size)
def forward(self, idx, targets=None):
B, T = idx.shape
# Look up token and position embeddings
tok_emb = self.token_embedding_table(idx)
pos_emb = self.position_embedding_table(torch.arange(T, device=device))
# Combine token identity with position information
x = tok_emb + pos_emb
# Pass through all transformer blocks
x = self.blocks(x)
# Final layer norm
x = self.ln_f(x)
# Project to vocabulary logits
logits = self.lm_head(x)
if targets is None:
loss = None
else:
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):
# Crop context to block_size tokens
idx_cond = idx[:, -block_size:]
# Get predictions
logits, loss = self(idx_cond)
# Focus on the last time step
logits = logits[:, -1, :]
# Convert to probabilities
probs = F.softmax(logits, dim=-1)
# Sample the next token
idx_next = torch.multinomial(probs, num_samples=1)
# Append to the running sequence
idx = torch.cat((idx, idx_next), dim=1)
return idxWhat does this code do?
This is the full model that ties everything together. It has three key parts:
- Embeddings: token_embedding_table converts each character into a 384-dimensional vector. position_embedding_table adds position information so the model knows where each token sits in the sequence.
- Transformer blocks: 6 stacked blocks of attention plus feedforward. Each block refines the representation, building deeper understanding of the text.
- Language model head: lm_head projects from 384 dimensions back to 65 (one score per character in the vocabulary). The highest score is the model's prediction for the next character.
The generate method works autoregressively: predict one token, append it to the sequence, then predict the next. This is exactly how ChatGPT and Claude produce text one token at a time.
- Save gpt.py.
- Verify there are no syntax errors by running:
python gpt.py- The script should complete without any traceback or error messages. It downloads the dataset (if not already present) and defines all classes, then exits cleanly.
Seeing a NameError or IndentationError?
- Make sure GPTLanguageModel is defined below the Block class. Python reads top to bottom, so each class must appear after the classes it references.
- Check that all method definitions inside the class are indented with 4 spaces relative to the class keyword.
- Verify that the forward and generate methods both use self as their first parameter.
Help me debug my GPTLanguageModel class definition.
✔️ Awesome, I've got everything!
Great. Double check you've saved your file, then continue to the next step.
ⓧ I'd like to double check the full code
Here's the complete gpt.py up to this point:
import torch
import torch.nn as nn
from torch.nn import functional as F
import urllib.request
import os
# Hyperparameters
batch_size = 64
block_size = 256
max_iters = 5000
eval_interval = 500
learning_rate = 3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 384
n_head = 6
n_layer = 6
dropout = 0.2
torch.manual_seed(1337)
# Download and load the dataset
data_path = 'input.txt'
if not os.path.exists(data_path):
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
urllib.request.urlretrieve(url, data_path)
with open(data_path, 'r', encoding='utf-8') as f:
text = f.read()
# Character-level tokenizer
chars = sorted(list(set(text)))
vocab_size = len(chars)
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])
# 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):
data_split = train_data if split == 'train' else val_data
ix = torch.randint(len(data_split) - block_size, (batch_size,))
x = torch.stack([data_split[i:i+block_size] for i in ix])
y = torch.stack([data_split[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)
logits, loss = model(X, Y)
losses[k] = loss.item()
out[split] = losses.mean()
model.train()
return out
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)
wei = q @ k.transpose(-2, -1) * k.shape[-1]**-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)
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):
"""A simple linear layer followed by a non-linearity."""
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(dropout),
)
def forward(self, x):
return self.net(x)
class Block(nn.Module):
"""Transformer block: communication followed by computation."""
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 GPTLanguageModel(nn.Module):
def __init__(self):
super().__init__()
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
self.position_embedding_table = nn.Embedding(block_size, n_embd)
self.blocks = nn.Sequential(*[Block(n_embd, n_head=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_table(idx)
pos_emb = self.position_embedding_table(torch.arange(T, device=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 = 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, loss = 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 idxYour full GPT architecture is complete. Next up, you'll write the training loop that teaches this model to predict the next character in Shakespeare's text.
Train Your GPT
Your transformer architecture is fully assembled. You have attention heads, feedforward networks, residual connections, and layer normalization all wired together into a complete GPTLanguageModel class.
Now it's time to bring your model to life. A neural network's weights start as random numbers, so your model currently produces pure noise. Training is the process of adjusting those weights so the model learns to predict the next character in Shakespeare's text.
In this step, get ready to:
- Instantiate the model and understand how the data loader feeds training batches.
- Write the training loop with the AdamW optimizer and cross-entropy loss.
- Train the model and monitor the loss decrease over 5000 iterations.
Instantiate the model and connect the data loader
Before training can begin, you need to create an instance of your GPTLanguageModel and move it to the right device (GPU if available, otherwise CPU).
You already wrote a get_batch function earlier that handles data loading. It randomly samples chunks of text from the dataset and returns input/target pairs. Each batch contains 64 sequences of 256 characters, giving the model diverse training examples every iteration.
- In gpt.py, scroll to the very bottom of the file (after the GPTLanguageModel class definition).
- Add the model instantiation code below the class:
model = GPTLanguageModel()
m = model.to(device)
print(f"{sum(p.numel() for p in m.parameters())/1e6:.2f}M parameters")What does this code do?
GPTLanguageModel() creates an instance of your full transformer model with all the layers and embeddings you defined.
model.to(device) moves all model parameters to the GPU (if available) or keeps them on CPU. The print statement counts every learnable parameter in the model. You should see approximately 10.8 million parameters.
How does get_batch feed data to the model?
The get_batch function you wrote earlier randomly picks 64 starting positions in the text. For each position, it grabs 256 consecutive characters as input (x) and the same 256 characters shifted by one as the target (y).
This random sampling means the model sees different parts of Shakespeare's text every iteration, which prevents overfitting to any one passage.
- Save gpt.py.
Write the training loop with AdamW optimizer
The training loop is where the model actually learns. Each iteration, the model makes a prediction, measures how wrong it was (using cross-entropy loss), and adjusts its weights to be less wrong next time.
- In gpt.py, add the following code directly below the print statement you just added:
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
for iter in range(max_iters):
# Evaluate loss on train and val sets at regular intervals
if iter % eval_interval == 0 or iter == max_iters - 1:
losses = estimate_loss()
print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
# Sample a batch of training data
xb, yb = get_batch('train')
# Forward pass, compute loss, backpropagate, update weights
logits, loss = model(xb, yb)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()What does this code do?
- AdamW is the optimizer that controls how aggressively the model updates its weights. It uses a learning rate of 3e-4, which is a standard choice for transformer training.
- Every 500 iterations, estimate_loss() averages the loss over 200 batches to give a stable measurement of how well the model is learning.
- The core loop grabs a batch, runs it through the model, computes how wrong the predictions were (loss), then uses loss.backward() to calculate gradients and optimizer.step() to update the weights.
- optimizer.zero_grad(set_to_none=True) clears old gradients before each new batch. Without this, gradients would accumulate and the model would learn incorrectly.
What is AdamW?
AdamW is an improved version of the Adam optimizer that properly handles weight decay (a regularization technique). It is the standard optimizer for training transformers, used in GPT-2, GPT-3, and most modern language models.
- Save gpt.py.
Seeing a NameError for model or optimizer?
Make sure the model instantiation code (model = GPTLanguageModel()) appears BEFORE the optimizer line. The optimizer needs the model to already exist so it can reference its parameters.
Help me debug my training loop setup.
Run training and monitor loss decrease
Your model starts with random weights, so its initial loss should be around 4.17. This is what you get from randomly guessing across 65 possible characters (ln(65) = 4.17). As training progresses, you should see this loss steadily decrease.
- Run your model training by executing this command in your terminal:
python gpt.pyWhat should I see?
First, you should see a line reporting approximately 10.79M parameters. Then every 500 steps, you will see train and validation loss printed.
On a GPU, the loss should drop from ~4.17 to approximately 1.47 on the validation set over 5000 iterations. On CPU, training will take significantly longer (30+ minutes with the full configuration).
Training on CPU is slow
If you are training on CPU and want faster results, you can temporarily reduce the hyperparameters at the top of gpt.py: try n_layer = 4, n_embd = 128, block_size = 64, batch_size = 12, and max_iters = 2000. This trains in about 3 minutes and reaches a validation loss of ~1.88.
- Watch the terminal output as training progresses. You should see loss values printed at step 0, 500, 1000, 1500, 2000, 2500, 3000, 3500, 4000, 4500, and 4999.
- Confirm that both train loss and val loss are decreasing over time. The first line should show a loss near 4.17. The final line should show a validation loss near 1.47 (GPU) or 1.88 (CPU with reduced config).
Getting CUDA out of memory?
If your GPU runs out of VRAM, reduce batch_size to 32 or lower n_embd to 128 at the top of your file. These reduce memory usage while still allowing the model to learn.
Help me fix my GPU memory issue during training.
✔️ Awesome, I've got everything!
Great. Double check you have saved gpt.py and your training has finished running.
ⓧ I'd like to double check the full code
Here is the full gpt.py file as it should look at this point:
import torch
import torch.nn as nn
from torch.nn import functional as F
import urllib.request
import os
# Hyperparameters
batch_size = 64
block_size = 256
max_iters = 5000
eval_interval = 500
learning_rate = 3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 384
n_head = 6
n_layer = 6
dropout = 0.2
torch.manual_seed(1337)
# Download and load the dataset
data_path = 'input.txt'
if not os.path.exists(data_path):
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
urllib.request.urlretrieve(url, data_path)
with open(data_path, 'r', encoding='utf-8') as f:
text = f.read()
# Character-level tokenizer
chars = sorted(list(set(text)))
vocab_size = len(chars)
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])
# 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):
data_split = train_data if split == 'train' else val_data
ix = torch.randint(len(data_split) - block_size, (batch_size,))
x = torch.stack([data_split[i:i+block_size] for i in ix])
y = torch.stack([data_split[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)
logits, loss = model(X, Y)
losses[k] = loss.item()
out[split] = losses.mean()
model.train()
return out
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)
wei = q @ k.transpose(-2, -1) * k.shape[-1]**-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)
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):
"""A simple linear layer followed by a non-linearity."""
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(dropout),
)
def forward(self, x):
return self.net(x)
class Block(nn.Module):
"""Transformer block: communication followed by computation."""
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 GPTLanguageModel(nn.Module):
def __init__(self):
super().__init__()
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
self.position_embedding_table = nn.Embedding(block_size, n_embd)
self.blocks = nn.Sequential(*[Block(n_embd, n_head=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_table(idx)
pos_emb = self.position_embedding_table(torch.arange(T, device=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 = 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, loss = 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
model = GPTLanguageModel()
m = model.to(device)
print(f"{sum(p.numel() for p in m.parameters())/1e6:.2f}M parameters")
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
for iter in range(max_iters):
if iter % eval_interval == 0 or iter == max_iters - 1:
losses = estimate_loss()
print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
xb, yb = get_batch('train')
logits, loss = model(xb, yb)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()Your model has learned to predict Shakespeare's text with significantly better accuracy than random guessing. Next up, you will use this trained model to generate brand new text and see Shakespeare-style prose emerge from your creation.
Generate Text and Analyze Results
Your model has finished training and the loss has dropped from ~4.17 (random guessing) down to around 1.47 on validation.
Now comes the payoff: making your GPT generate brand new text. The model has learned statistical patterns in Shakespeare's writing, and you're about to see it produce novel sequences character by character.
In this step, get ready to:
- Understand how the autoregressive generation function predicts one token at a time.
- Generate Shakespeare-style text and try different prompts.
- Experiment with model hyperparameters and observe effects on output quality.
Understand the autoregressive generation function
Your GPTLanguageModel class already contains a generate method that you wrote in Step 3. Before calling it, let's walk through how autoregressive generation works.
- In gpt.py, scroll to the generate method inside the GPTLanguageModel class. Find this code:
def generate(self, idx, max_new_tokens):
for _ in range(max_new_tokens):
idx_cond = idx[:, -block_size:]
logits, loss = 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 idxHow does autoregressive generation work?
The model generates text one character at a time in a loop. Each iteration follows the same pattern: predict, sample, append.
- idx_cond = idx[:, -block_size:] crops the context to the last block_size tokens (256 characters). The model can only look at this many tokens at once.
- logits[:, -1, :] extracts the prediction for the very last position. This is the model's guess for what comes next.
- torch.multinomial(probs, num_samples=1) samples from the probability distribution rather than always picking the highest probability. This introduces creative variety in the output.
- torch.cat appends the new token to the sequence, then the loop repeats with the extended context.
Now you'll call this method to produce text from your trained model.
- Scroll to the very bottom of gpt.py (after the training loop's optimizer.step() line).
- Add the generation call by copying this code at the end of the file:
# Generate text
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))What does this code do?
- torch.zeros((1, 1), dtype=torch.long, device=device) creates a starting token of index 0, which corresponds to the newline character in our vocabulary. This gives the model a blank slate to begin generating.
- m.generate(context, max_new_tokens=500) asks the model to produce 500 new characters one at a time.
- decode(...[0].tolist()) converts the integer token sequence back into readable characters using your tokenizer's decode function.
- Save gpt.py.
Generate text with different prompts
With the generation code in place, it's time to run the full script and see your model produce Shakespeare-style text.
- Run the script by executing this command in your terminal:
python gpt.pyWhat should I see?
The script will train the model first (printing loss values every 500 steps), then print 500 characters of generated text at the end. Training takes around 3-8 minutes on GPU or 30+ minutes on CPU with the full configuration.
The generated text will look like Shakespeare. You'll see character names, iambic-style dialogue, and stage directions. It won't be perfect English, but it captures the rhythm and vocabulary of the training data.
Seeing gibberish or random characters?
- Check your final validation loss. If it's still above 3.0, the model hasn't trained enough. Make sure max_iters is set to 5000.
- If you're running on CPU and reduced the hyperparameters for speed, the output quality will be lower. A validation loss around 1.88 is expected for reduced configurations.
- Make sure the generation code is placed after the training loop completes, not inside it.
My generated text doesn't look like Shakespeare at all. Help me debug.
Now try providing a specific starting prompt instead of a blank slate. The model will continue from whatever text you give it.
- In gpt.py, find this line at the bottom of the file:
context = torch.zeros((1, 1), dtype=torch.long, device=device)- Try replacing it with an encoded prompt, like this:
# Try a custom prompt
context = torch.tensor([encode("ROMEO:")], dtype=torch.long, device=device)
print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))- Save and run the script again. The model will complete the text from "ROMEO:" as if it's continuing a Shakespeare scene.
- Experiment with other prompts like "JULIET:", "To be or", or "The king" to see how the model responds to different starting contexts.
Why does the prompt affect the output?
The generate method uses the prompt as the initial context window. The model's attention mechanism looks at all provided tokens when predicting the next one. A prompt like "ROMEO:" biases the model toward dialogue patterns it saw following that character name during training.
- Once you've experimented, change the context line back to the original version before moving on:
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))- Save gpt.py.
Experiment with model hyperparameters
The quality of your model's output is directly tied to its size and architecture. Changing the hyperparameters at the top of the file lets you see how model capacity affects text generation.
- In gpt.py, find the hyperparameters section at the top of the file:
n_embd = 384
n_head = 6
n_layer = 6- Try reducing the model size by changing these values to n_embd = 128, n_head = 4, and n_layer = 4.
- Run the script again and compare the generated text quality to the original output.
What do these hyperparameters control?
- n_embd controls the embedding dimension. This is the width of information each token carries through the network. Larger values capture richer representations.
- n_head controls how many parallel attention heads operate at each layer. More heads let the model attend to different types of relationships simultaneously.
- n_layer controls how many transformer blocks stack on top of each other. More layers give the model deeper reasoning about token relationships.
With the reduced config, your model drops from ~10.8M parameters to roughly ~1M. Training is much faster, but the generated text will be noticeably less coherent.
- After experimenting, set n_embd back to 384, n_head back to 6, and n_layer back to 6.
- Save gpt.py.
The scaling law in action
What you just observed is a fundamental principle in deep learning: bigger models with more parameters produce better results, but require more compute. This same tradeoff drives the scaling decisions behind production LLMs like GPT-4 and Claude.
✔️ Awesome, I've got everything!
Great. Double check you've saved gpt.py with the original hyperparameters restored and the generation code at the bottom.
ⓧ I'd like to double check the full code
Here's the complete gpt.py file for comparison:
import torch
import torch.nn as nn
from torch.nn import functional as F
import urllib.request
import os
# Hyperparameters
batch_size = 64
block_size = 256
max_iters = 5000
eval_interval = 500
learning_rate = 3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 384
n_head = 6
n_layer = 6
dropout = 0.2
torch.manual_seed(1337)
# Download and load the dataset
data_path = 'input.txt'
if not os.path.exists(data_path):
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
urllib.request.urlretrieve(url, data_path)
with open(data_path, 'r', encoding='utf-8') as f:
text = f.read()
# Character-level tokenizer
chars = sorted(list(set(text)))
vocab_size = len(chars)
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])
# 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):
data_split = train_data if split == 'train' else val_data
ix = torch.randint(len(data_split) - block_size, (batch_size,))
x = torch.stack([data_split[i:i+block_size] for i in ix])
y = torch.stack([data_split[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)
logits, loss = model(X, Y)
losses[k] = loss.item()
out[split] = losses.mean()
model.train()
return out
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)
wei = q @ k.transpose(-2, -1) * k.shape[-1]**-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)
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):
"""A simple linear layer followed by a non-linearity."""
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(dropout),
)
def forward(self, x):
return self.net(x)
class Block(nn.Module):
"""Transformer block: communication followed by computation."""
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 GPTLanguageModel(nn.Module):
def __init__(self):
super().__init__()
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
self.position_embedding_table = nn.Embedding(block_size, n_embd)
self.blocks = nn.Sequential(*[Block(n_embd, n_head=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_table(idx)
pos_emb = self.position_embedding_table(torch.arange(T, device=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 = 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, loss = 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
model = GPTLanguageModel()
m = model.to(device)
print(f"{sum(p.numel() for p in m.parameters())/1e6:.2f}M parameters")
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
for iter in range(max_iters):
if iter % eval_interval == 0 or iter == max_iters - 1:
losses = estimate_loss()
print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
xb, yb = get_batch('train')
logits, loss = model(xb, yb)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
# Generate text
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))Secret mission
Modernize Your Transformer
Your GPT works, but it uses techniques from the original 2017 transformer paper. Production LLMs like LLaMA and GPT-3 use improved components that train faster and produce better output. Can you upgrade your model with RMSNorm and GELU to see the difference for yourself?
Clean Up Your Resources
Clean Up Your Resources
This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your files for future experimentation or delete them entirely.
Resources you used:
- Project folder on your Desktop containing gpt.py and input.txt.
- Python virtual environment inside the project folder.
Keep everything running
No action needed. Choose this if you want to keep experimenting with hyperparameters, training for more iterations, or extending your model further.
- Your gpt.py file and input.txt dataset stay on your Desktop.
- Your virtual environment stays ready so you can rerun the script anytime.
Pause - I'll come back to this later
Since this is a local project with no running processes, there is nothing to pause. Your files sit idle on disk until you return.
- Deactivate your virtual environment by running:
deactivate- When you come back, reactivate it with source venv/bin/activate (macOS/Linux) or venv\Scripts\activate (Windows).
Delete - I don't want to use this again
Remove all project files and start fresh. This deletes your trained model code and the downloaded dataset.
- Deactivate your virtual environment first (if active) by running:
deactivate- Navigate to your Desktop and remove the project folder by running these commands.
On macOS/Linux:
cd ~/Desktop
rm -rf gpt-from-scratchOn Windows:
cd ~\Desktop
rmdir /s /q gpt-from-scratch- This removes gpt.py, input.txt, and the virtual environment in one step.
Nice Work!
Nice Work!
You just built a GPT language model entirely from scratch. Your model understands the sequential structure of Shakespeare's writing well enough to generate new text in the same style.
Here's what you learned:
- Implemented the complete transformer architecture piece by piece, including self-attention, multi-head attention, feedforward networks, and residual connections.
- Built a full training loop with the AdamW optimizer and cross-entropy loss to teach your model to predict the next character.
- Implemented autoregressive text generation with causal masking to produce novel Shakespeare-style text from your trained model.
- Modernized your architecture with RMSNorm and GELU activation, understanding why these improvements are used in production LLMs.
Ready to quiz yourself?