Try it live · PyTorch · Streamlit
A decoder-only transformer language model, written from first principles —
every piece that actually does the "transformer" work (causal self-attention,
the multi-head split/merge, the feed-forward block) is built from raw tensor
ops, not assembled from torch.nn.MultiheadAttention or
torch.nn.TransformerDecoder.
Trained at the character level on the complete works of Shakespeare (public domain, via Project Gutenberg) to demonstrate the whole pipeline end-to-end on CPU: tokenizer → model → training loop → sampling.
Most of my other projects use LLMs through an API (Claude, Groq) and focus on the systems around them — retrieval, agent handoffs, evaluation. This one is the opposite: no API, no pretrained weights, nothing hidden behind a library call. The point is to show the actual mechanism — attention scores, causal masking, backprop through the whole stack — works correctly, not just that it's shaped right.
tokens
│
▼
token embedding + positional embedding
│
▼
┌─────────────────────────────┐
│ Transformer Block × 4 │
│ ┌────────────────────────┐ │
│ │ LayerNorm │ │
│ │ Causal Self-Attention │──┐ residual
│ │ LayerNorm │ │ add
│ │ Feed-Forward (MLP) │──┘
│ └────────────────────────┘ │
└─────────────────────────────┘
│
▼
LayerNorm → Linear → logits over vocab
Causal self-attention (model.py) is implemented manually:
project to Q/K/V, split into heads, compute scaled dot-product scores,
mask out any position j > i with -inf before the softmax so token i
can never attend to a token that comes after it, then merge the heads back.
Shape tests catch a wrong dimension; they don't catch an attention mask that's silently wrong. tests/test_model.py includes a test that changes only the last token of the input and asserts every earlier position's output logits are byte-for-byte unchanged — the direct, falsifiable version of "no future leakage," not an assumption.
There's also a training-loop sanity check (test_overfits_a_tiny_fixed_batch):
a small model trained on one fixed batch for 200 steps should crush the loss
by at least 10x. If the forward/backward wiring were broken, this would fail
regardless of how well shape tests pass.
pip install -r requirements.txt
pytest tests/ -v833,280 parameters, 4 layers, 4 heads, 128-dim embeddings, block size 128. 2000 steps, batch size 64, trained on ~1.35M characters (90/10 train/val split) — CPU only, ~26 minutes.
| step | train loss | val loss |
|---|---|---|
| 0 | 4.561 | 4.564 |
| 200 | 2.535 | 2.562 |
| 600 | 2.188 | 2.246 |
| 1000 | 1.930 | 2.010 |
| 1400 | 1.761 | 1.883 |
| 2000 | 1.612 | 1.770 |
Loss at step 0 (4.56) is almost exactly ln(92) ≈ 4.52 — the vocab size —
which is what a model predicting uniformly at random should score. It drops
steadily to 1.61 with no train/val divergence, i.e. it's still generalizing,
not memorizing, at 2000 steps.
Sample generation (temperature 0.8, top-k 40), starting from an empty prompt:
DOL.
Why one his too.
POSTHUMUS.
And hast sir, what there man, with be goints and
Of me, shall his countly, and for thy confortuor hand full he incomplace gracks that bost reath the runjuse rom on my live of
Follant all of the greated plise pyments,
Eand his like the grand high wayth they falless o
Not coherent English — an 833K-parameter character-level model over 26
minutes of CPU training isn't going to produce that. What it did learn:
speaker names in caps followed by a period and dialogue (DOL., POSTHUMUS.),
line-break rhythm that looks like verse, and archaic-flavored invented words
with plausible English letter statistics. That's the model correctly
learning the structure of the corpus from raw characters, which is what
this project is actually testing.
pip install -r requirements.txt
python data/prepare.py # downloads Shakespeare's complete works, trims to ~1.5M chars
python train.py # ~25-30 min on CPU, saves checkpoint.pt
python generate.py "ROMEO:"Why character-level, not BPE. A subword tokenizer would need its own training pass and would make it harder to see what the transformer itself learned vs. what the tokenizer already encoded. Character-level keeps the vocabulary tiny (92 tokens) and puts all the modeling burden on the transformer.
Why Shakespeare. Public domain, large enough to need a real model, small enough (trimmed to ~1.5M characters) to train on CPU in well under half an hour, and stylistically distinctive enough that "did it learn the structure" is an easy visual check.
Why 2000 steps and not more. This project is about proving the mechanism works, not chasing a state-of-the-art loss number. The loss curve is still falling at step 2000 with no sign of overfitting — more steps would likely help, but the honest result at a CPU-friendly budget is more useful here than a cherry-picked longer run.
- Character-level, not word/subword — no chance of learning real semantics at this scale, only surface structure and letter statistics.
- No positional extrapolation: generation is capped at
block_size(128) tokens of context. - CPU-only training budget (2000 steps). A GPU and a longer run would produce noticeably more coherent output; that's a deliberate scope boundary, not a bug.
MIT