Transformer, sample-trained on LOTR (Decoder-only)

This repository contains a from‑scratch PyTorch implementation of decoder‑only transformers (inspired by ‘Attention Is All You Need’). It uses deterministic unit tests and reproducible training (Docker + CI).

The model is configured by default to train on a Lord of the Rings corpus with spaCy tokenization.

Features

• Decoder-only Transformer (multi-head self-attention with causal masking)
• spaCy tokenization with a simple fallback tokenizer
• Custom Vocab with <UNK> handling and minimum-freq filtering
• Dataset preparation utilities for token sequences
• Training loop with label smoothing, evaluation sampler, simple CLI (scripts.run)
• CPU and GPU Dockerfiles for reproducible environments
• Unit tests (attention parity, mask correctness, dataset/target offset), deterministic sampling tests, and CI integration.

Quickstart (local)

Create and activate a virtual env, install dependencies, and download spaCy model:

python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
python -m spacy download en_core_web_sm

Download example LOTR data and start training:

# download data
python -m scripts.download_data --out_dir data/raw

# train (adjust config as needed)
python -m scripts.run train --config configs/default.yaml

Quick smoke run

To do a fast sanity/smoke run: create a small file under data/raw/text_files/ and edit configs/default.yaml to use a small batch_size and n_steps, then run the same scripts.run train command above.

Practical notes and troubleshooting

• spaCy model: The tokenization requires the en_core_web_sm model to be downloaded locally. If you run in Docker the image build already downloads the model so you don't need to install it on the host.

• venv activation on Windows: activation scripts are in .venv\Scripts\ (not .venv/bin/). From Git Bash use:

source .venv/Scripts/activate

• Use Docker to avoid installing heavy packages on the host (recommended), e.g., build the CPU image and run tests inside it.

# build CPU image (installs deps inside the image)
docker build -f docker/Dockerfile.cpu -t lotr-transformer:cpu .

# run the unit tests inside the built image
docker run --rm -it -v "$(pwd)":/app -w /app lotr-transformer:cpu bash -lc "pytest -q"

If you have NVIDIA Container Toolkit and want GPU training, the provided docker-compose.yml includes a train target you can run with:

docker compose -f docker-compose.yml up --build train

• Data placement: the training script expects the corpus at data/raw/text_files/lotr.txt by default (see configs/default.yaml). Use python -m scripts.download_data --out_dir data/raw to fetch the example archive, or drop your .txt files into that folder and update the filename in the config.

• Tests and PyTorch: unit tests import torch. To run pytest on your host you must have a Python environment where torch is installed. Running tests inside the Docker CPU image is usually easiest

Usage examples

Sample generation (replace outputs/last.ckpt with your checkpoint):

python -m scripts.run sample --ckpt outputs/last.ckpt --start_tokens "The Where fight Gandalf" --num_tokens 128

Config

See configs/default.yaml for hyperparameters, dataset paths, and training options.

Tests and coverage

Run tests with coverage locally:

PYTHONPATH=./src pytest --cov=lotr_transformer --cov-report=term-missing -q

Notes

Based on "Attention is All You Need" (Vaswani et al., 2017). This repo uses a decoder-only design since next-token prediction treats input and output sequences the same. For an illustrated explanation of the original Transformer, see: http://jalammar.github.io/illustrated-transformer/