Spaces:
Sleeping
Sleeping
File size: 4,390 Bytes
d6a7ec4 7784b50 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | ---
title: Tiny Transformer
sdk: gradio
app_file: app.py
pinned: false
---
# Tiny Transformer
A compact GPT-style language model built from scratch in PyTorch. This repo is designed to show the fundamentals recruiters actually care about: clean architecture, readable math, reproducible training, tests, and an end-to-end demo path from raw text to generated tokens.
## What Makes This Worth Looking At
- Implements a decoder-only Transformer without Hugging Face or high-level training frameworks.
- Includes causal self-attention, multi-head attention, residual blocks, layer norm, embeddings, generation, and checkpointing.
- Ships with character and byte-pair encoding tokenizers so the model can train on any plain-text file.
- Keeps the code small enough to understand in one sitting, but structured like production Python.
- Includes smoke tests for masking, shapes, tokenization, attention export, and generation behavior.
## Quickstart
```bash
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -e ".[dev]"
```
Train on the included sample text:
```bash
tiny-transformer train --data data/tiny_shakespeare_excerpt.txt --steps 300 --device cpu
```
Use the optional BPE tokenizer, gradient accumulation, and mixed precision when you want a stronger local run:
```bash
tiny-transformer train \
--data data/tiny_shakespeare_excerpt.txt \
--tokenizer bpe \
--bpe-vocab-size 128 \
--grad-accum-steps 4 \
--amp \
--device mps
```
Generate text from a checkpoint:
```bash
tiny-transformer generate --checkpoint runs/tiny-transformer.pt --prompt "To be" --max-new-tokens 160
```
Export an attention heatmap:
```bash
tiny-transformer attention --checkpoint runs/tiny-transformer.pt --prompt "To be" --output runs/attention.svg
```
Launch the local playground:
```bash
tiny-transformer serve --checkpoint runs/tiny-transformer.pt
```
Deploy the hosted playground:
```bash
pip install huggingface_hub
hf auth login
hf repos create axay28/tiny-transformer --type space --space-sdk gradio --public --exist-ok
git remote add space https://huggingface.co/spaces/axay28/tiny-transformer
git push space main
```
Run tests:
```bash
pytest
```
## Project Layout
```text
src/tiny_transformer/
cli.py Command line interface for training and generation
config.py Model and training configuration
data.py Text dataset and batching utilities
model.py GPT-style Transformer implementation
tokenizer.py Character-level tokenizer
train.py Training loop, evaluation, checkpointing
visualize.py Attention heatmap export
web.py Local generation playground
tests/ Unit and smoke tests
data/ Tiny sample corpus
```
## Architecture
The model is intentionally small, but it follows the same structure as larger decoder-only LLMs:
1. Token and positional embeddings convert IDs into vectors.
2. Each Transformer block applies pre-norm causal self-attention.
3. Feed-forward layers expand and compress the hidden dimension.
4. Residual connections preserve gradient flow.
5. A tied-size language modeling head predicts the next token.
The attention mask is causal, so each position can only attend to itself and previous positions.
```mermaid
flowchart LR
A["Raw text corpus"] --> B["Char or BPE tokenizer"]
B --> C["Token IDs"]
C --> D["Contiguous train/val batches"]
D --> E["Token + position embeddings"]
E --> F1["LayerNorm"]
F1 --> F2["Masked multi-head self-attention"]
F2 --> F3["Residual add"]
F3 --> F4["LayerNorm"]
F4 --> F5["Feed-forward MLP"]
F5 --> F6["Residual add"]
F6 --> F7["Repeat for N layers"]
F7 --> G["Final layer norm"]
G --> H["Language modeling head"]
H --> I["Next-token logits"]
I --> J["Cross-entropy loss during training"]
I --> K["Top-k sampling during generation"]
F2 --> L["Attention heatmap export"]
K --> M["Local web playground"]
```
## Example Configuration
The CLI defaults train quickly on CPU. For the included tiny corpus, the command uses a
32-token context window; for larger text files, 128 tokens is a good next step:
```python
ModelConfig(
vocab_size=128,
block_size=128,
n_layer=4,
n_head=4,
n_embd=128,
dropout=0.1,
)
```
Increase `n_layer`, `n_head`, and `n_embd` for a stronger demo once the training loop is validated.
|