nanoGPT — Character-Level Transformer from Scratch

Live Demo:https://nanogpt.streamlit.app

A CPU-friendly character-level GPT from scratch in PyTorch featuring modern architecture (RoPE, RMSNorm, GQA) and a modular RAG pipeline with Chroma DB. No Hugging Face, no magic — every line is readable and educational.

Note: This model is currently trained exclusively on Shakespearean text. It generates text in the style and vocabulary of Shakespeare and is not a general-purpose instruction-following model.

~250 lines of model code  |  ~200 lines of training code

---

System Architecture

---

App Interface

---

Quickstart

1. Create the virtual environment

python -m venv .venv

2. Change execution policy to allow script activation (Windows/PowerShell only)

Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser

3. Activate the virtual environment

.venv/Scripts/activate

3. Install dependencies

pip install -r requirements.txt

4. Download a sample corpus (Shakespeare ~1MB)

python data.py --download shakespeare
# → downloads input.txt

Or use your own text file — just point --data at it.

5. Train

# CPU-friendly tiny model (~3 min on CPU)
python train.py --preset tiny

# Default model (~10M params, needs GPU or ~30 min CPU)
python train.py

# Your own corpus
python train.py --data my_notes.txt --preset tiny

Training prints loss every 100 steps and generates text samples every 500 steps:

step   500 | train loss 1.8234 | val loss 1.9102 | lr 3.00e-04 | elapsed 42s

──── sample @ step 500 ────
GLOUCESTER:
What, art thou so brave? then let us hear
The king himself hath spoke it; and for my part...

6. Generate text

# Single prompt
python generate.py \
    --checkpoint checkpoints/ckpt_02000.pt \
    --prompt "To be or"

# Interactive REPL
python generate.py \
    --checkpoint checkpoints/ckpt_02000.pt \
    --interactive

---

Configuration

Edit config.py or pass flags. Key knobs:

Parameter	Default	Effect
block_size	256	Context window. Longer = more memory
n_embd	384	Embedding size. Larger = smarter, slower
n_head	6	Attention heads. Must divide n_embd
n_layer	6	Transformer depth
dropout	0.2	Regularisation. 0 for tiny datasets
batch_size	64	Reduce if OOM
learning_rate	3e-4	AdamW LR
temperature	0.8	Generation: higher = more random
top_k	40	0 = pure sampling, 40 = focused

Presets

# In config.py
TINY_CONFIG   # n_embd=128, n_layer=4  — trains in minutes on CPU
SMALL_CONFIG  # n_embd=384, n_layer=6  — ~10M params, needs GPU

---

Architecture Explained

Input token IDs  →  Token Embedding
                 +  Positional Embedding
                        ↓
              [TransformerBlock × n_layer]
              ┌─────────────────────────────┐
              │  LayerNorm                  │
              │  MultiHeadSelfAttention ←── │── causal mask prevents
              │  + residual                 │   seeing future tokens
              │  LayerNorm                  │
              │  FeedForward (MLP)       ←──│── per-token processing
              │  + residual                 │
              └─────────────────────────────┘
                        ↓
              LayerNorm  →  Linear head  →  logits (vocab_size)
                        ↓
              softmax  →  P(next character)

Key concepts implemented:

• Causal (masked) self-attention — tokens can only attend to past positions
• Multi-head attention — multiple attention patterns in parallel
• Positional embeddings — learned position encodings
• Residual connections — skip connections around each sub-layer
• Pre-LayerNorm — normalise before sub-layers (more stable than original)
• Weight tying — token embedding and LM head share weights
• Cosine LR schedule — warmup + cosine decay
• Gradient clipping — prevents exploding gradients
• Top-k sampling — controlled randomness at generation time

---

Learning Path

After you've trained this:

1. Swap embeddings — try sinusoidal positional encodings instead of learned
2. Flash Attention — torch.nn.functional.scaled_dot_product_attention (one line swap, much faster)
3. BPE tokenization — replace CharTokenizer with tiktoken for subword tokens
4. Fine-tuning — load a pretrained gpt2 via Hugging Face, freeze some layers
5. RAG — feed retrieved document chunks into the prompt instead of baking knowledge into weights

