diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2dc6768 --- /dev/null +++ b/.gitignore @@ -0,0 +1,42 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +env/ +venv/ +ENV/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# Data and models +*.pt +*.pth +*.ckpt +data/ +*.txt +!requirements.txt + +# Logs +*.log diff --git a/README.md b/README.md index bd95aab..9789352 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,196 @@ -# GPT -A character-level GPT model built from scratch. +# Character-Level GPT Model + +A minimal implementation of GPT (Generative Pre-trained Transformer) for character-level language modeling, built from scratch using PyTorch. + +## Overview + +This project implements a character-level language model based on the GPT (Generative Pre-trained Transformer) architecture. Unlike word-level or subword tokenization, this model operates directly on individual characters, making it simpler and more flexible for various text generation tasks. + +## Features + +- **Pure PyTorch Implementation**: Built from scratch with no external dependencies except PyTorch and NumPy +- **Character-Level Tokenization**: Simple tokenization that works with any text +- **Multi-Head Self-Attention**: Implements the core transformer attention mechanism +- **Positional Embeddings**: Encodes position information for sequence modeling +- **Layer Normalization**: Stabilizes training +- **Residual Connections**: Enables training of deeper networks +- **Text Generation**: Supports temperature-based sampling and top-k sampling + +## Architecture + +The model consists of: +- Token embedding layer +- Positional embedding layer +- Multiple transformer blocks, each containing: + - Multi-head self-attention mechanism + - Feed-forward neural network + - Layer normalization + - Residual connections +- Output projection layer + +## Installation + +```bash +pip install -r requirements.txt +``` + +Requirements: +- Python 3.7+ +- PyTorch 2.0+ +- NumPy 1.24+ + +## Usage + +### Basic Training + +```python +from train import train, generate_text +import torch + +# Prepare your text +text = """Your training text goes here...""" + +# Set device +device = 'cuda' if torch.cuda.is_available() else 'cpu' + +# Train the model +model, tokenizer = train( + text, + n_embd=384, # Embedding dimension + num_heads=6, # Number of attention heads + num_layers=6, # Number of transformer layers + block_size=256, # Maximum context length + batch_size=64, # Batch size + max_iters=5000, # Training iterations + learning_rate=3e-4, + device=device +) + +# Generate text +generated = generate_text( + model, tokenizer, + prompt="Your prompt", + max_new_tokens=500, + temperature=0.8, + device=device +) +print(generated) +``` + +### Running the Training Script + +```bash +python train.py +``` + +This will train a model on sample text and generate some examples. + +### Running Examples + +```bash +python example.py +``` + +This demonstrates different use cases including Shakespeare-style text generation. + +## Model Components + +### 1. `model.py` - Core Model Architecture + +- **Head**: Single attention head +- **MultiHeadAttention**: Multiple attention heads running in parallel +- **FeedForward**: Position-wise feed-forward network +- **Block**: Complete transformer block +- **GPTLanguageModel**: Main model class + +### 2. `train.py` - Training and Utilities + +- **CharacterTokenizer**: Simple character-level tokenizer +- **train()**: Main training function +- **generate_text()**: Text generation function +- **get_batch()**: Mini-batch generation +- **estimate_loss()**: Loss estimation on train/val sets + +### 3. `example.py` - Usage Examples + +Demonstrates different use cases and configurations. + +## Hyperparameters + +Key hyperparameters you can tune: + +- `n_embd`: Embedding dimension (default: 384) +- `num_heads`: Number of attention heads (default: 6) +- `num_layers`: Number of transformer blocks (default: 6) +- `block_size`: Maximum context length (default: 256) +- `batch_size`: Training batch size (default: 64) +- `learning_rate`: Learning rate (default: 3e-4) +- `dropout`: Dropout rate (default: 0.2) + +## Text Generation + +The model supports several generation strategies: + +```python +# Basic generation +generated = model.generate(context, max_new_tokens=100) + +# With temperature (higher = more random) +generated = model.generate(context, max_new_tokens=100, temperature=0.8) + +# With top-k sampling (only sample from k most likely tokens) +generated = model.generate(context, max_new_tokens=100, temperature=0.8, top_k=10) +``` + +## Training Tips + +1. **Start Small**: Begin with smaller models (fewer layers, smaller embedding dimension) to iterate quickly +2. **Monitor Loss**: Watch both training and validation loss to detect overfitting +3. **Adjust Learning Rate**: If loss plateaus, try reducing the learning rate +4. **Increase Data**: More training data generally leads to better results +5. **Context Length**: Longer `block_size` allows the model to capture longer dependencies but requires more memory + +## Example Output + +After training on Shakespeare text: + +``` +Prompt: 'First' +Generated: First Citizen: +We know't, we know the people. + +All: +Resolved. resolved. + +First Citizen: +Let us kill him... +``` + +## Project Structure + +``` +GPT/ +├── model.py # GPT model architecture +├── train.py # Training script and utilities +├── example.py # Usage examples +├── requirements.txt # Python dependencies +├── .gitignore # Git ignore file +└── README.md # This file +``` + +## How It Works + +1. **Tokenization**: Text is converted to a sequence of integer indices (one per character) +2. **Embedding**: Each character index is converted to a dense vector +3. **Positional Encoding**: Position information is added to each token +4. **Transformer Blocks**: Self-attention and feed-forward layers process the sequence +5. **Prediction**: Output layer predicts the next character +6. **Generation**: Sample from predictions to generate new text + +## License + +MIT License - Feel free to use this code for learning and experimentation. + +## Acknowledgments + +Based on the GPT architecture from the paper "Attention is All You Need" and subsequent GPT models by OpenAI. Inspired by Andrej Karpathy's educational content on transformers and language models. diff --git a/example.py b/example.py new file mode 100644 index 0000000..44b9142 --- /dev/null +++ b/example.py @@ -0,0 +1,153 @@ +""" +Example script demonstrating how to use the character-level GPT model +""" + +import torch +from model import GPTLanguageModel +from train import CharacterTokenizer, train, generate_text + + +def example_shakespeare(): + """Example using Shakespeare-like text""" + + # Sample Shakespeare text + text = """ + First Citizen: + Before we proceed any further, hear me speak. + + All: + Speak, speak. + + First Citizen: + You are all resolved rather to die than to famish? + + All: + Resolved. Resolved. + + First Citizen: + First, you know Caius Marcius is chief enemy to the people. + + All: + We know't, we know't. + + First Citizen: + Let us kill him, and we'll have corn at our own price. + Is't a verdict? + + All: + No more talking on't; let it be done: away, away! + + Second Citizen: + One word, good citizens. + + First Citizen: + We are accounted poor citizens, the patricians good. + What authority surfeits on would relieve us: if they + would yield us but the superfluity, while it were + wholesome, we might guess they relieved us humanely; but + they think we are too dear: the leanness that afflicts + us, the object of our misery, is as an inventory to + particularise their abundance; our sufferance is a gain + to them Let us revenge this with our pikes, ere we + become rakes: for the gods know I speak this in hunger + for bread, not in thirst for revenge. + """ + + device = 'cuda' if torch.cuda.is_available() else 'cpu' + print(f"Using device: {device}\n") + + # Train model + print("Training on Shakespeare text...") + model, tokenizer = train( + text, + n_embd=128, + num_heads=4, + num_layers=4, + block_size=64, + batch_size=16, + max_iters=2000, + learning_rate=1e-3, + eval_interval=500, + eval_iters=50, + dropout=0.1, + device=device + ) + + # Generate text with different prompts + prompts = ["First", "Let us", "We are"] + + print("\n" + "=" * 70) + print("GENERATED TEXT SAMPLES") + print("=" * 70) + + for prompt in prompts: + generated = generate_text( + model, tokenizer, prompt, + max_new_tokens=150, + temperature=0.8, + top_k=10, + device=device + ) + print(f"\nPrompt: '{prompt}'") + print("-" * 70) + print(generated) + print("-" * 70) + + +def example_custom_text(): + """Example using custom text""" + + # You can use any text here + custom_text = """ + The quick brown fox jumps over the lazy dog. + Pack my box with five dozen liquor jugs. + How vexingly quick daft zebras jump! + The five boxing wizards jump quickly. + Sphinx of black quartz, judge my vow. + Two driven jocks help fax my big quiz. + """ + + device = 'cuda' if torch.cuda.is_available() else 'cpu' + print(f"Using device: {device}\n") + + print("Training on custom text...") + model, tokenizer = train( + custom_text, + n_embd=64, + num_heads=4, + num_layers=3, + block_size=32, + batch_size=8, + max_iters=1000, + learning_rate=1e-3, + eval_interval=250, + eval_iters=50, + dropout=0.1, + device=device + ) + + # Generate + print("\nGenerating text...") + generated = generate_text( + model, tokenizer, "The", + max_new_tokens=100, + temperature=1.0, + device=device + ) + print(f"\nGenerated:\n{generated}") + + +if __name__ == '__main__': + print("=" * 70) + print("CHARACTER-LEVEL GPT - EXAMPLES") + print("=" * 70) + print("\nExample 1: Shakespeare-style text generation") + print("=" * 70) + + example_shakespeare() + + print("\n\n" + "=" * 70) + print("Example 2: Custom text generation") + print("=" * 70) + + example_custom_text() diff --git a/model.py b/model.py new file mode 100644 index 0000000..37130a1 --- /dev/null +++ b/model.py @@ -0,0 +1,164 @@ +""" +Character-level GPT Model +A minimal implementation of GPT (Generative Pre-trained Transformer) for character-level language modeling. +""" + +import torch +import torch.nn as nn +from torch.nn import functional as F + + +class Head(nn.Module): + """Single head of self-attention""" + + def __init__(self, n_embd, head_size, block_size, dropout=0.1): + super().__init__() + self.key = nn.Linear(n_embd, head_size, bias=False) + self.query = nn.Linear(n_embd, head_size, bias=False) + self.value = nn.Linear(n_embd, head_size, bias=False) + self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size))) + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + B, T, C = x.shape + k = self.key(x) # (B, T, head_size) + q = self.query(x) # (B, T, head_size) + + # Compute attention scores + head_size = k.shape[-1] + wei = q @ k.transpose(-2, -1) * (head_size ** -0.5) # (B, T, T) + wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T) + wei = F.softmax(wei, dim=-1) # (B, T, T) + wei = self.dropout(wei) + + # Perform weighted aggregation of values + v = self.value(x) # (B, T, head_size) + out = wei @ v # (B, T, head_size) + return out + + +class MultiHeadAttention(nn.Module): + """Multiple heads of self-attention in parallel""" + + def __init__(self, n_embd, num_heads, head_size, block_size, dropout=0.1): + super().__init__() + self.heads = nn.ModuleList([Head(n_embd, head_size, block_size, dropout) for _ in range(num_heads)]) + self.proj = nn.Linear(n_embd, n_embd) + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + out = torch.cat([h(x) for h in self.heads], dim=-1) + out = self.dropout(self.proj(out)) + return out + + +class FeedForward(nn.Module): + """Simple feed-forward network""" + + def __init__(self, n_embd, dropout=0.1): + super().__init__() + self.net = nn.Sequential( + nn.Linear(n_embd, 4 * n_embd), + nn.ReLU(), + nn.Linear(4 * n_embd, n_embd), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +class Block(nn.Module): + """Transformer block: communication followed by computation""" + + def __init__(self, n_embd, num_heads, block_size, dropout=0.1): + super().__init__() + head_size = n_embd // num_heads + self.sa = MultiHeadAttention(n_embd, num_heads, head_size, block_size, dropout) + self.ffwd = FeedForward(n_embd, dropout) + self.ln1 = nn.LayerNorm(n_embd) + self.ln2 = nn.LayerNorm(n_embd) + + def forward(self, x): + x = x + self.sa(self.ln1(x)) + x = x + self.ffwd(self.ln2(x)) + return x + + +class GPTLanguageModel(nn.Module): + """Character-level GPT Language Model""" + + def __init__(self, vocab_size, n_embd=384, num_heads=6, num_layers=6, block_size=256, dropout=0.1): + super().__init__() + self.block_size = block_size + + # Token and position embeddings + self.token_embedding_table = nn.Embedding(vocab_size, n_embd) + self.position_embedding_table = nn.Embedding(block_size, n_embd) + + # Transformer blocks + self.blocks = nn.Sequential(*[Block(n_embd, num_heads, block_size, dropout) for _ in range(num_layers)]) + + # Final layer norm and output projection + self.ln_f = nn.LayerNorm(n_embd) + self.lm_head = nn.Linear(n_embd, vocab_size) + + def forward(self, idx, targets=None): + B, T = idx.shape + + # idx and targets are both (B, T) tensor of integers + tok_emb = self.token_embedding_table(idx) # (B, T, C) + pos_emb = self.position_embedding_table(torch.arange(T, device=idx.device)) # (T, C) + x = tok_emb + pos_emb # (B, T, C) + x = self.blocks(x) # (B, T, C) + x = self.ln_f(x) # (B, T, C) + logits = self.lm_head(x) # (B, T, vocab_size) + + if targets is None: + loss = None + else: + B, T, C = logits.shape + logits = logits.view(B*T, C) + targets = targets.view(B*T) + loss = F.cross_entropy(logits, targets) + + return logits, loss + + def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): + """ + Generate new tokens given a context + + Args: + idx: (B, T) array of indices in the current context + max_new_tokens: number of new tokens to generate + temperature: sampling temperature (higher = more random) + top_k: if set, only sample from top k most likely tokens + """ + for _ in range(max_new_tokens): + # Crop idx to the last block_size tokens + idx_cond = idx[:, -self.block_size:] + + # Get predictions + logits, _ = self(idx_cond) + + # Focus only on the last time step + logits = logits[:, -1, :] # (B, C) + + # Apply temperature + logits = logits / temperature + + # Optionally crop to top k + if top_k is not None: + v, _ = torch.topk(logits, min(top_k, logits.size(-1))) + logits[logits < v[:, [-1]]] = -float('Inf') + + # Apply softmax to get probabilities + probs = F.softmax(logits, dim=-1) # (B, C) + + # Sample from the distribution + idx_next = torch.multinomial(probs, num_samples=1) # (B, 1) + + # Append sampled index to the running sequence + idx = torch.cat((idx, idx_next), dim=1) # (B, T+1) + + return idx diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..5259e95 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +torch>=2.0.0 +numpy>=1.24.0 diff --git a/train.py b/train.py new file mode 100644 index 0000000..7335ba5 --- /dev/null +++ b/train.py @@ -0,0 +1,214 @@ +""" +Training script for character-level GPT model +""" + +import torch +import torch.nn as nn +from model import GPTLanguageModel + + +class CharacterTokenizer: + """Simple character-level tokenizer""" + + def __init__(self, text): + # Get all unique characters from the text + self.chars = sorted(list(set(text))) + self.vocab_size = len(self.chars) + + # Create mappings from characters to integers and vice versa + self.stoi = {ch: i for i, ch in enumerate(self.chars)} + self.itos = {i: ch for i, ch in enumerate(self.chars)} + + def encode(self, text): + """Convert text to a list of integers""" + return [self.stoi[c] for c in text] + + def decode(self, indices): + """Convert a list of integers back to text""" + return ''.join([self.itos[i] for i in indices]) + + +def get_batch(data, block_size, batch_size, device='cpu'): + """Generate a small batch of data of inputs x and targets y""" + ix = torch.randint(len(data) - block_size, (batch_size,)) + x = torch.stack([data[i:i+block_size] for i in ix]) + y = torch.stack([data[i+1:i+block_size+1] for i in ix]) + x, y = x.to(device), y.to(device) + return x, y + + +@torch.no_grad() +def estimate_loss(model, train_data, val_data, block_size, batch_size, eval_iters, device='cpu'): + """Estimate loss on train and val sets""" + out = {} + model.eval() + for split, data in [('train', train_data), ('val', val_data)]: + losses = torch.zeros(eval_iters) + for k in range(eval_iters): + X, Y = get_batch(data, block_size, batch_size, device) + logits, loss = model(X, Y) + losses[k] = loss.item() + out[split] = losses.mean() + model.train() + return out + + +def train(text, + n_embd=384, + num_heads=6, + num_layers=6, + block_size=256, + batch_size=64, + max_iters=5000, + learning_rate=3e-4, + eval_interval=500, + eval_iters=200, + dropout=0.2, + device='cpu'): + """ + Train a character-level GPT model + + Args: + text: Input text to train on + n_embd: Embedding dimension + num_heads: Number of attention heads + num_layers: Number of transformer layers + block_size: Maximum context length + batch_size: Batch size for training + max_iters: Maximum number of training iterations + learning_rate: Learning rate + eval_interval: How often to evaluate + eval_iters: Number of iterations for evaluation + dropout: Dropout rate + device: Device to train on ('cpu' or 'cuda') + """ + + # Initialize tokenizer + tokenizer = CharacterTokenizer(text) + print(f"Vocabulary size: {tokenizer.vocab_size}") + print(f"Unique characters: {''.join(tokenizer.chars)}") + + # Encode the entire text + data = torch.tensor(tokenizer.encode(text), dtype=torch.long) + + # Split into train and validation sets + n = int(0.9 * len(data)) + train_data = data[:n] + val_data = data[n:] + + # Initialize model + model = GPTLanguageModel( + vocab_size=tokenizer.vocab_size, + n_embd=n_embd, + num_heads=num_heads, + num_layers=num_layers, + block_size=block_size, + dropout=dropout + ) + model = model.to(device) + + # Print model size + print(f"Model parameters: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M") + + # Create optimizer + optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate) + + # Training loop + for step in range(max_iters): + # Every eval_interval, evaluate the loss on train and val sets + if step % eval_interval == 0 or step == max_iters - 1: + losses = estimate_loss(model, train_data, val_data, block_size, batch_size, eval_iters, device) + print(f"Step {step}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}") + + # Sample a batch of data + xb, yb = get_batch(train_data, block_size, batch_size, device) + + # Evaluate the loss + logits, loss = model(xb, yb) + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + + return model, tokenizer + + +def generate_text(model, tokenizer, prompt, max_new_tokens=500, temperature=1.0, top_k=None, device='cpu'): + """Generate text from a trained model""" + model.eval() + + # Encode the prompt + context = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long, device=device) + + # Generate + generated = model.generate(context, max_new_tokens=max_new_tokens, temperature=temperature, top_k=top_k) + + # Decode and return + return tokenizer.decode(generated[0].tolist()) + + +if __name__ == '__main__': + # Example usage + print("Character-level GPT Training") + print("=" * 50) + + # Check if CUDA is available + device = 'cuda' if torch.cuda.is_available() else 'cpu' + print(f"Using device: {device}") + + # Sample text (you can replace this with any text file) + sample_text = """ + To be, or not to be, that is the question: + Whether 'tis nobler in the mind to suffer + The slings and arrows of outrageous fortune, + Or to take arms against a sea of troubles + And by opposing end them. To die—to sleep, + No more; and by a sleep to say we end + The heart-ache and the thousand natural shocks + That flesh is heir to: 'tis a consummation + Devoutly to be wish'd. To die, to sleep; + To sleep, perchance to dream—ay, there's the rub: + For in that sleep of death what dreams may come, + When we have shuffled off this mortal coil, + Must give us pause—there's the respect + That makes calamity of so long life. + """ + + print(f"\nTraining on {len(sample_text)} characters...") + print("\nStarting training...") + + # Train the model (using smaller parameters for quick testing) + model, tokenizer = train( + sample_text, + n_embd=128, + num_heads=4, + num_layers=4, + block_size=64, + batch_size=16, + max_iters=1000, + learning_rate=1e-3, + eval_interval=200, + eval_iters=50, + dropout=0.1, + device=device + ) + + print("\n" + "=" * 50) + print("Generating text...") + print("=" * 50) + + # Generate some text + prompt = "To be" + generated = generate_text(model, tokenizer, prompt, max_new_tokens=200, temperature=0.8, device=device) + print(f"\nPrompt: '{prompt}'") + print(f"\nGenerated text:\n{generated}") + + # Save the model + torch.save({ + 'model_state_dict': model.state_dict(), + 'vocab_size': tokenizer.vocab_size, + 'chars': tokenizer.chars, + 'stoi': tokenizer.stoi, + 'itos': tokenizer.itos, + }, 'gpt_model.pt') + print("\n" + "=" * 50) + print("Model saved to gpt_model.pt")