aggrogpt

A typed, object-structured re-implementation of microgpt, targeting Shakespeare instead of names. PyTorch-backed math; same architecture as microgpt; clearer code.

This project keeps the algorithm intact but introduces enough abstraction that:

• Every shape is encoded in a class.
• Every public attribute and method is fully type-annotated.
• The Pyright type checker passes in strict mode.
• The math sits behind a thin tensor surface (Vector, Matrix) that delegates to torch.Tensor underneath. Higher layers (tokenizer, layers, model, train) never import torch directly.

File layout

aggrogpt/
  tensors.py     Vector, Matrix and the ops they support. Backed by
                 torch.Tensor; the only file that imports torch
                 internally (plus train.py for the optimizer).
  tokenizer.py   TokenId, PositionId, Tokenizer (character-level).
  corpus.py      Downloads Project Gutenberg's Complete Works of
                 Shakespeare (pg100.txt) and strips the license front
                 and back matter.
  layers.py      TokenEmbedding, PositionEmbedding, LinearProjection,
                 RMSNorm, MultiHeadAttention, MLP, TransformerLayer,
                 KVCacheLayer.
  model.py       KVCache, GPT.
  train.py       Entry point: load corpus, build model, train with
                 torch.optim.Adam, sample.

Running

uv sync --dev
uv run python train.py

The first run downloads pg100.txt (~5.5 MB) from Project Gutenberg and writes a cleaned shakespeare.txt next to it. Subsequent runs read the cleaned cache. Training the defaults (NUM_TRAINING_STEPS = 1000) takes about 3 minutes on a laptop CPU.

What "cleaned" means

corpus.py strips the two canonical Project Gutenberg marker lines (*** START OF THE PROJECT GUTENBERG EBOOK ... and the matching *** END OF ...) and discards everything outside them — about 12 KB of licensing text the model has no business learning.

The Project-Gutenberg-supplied table of contents (a short list of play titles) is left intact at the top of the file. It is small enough that the model is unlikely to overfit to it, and detecting the boundary between TOC and the first play reliably is fiddly.

Hyperparameters

Defaults in train.py:

name	value	meaning
EMBEDDING_DIM	32	per-token vector width
N_HEAD	4	attention heads per layer
N_LAYER	2	stacked transformer layers
BLOCK_SIZE	32	context window length
LEARNING_RATE	0.01	initial Adam step size (linearly decayed)
NUM_TRAINING_STEPS	1000	training steps
TEMPERATURE	0.5	sampling temperature

About 32,000 parameters. Expect outputs that look like Shakespeare in shape — character names followed by colons, blank lines between speakers, roughly the right vocabulary — but make no real semantic sense. For better quality, bump EMBEDDING_DIM, N_LAYER, and NUM_TRAINING_STEPS.

Type checking and formatting

uv run pyright .   # strict mode, with three torch-specific rules relaxed
uv run ruff format .
uv run ruff check .

Configured in pyproject.toml. Pyright runs in strict mode but with reportUnknownMemberType, reportUnknownArgumentType, and reportUnknownVariableType disabled — these are necessary because torch's type stubs leave many parameter and return types as Unknown.

VS Code

.vscode/launch.json is set up to run train.py under debugpy. Recommended extensions install on open.

Why per-token forward (and not batched)?

In real PyTorch GPT implementations, training runs forward on an entire (block_size,) sequence in parallel and applies a causal mask. This project keeps the microgpt-style per-token forward — one call per position, growing the KV cache as it goes — for two reasons:

1. Readability. The GPT.__call__ signature is (token, position, cache) -> logits. The body is short and reads linearly. A batched version threads a full sequence and mask through every sub-block; harder to follow.
2. Architectural symmetry between training and inference. The same __call__ runs at both training and sampling time, with the same cache growing behaviour. No separate code paths to reconcile.

The cost is a constant factor: per-token training is roughly 30× slower on CPU than batched training would be. With the current backend that still works out to a few minutes for 1000 steps, which is fast enough for an evening of tinkering.

Swapping or extending the backend

To run on GPU: set the parameter tensors to .cuda() after creation and ensure all input tensors live on the same device. The Matrix.parameter_normal factory would need a device= argument.

To switch to batched training: change GPT.__call__ to take a 1-D LongTensor of token ids of length block_size and emit a 2-D logits tensor of shape (block_size, vocab_size). The per-token KV-cache machinery is retained for inference only.

Credits

Algorithm: Andrej Karpathy's microgpt and nanoGPT. The transformer architecture is the standard pre-norm, RMSNorm, multi-head-attention GPT-2 recipe with no biases and learned positional embeddings.

Corpus: Project Gutenberg eBook #100, The Complete Works of William Shakespeare, in the public domain.