---

Resume Training

python train.py --resume checkpoints/ckpt_02000.pt

---

Expected Loss (Shakespeare)

Steps	Train Loss	Val Loss	Quality
0	~4.2	~4.2	Random characters
500	~2.0	~2.1	Recognisable words
2000	~1.5	~1.6	Shakespearean structure
5000	~1.2	~1.4	Convincing prose

---

Stage 2: Retrieval-Augmented Generation (RAG)

A fully modular RAG pipeline that allows you to perform semantic search and context injection over your corpus using ChromaDB and sentence-transformers.

1. Install RAG dependencies

pip install -r rag/requirements.txt

2. Run the Interactive RAG Demo

python rag/demo.py

Sample Demo Output

------------------------- RAG Demo -------------------------
Corpus : D:\Code\Repo\SLM\nanogpt\input.txt
Index  : D:\Code\Repo\SLM\nanogpt\rag_index

[1/4] Building ChromaDB index from corpus ...
      This embeds ~4000 chunks — takes ~2 min on CPU, once only.

Loading embedding model: sentence-transformers/all-MiniLM-L6-v2
  Embedding dim : 384
  Embedding 4648 chunks from 'input.txt' ...
  Added 4648 chunks | Total: 4648
VectorStore saved -> D:\Code\Repo\SLM\nanogpt\rag_index/ (4648 vectors, dim=384)

Index built in 64.2s and saved to 'D:\Code\Repo\SLM\nanogpt\rag_index/'

Index ready: 4648 chunks indexed

-------------- Retrieval Demo (no LLM needed) --------------

Query: 'What does Hamlet say about death and dying?'
  [1] score=0.5423 ##########
       kill your husband.

LADY ANNE:
Why, then he is alive.

GLOUCESTER:
Nay, he is dead; and slain by Edward's hand.

LADY ANNE:
In thy foul throat thou li...

Query: 'A betrayal and murder of a king'
  [1] score=0.6159 ############
       die your king,
As ours by murder, to make him a king!
Edward thy son, which now is Prince of Wales,
For Edward my son, which was Prince of Wales,
Die...

Query: 'Love and romantic feelings'
  [1] score=0.4751 #########
       ng first create!
O heavy lightness! serious vanity!
Mis-shapen chaos of well-seeming forms!
Feather of lead, bright smoke, cold fire,
sick health!
Sti...

To run with full LLM-augmented generation:

1. Download Ollama
2. Run ollama pull llama3.2:1b
3. Start the Ollama server: ollama serve
4. Run python rag/demo.py again.

Streamlit Deployment

You can easily deploy this model and the RAG pipeline as a web application on Streamlit Cloud for free.

1. Ensure your trained checkpoint (e.g., ckpt_01000.pt) and tokenizer.pkl are in checkpoints_v2/ and are tracked by Git (you may need to force-add them or update .gitignore).
2. If using RAG, ensure rag_index/chroma.sqlite3 is committed.
3. Push your repository to GitHub.
4. Go to share.streamlit.io and click New app.
5. Select your GitHub repository, branch, and set the Main file path to app.py.
6. Click Deploy. Streamlit Cloud will automatically install dependencies from requirements.txt and launch the app.

Troubleshooting

• _pickle.UnpicklingError when loading checkpoint: Recent versions of PyTorch default to weights_only=True for security. Because our checkpoint contains a custom GPTConfig object, ensure you use torch.load(..., weights_only=False) in app.py.
• App crashes on memory on Streamlit Cloud: The free tier has a 1GB memory limit. Ensure you are using the --preset tiny checkpoint (~3.5M params) rather than the default 10M parameter model, which might spike RAM during generation.
• Missing Tokenizer Error: The tokenizer is saved separately from the checkpoint (tokenizer.pkl). Ensure both files are committed in the same directory.
• RAG takes a long time on first load: sentence-transformers and chromadb are heavy dependencies. The very first load on Streamlit Cloud might take a minute or two as it downloads the embedding models.

License

This project is licensed under the MIT License - see the LICENSE file for details.