Build an LLM.

A token is a unit of text. Here, one normalized Unicode character = one token. The book later uses BPE, where a token can contain multiple characters. Vocabulary is built from the whole corpus; only the first 80% is used for updates.

chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
ids = [stoi[c] for c in text]
assert ''.join(chars[i] for i in ids) == text

Context is the maximum number of tokens read at once. Targets are the same sequence shifted by one token. Training and validation are split before making windows: no window crosses the boundary. Validation measures prediction on held-out text, not understanding.

train, validation = ids[:split], ids[split:]
x = train[start:start + context]
y = train[start + 1:start + context + 1]
assert x[1:] == y[:-1]

Queries Q and keys K produce similarity scores. Future scores are masked before softmax, which converts scores into weights summing to 1. Values V are then combined with these weights. This example uses fixed Q/K to isolate the mask; the model learns its Q/K/V.

scores = q @ k.transpose(-2, -1) / sqrt(head_dim)
scores = scores.masked_fill(future, float('-inf'))
weights = scores.softmax(dim=-1)
context = weights @ v

Embeddings map each token and its position to vectors. Two attention heads read different relationships. An MLP transforms the result; residual connections and LayerNorm stabilize computation. Logits are scores, one per possible token. We use one block and no dropout to keep the model small and reproducible.

x = token_embedding(ids) + position_embedding(positions)
x = x + attention(norm1(x))
x = x + mlp(norm2(x))
logits = output(final_norm(x))

Loss is the mean prediction error (cross entropy, in nats). Lower loss means better prediction of these targets. Gradients indicate how to change weights; Adam updates them at rate 0.003 with gradient norm clipped to 1. Measurements always use the same 8 windows per split. Falling training loss with rising validation loss indicates overfitting. The best checkpoint has the lowest measured validation loss.

loss = cross_entropy(logits.flatten(0, 1), targets.flatten())
optimizer.zero_grad()
loss.backward()
clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()

Temperature divides logits before softmax. Lower values concentrate the distribution; higher values spread it. Sampling uses the full vocabulary, without top-k. Displayed probabilities describe the last generation decision. The same seed and weights produce the same sequence. A few lines of text produce a demonstration model, not an assistant or reasoning capability.

logits = model(ids[:, -context:])[:, -1, :]
probabilities = (logits / temperature).softmax(-1)
next_id = multinomial(probabilities, 1)
ids = cat((ids, next_id), dim=1)