Build Your Own Mini GPT — A Language Model in 300 Lines

Write a character-level GPT in PyTorch from an empty file: tokenizer, causally masked self-attention, training loop, and temperature sampling — then watch Shakespeare's formatting emerge from nothing but next-character prediction.

Attention Is All You Need

What fits in 300 lines

"Building a GPT" sounds like a story about thousands of GPUs and trillions of tokens. For training a commercial model, it is. But the architecture itself is startlingly small. A tokenizer that turns characters into numbers, embeddings, self-attention, a feed-forward block, a training loop, a generation loop — all of it fits in roughly 300 lines of PyTorch.

What we build here is a character-level GPT: it predicts one character at a time instead of one word. The vocabulary is only a few dozen symbols, so the output layer stays small and the tokenizer takes twenty lines. That leaves all the attention on the Transformer itself. For data we use the complete works of Shakespeare flattened into a single text file — the familiar tiny-shakespeare corpus, on the order of a million characters with about 65 distinct symbols in it.

The goal of this article is not to copy the code. It is to be able to say what lives on each of those 300 lines, and what breaks if you delete any one of them. That is the exact boundary between calling model.generate() from a library and choosing your own settings with reasons behind them.

A GPT is a probability distribution over the next character

Strip everything away and a GPT is one function: it takes the text so far and returns a probability distribution over the next single character. That's all.

Show it "To be or not to b" and it should give "e" a high probability and "z" a low one — a number for every symbol in the vocabulary. Generating text is nothing more than calling that function, sampling one character, appending it, and calling again. The model is not "writing a sentence"; it is being re-invoked once per roll of the dice.

Training means adjusting the weights so that the character that actually came next gets as much probability as possible. The whole objective is one cross-entropy term.

In words: at every position $t$, take the probability the model assigned to the character that actually appeared, take its logarithm, flip the sign, and average. Here $x_{<t}$ means "every character before position $t$", $\theta$ is the full set of model parameters, and $T$ is how many characters we look at in one go. Assign probability 1 to the truth and $\log 1 = 0$ costs nothing; the more the model underrates the truth, the faster the number climbs.

One efficiency note worth internalising. Pushing a length-$T$ passage through the model once gives you predictions and losses at all $T$ positions simultaneously. You are not training one character at a time. That is why next-character prediction can consume enormous amounts of text — and the causal mask we meet later is precisely the trick that makes this simultaneous computation legitimate.

The raw numbers the model emits are called logits, and softmax turns them into probabilities. At generation time we slip a temperature $\tau$ into that step:

$$p_i = \frac{\exp(z_i/\tau)}{\sum_j \exp(z_j/\tau)}$$

In words, this just divides each logit $z_i$ by $\tau$ before the softmax. A small $\tau$ exaggerates the largest logits and makes the distribution spiky; a large $\tau$ flattens it. That single division is the dial between "safe but dull" and "inventive but incoherent". Try it:

The parts list — where the 300 lines go

The rest is imports and configuration. No new concepts appear from here on. What follows is simply filling in that table from the top.

Turning characters into integers

Neural networks only handle numbers, so the first job is assigning an integer to each character. Collect the unique characters from the corpus, sort them, and build two dictionaries.

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

text  = open('input.txt', encoding='utf-8').read()
chars = sorted(set(text))                      # about 65 symbols
stoi  = {c: i for i, c in enumerate(chars)}
itos  = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: ''.join(itos[i] for i in ids)

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

A real GPT swaps this for subword splitting such as BPE, pushing the vocabulary into the tens of thousands. The difference is purely a trade between vocabulary size and sequence length; everything downstream is identical. Character-level keeps the vocabulary tiny, but the same passage becomes a much longer sequence, so the model has to reach further back. How to choose a splitting scheme is covered in Tokenizers from Scratch.

The input is two copies offset by one

This is the part of the implementation people most often misread. The input x and the target y are the same string shifted by a single character.

def get_batch(data, block_size, batch_size):
    ix = torch.randint(len(data) - block_size - 1, (batch_size,))
    x = torch.stack([data[i:i + block_size]         for i in ix])
    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
    return x, y