ScratchGPT is a Python project that implements a small-scale transformer-based language model from scratch. It is designed for educational purposes, allowing developers to explore the internals of a transformer model without the complexity of large-scale frameworks. The project provides functionality for training the model on custom datasets and generating text from a prompt.
We want to allow people to experiment easily with any sequence-to-sequence problems. This package is simple to understand, simple to use - show us your projects using ScratchGPT.
- Custom transformer architecture implementation
- Training on user-provided text data
- Text generation using the trained model
- Command-line interfaces for training and inference
- Custom Transformer Architecture: A from-the-ground-up implementation of a decoder-only transformer, including Multi-Head Self-Attention , Feed-Forward layers, and Layer Normalization.
- Flexible Tokenization: Includes a simple character-level tokenizer and a wrapper for using any tokenizer from the Hugging Face Hub.
- Configurable Training: Easily configure model architecture (e.g., embedding_size, num_heads) and training parameters (e.g., learning_rate, batch_size) via a scratch_gpt.yaml file.
- Command-Line Interfaces: Comes with user-friendly CLIs for both training the model and performing inference.
- Pre-tokenization Caching: Caches tokenized datasets to disk for significantly faster startup on subsequent training runs.
- Python 3.12+
uvfor dependency management
-
Clone the repository:
git clone https://github.com/LabStrangeLoop/scratchgpt.git cd scratchgpt -
Install dependencies using uv:
uv sync --all-groups -
Install from pip:
pip install scratchgpt
Please take a look at the simple example in the examples folder.
Note: Some examples require additional dependencies. To run all examples, install the optional dependencies:
uv sync --extra examples-dependenciesTo train the model on your custom dataset, run the train command. This will create an experiment folder containing the model weights, tokenizer files, and configuration.
uv run train -t <path_to_training_data> -e <experiment_folder>
-d, --data_source: Path to the training data file or folder-e, --experiment: Path to the folder where experiment checkpoints will be saved-t, --tokenizer: (Optional) The Hugging Face Hub tokenizer to use (default: "gpt2")
To generate text using a trained model, use infer command:
uv run infer -e <experiment_folder> [-dv <device>] [-m <max_tokens>]
-e, --experiment: Path to the folder containing the trained model-dv, --device: Device to run the model on (default: "cuda")-m, --max_tokens: Maximum number of tokens to generate (default: 512)
This project allows you to create your own tokenizers easily or bootstraps huggingface tokenizers for you to use.
The repository is organized to separate concerns, making it easy to navigate.
scratchgpt/train.py: Main training script.scratchgpt/infer.py: Inference script for text generation.scratchgpt/config.py: Contains all Pydantic configuration models.scratchgpt/model/model.py: The core Transformer model implementation.scratchgpt/training/trainer.py: Orchestrates the training and validation loops.scratchgpt/tokenizer/: Tokenizer implementations, including wrappers for Hugging Face.scratchgpt/model_io.py: Utilities for saving and loading models and tokenizers.tests/: Unit tests for the project.
This project uses various development tools:
mypyfor static type checkingrufffor formatting and standard adherencepytestfor testing
Run the following commands to ensure code quality:
uv run ruff check --fix .
uv run mypy scratchgpt
uv run pytest ./tests/
- Apply SOTA optimizations
Contributions are welcome! Please feel free to submit a Pull Request.
export UV_PUBLISH_USERNAME=__token__
export UV_PUBLISH_PASSWORD=
uv build -vv --wheel
uv publish --publish-url https://upload.pypi.org/legacy/
- Aleksandr Yeganov
- Dario Cazzani