A character-level language diffusion model for text generation. The model is a modified version of the nanochat gpt implementation and is trained on Tiny Shakespeare! It is only 10.7 million parameters, so you can try it out locally!
Installation
# Clone the repository git clone < repository-url > cd tiny-diffusion # Install dependencies (Python 3.10+) uv sync
Quick Start
The file training.py puts the weights in weights/diffusion_model.pt . The sample and animation files load the model from this file.
Train Your Own Model
Currently, the weights are already provided for you! It took me around half an hour to train this model for 20,000 steps on 4xA100s. But if you want to retrain the model again, run:
# Train from scratch on Shakespeare uv run training.py # Training will save checkpoints to weights/diffusion_model.pt
Generate Text
To generate a continuous stream of output (currently 30 context lengths), run:
... continue reading