Try It Yourself

# Check what we're working with
import torch

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")

if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print("Great! Training will be fast.")
else:
    print("Device: CPU")
    print("No GPU found. Training will work but be slower (10-20x).")
    print("Consider reducing num_samples in the data loading step.")

# Import everything we need
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer, get_linear_schedule_with_warmup
from datasets import load_dataset
from tqdm.auto import tqdm
import numpy as np

print("All imports successful!")
print("\nIf you got any errors above, install missing packages with:")
print("  pip install transformers datasets torch tqdm")

# Load GPT-2 (small, 124M parameters)
model_name = "gpt2"

print(f"Loading {model_name}...")
print("This downloads ~500MB the first time, then caches locally.\n")

model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# GPT-2 doesn't have a padding token by default, so we add one
# We just reuse the EOS token. common practice and works fine
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = tokenizer.eos_token_id

# Move to GPU if available
if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")
model = model.to(device)

# How big is this thing?
total_params = sum(p.numel() for p in model.parameters())
print(f"✓ Model loaded!")
print(f"  Parameters: {total_params:,}")
print(f"  Device: {device}")
print(f"  Memory: ~{total_params * 4 / 1e9:.1f} GB (in float32)")

def generate_response(model, tokenizer, instruction, max_new_tokens=100):
    """
    Generate a response to an instruction.
    
    We format the prompt in "Alpaca style". A specific template that works well
    for instruction following. You'll see this same format in the training data.
    """
    prompt = f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{instruction}

### Response:
"""
    
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():  # Don't track gradients for inference
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.7,  # Some randomness (0 = deterministic, 1 = very random)
            top_p=0.9,  # Nucleus sampling
            do_sample=True,  # Use sampling instead of greedy
            pad_token_id=tokenizer.eos_token_id
        )
    
    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Extract just the response part (after "### Response:")
    response = full_text.split("### Response:\n")[-1].strip()
    
    return response

# Test base model on a few questions
test_instructions = [
    "What is the capital of France?",
    "Write a haiku about programming.",
    "Explain machine learning in one sentence.",
]

print("BASE MODEL (before fine-tuning):")
print("=" * 70)
print("Watch how it fails to actually answer the questions...\n")

for instruction in test_instructions:
    print(f"Q: {instruction}")
    response = generate_response(model, tokenizer, instruction)
    # Truncate long responses
    if len(response) > 200:
        response = response[:200] + "..."
    print(f"A: {response}")
    print("-" * 70)

# Load the Alpaca dataset
print("Loading Alpaca dataset from HuggingFace...")
raw_dataset = load_dataset("yahma/alpaca-cleaned", split="train")

# Use a small subset for quick training
# Feel free to adjust this:
#   - 100 samples: Very fast, model learns a bit (2-3 min on GPU)
#   - 500 samples: Good learning, reasonable time (10-15 min on GPU)
#   - 5000+ samples: Better results, longer training (1+ hour)
num_samples = 500

raw_dataset = raw_dataset.select(range(num_samples))

print(f"\n✓ Dataset loaded: {len(raw_dataset)} training examples")
print(f"\nHere's what one example looks like:")
print(f"  Instruction: {raw_dataset[0]['instruction']}")
if raw_dataset[0]['input']:
    print(f"  Input: {raw_dataset[0]['input']}")
print(f"  Output: {raw_dataset[0]['output'][:100]}...")
print(f"\nThe model will learn to generate 'Output' given 'Instruction' (and 'Input' if present).")

class InstructionDataset(Dataset):
    """
    Dataset for instruction fine-tuning.
    
    The key insight here is in the LABEL MASKING. We only compute loss on the response
    tokens, not the instruction tokens. Why? Because we want the model to learn
    to GENERATE responses, not to predict the instruction itself.
    
    This is crucial. Without it, the model would waste capacity learning to 
    predict the instruction template, which is useless.
    """
    
    def __init__(self, data, tokenizer, max_length=256):
        self.data = data
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.data)
    
    def format_example(self, example):
        """Format in Alpaca style (same as our generate function)."""
        if example['input']:
            # Some examples have an additional 'input' field for context
            prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{example['instruction']}

### Input:
{example['input']}

### Response:
"""
        else:
            prompt = f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{example['instruction']}

### Response:
"""
        return prompt, example['output']
    
    def __getitem__(self, idx):
        example = self.data[idx]
        prompt, response = self.format_example(example)
        
        # Tokenize prompt and response separately (important!)
        prompt_tokens = self.tokenizer.encode(prompt, add_special_tokens=True)
        response_tokens = self.tokenizer.encode(response, add_special_tokens=False)
        
        # Combine: [prompt tokens] + [response tokens] + [EOS]
        input_ids = prompt_tokens + response_tokens + [self.tokenizer.eos_token_id]
        
        # Create labels: -100 for prompt (ignored in loss), actual tokens for response
        # This is the key to supervised fine-tuning!
        labels = [-100] * len(prompt_tokens) + response_tokens + [self.tokenizer.eos_token_id]
        
        # Truncate if too long
        if len(input_ids) > self.max_length:
            input_ids = input_ids[:self.max_length]
            labels = labels[:self.max_length]
        
        # Pad to max_length (makes batching easier)
        padding_length = self.max_length - len(input_ids)
        input_ids = input_ids + [self.tokenizer.pad_token_id] * padding_length
        labels = labels + [-100] * padding_length  # Ignore padding in loss
        attention_mask = [1] * (self.max_length - padding_length) + [0] * padding_length
        
        return {
            'input_ids': torch.tensor(input_ids),
            'attention_mask': torch.tensor(attention_mask),
            'labels': torch.tensor(labels),
        }

# Create dataset and dataloader
train_dataset = InstructionDataset(raw_dataset, tokenizer, max_length=256)
train_loader = DataLoader(
    train_dataset, 
    batch_size=4,  # Small batch size to fit in memory
    shuffle=True   # Randomize order each epoch
)

print(f"✓ Created dataset with {len(train_dataset)} samples")
print(f"  Batches per epoch: {len(train_loader)}")
print(f"  Batch size: 4")
print(f"\nEach batch contains:")
print(f"  - input_ids: The tokenized text (prompt + response)")
print(f"  - attention_mask: Which tokens to pay attention to (1) vs ignore (0)")
print(f"  - labels: What to predict (-100 for prompt, actual tokens for response)")

# Training hyperparameters
learning_rate = 5e-5  # Standard for fine-tuning (0.00005)
num_epochs = 1        # One pass through the data
warmup_steps = 50     # Gradually increase LR for first 50 steps
max_grad_norm = 1.0   # Clip gradients to prevent instability

# Set up optimizer (AdamW is standard for transformers)
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# Learning rate scheduler (warmup then linear decay)
total_steps = len(train_loader) * num_epochs
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=total_steps
)

print("Training configuration:")
print(f"  Learning rate: {learning_rate}")
print(f"  Epochs: {num_epochs}")
print(f"  Steps per epoch: {len(train_loader)}")
print(f"  Total steps: {total_steps}")
print(f"  Warmup steps: {warmup_steps}")
print(f"  Gradient clipping: {max_grad_norm}")
print(f"\nEstimated time: ~10-15 minutes on GPU, ~2 hours on CPU")

# The actual training loop
print("\nStarting training...")
print("Watch the loss go down! (Lower = better)")
print("\nMetrics explained:")
print("  - loss: How wrong the model is (lower = better)")
print("  - avg_loss: Running average of loss")
print("  - ppl: Perplexity (e^loss), another way to measure quality")
print("\nGo grab a coffee. This'll take a few minutes...\n")

model.train()  # Put model in training mode

for epoch in range(num_epochs):
    total_loss = 0
    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}")
    
    for step, batch in enumerate(progress_bar):
        # Move batch to GPU
        batch = {k: v.to(device) for k, v in batch.items()}
        
        # Forward pass: compute loss
        outputs = model(
            input_ids=batch['input_ids'],
            attention_mask=batch['attention_mask'],
            labels=batch['labels']
        )
        loss = outputs.loss
        
        # Backward pass: compute gradients
        optimizer.zero_grad()  # Clear old gradients
        loss.backward()        # Compute new gradients
        
        # Clip gradients (prevent explosions)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
        
        # Update weights
        optimizer.step()
        scheduler.step()
        
        # Track metrics
        total_loss += loss.item()
        avg_loss = total_loss / (step + 1)
        
        # Update progress bar
        progress_bar.set_postfix({
            'loss': f'{loss.item():.4f}',
            'avg_loss': f'{avg_loss:.4f}',
            'ppl': f'{np.exp(avg_loss):.2f}'
        })
    
    print(f"\n✓ Epoch {epoch+1} complete!")
    print(f"  Final average loss: {avg_loss:.4f}")
    print(f"  Final perplexity: {np.exp(avg_loss):.2f}")

print("\n🎉 Training complete!")
print("\nThe model has now seen 500 examples of how to follow instructions.")
print("Let's see if it actually learned anything...")

# Switch to evaluation mode (disables dropout, etc.)
model.eval()

print("FINE-TUNED MODEL (after training):")
print("=" * 70)
print("Same questions as before. Notice the difference?\n")

for instruction in test_instructions:
    print(f"Q: {instruction}")
    response = generate_response(model, tokenizer, instruction)
    print(f"A: {response}")
    print("-" * 70)

print("\nMuch better, right?")
print("\nThe model isn't perfect (it's only seen 500 examples), but it's actually")
print("trying to answer the questions now instead of just rambling.")
print("\nThat's the power of supervised fine-tuning!")

# Let's try some more examples to really see what it can do
additional_tests = [
    "List three benefits of exercise.",
    "What is Python used for?",
    "Explain what a neural network is in simple terms.",
    "Write a short poem about the ocean.",
]

print("\nLet's try some different questions:")
print("=" * 70)

for instruction in additional_tests:
    print(f"\nQ: {instruction}")
    response = generate_response(model, tokenizer, instruction)
    print(f"A: {response}")
    print("-" * 70)

print("\n**Key observation:** The model has learned the *pattern* of instruction-following,")
print("not just memorized specific facts. It generalizes to new questions!")
print("\nThough sometimes it gets a bit.. creative. (That's LLMs for you.)")

def compute_perplexity(model, dataloader, device):
    """
    Compute perplexity on a dataset.
    
    Perplexity = e^(average loss)
    
    Think of it as: "On average, how many equally-likely tokens could come next?"
    Lower is better. Random guessing on a 50k vocab = perplexity of 50,000.
    A well-trained model on instructions = perplexity of 5-10.
    """
    model.eval()
    total_loss = 0
    total_tokens = 0
    
    with torch.no_grad():  # No gradients needed for evaluation
        for batch in tqdm(dataloader, desc="Computing perplexity"):
            batch = {k: v.to(device) for k, v in batch.items()}
            
            outputs = model(
                input_ids=batch['input_ids'],
                attention_mask=batch['attention_mask'],
                labels=batch['labels']
            )
            
            # Count non-masked tokens (only response tokens, not prompt)
            num_tokens = (batch['labels'] != -100).sum().item()
            total_loss += outputs.loss.item() * num_tokens
            total_tokens += num_tokens
    
    avg_loss = total_loss / total_tokens
    perplexity = np.exp(avg_loss)
    
    return perplexity, avg_loss

# Compute perplexity on the training set
# (In practice, you'd use a held-out validation set, but we're keeping it simple)
print("\nEvaluating model quality...")
perplexity, loss = compute_perplexity(model, train_loader, device)

print(f"\n✓ Evaluation complete!")
print(f"  Loss: {loss:.4f}")
print(f"  Perplexity: {perplexity:.2f}")
print(f"\nInterpretation:")
print(f"  - Perplexity < 10: Excellent")
print(f"  - Perplexity 10-20: Good")
print(f"  - Perplexity 20-50: Okay")
print(f"  - Perplexity > 50: Needs more training")
print(f"\nYour model: ", end="")
if perplexity < 10:
    print("Excellent! 🎉")
elif perplexity < 20:
    print("Good! 👍")
elif perplexity < 50:
    print("Okay. More training would help.")
else:
    print("Needs more training. Try more epochs or more data.")

def compute_diversity(responses):
    """
    Compute diversity metrics for generated text.
    
    Distinct-1: Percentage of unique words (unigrams)
    Distinct-2: Percentage of unique word pairs (bigrams)
    
    Why does this matter? If the model always says "the the the the" you'd have
    low diversity even if perplexity looks okay. Diversity catches mode collapse.
    """
    all_unigrams = []
    all_bigrams = []
    
    for response in responses:
        tokens = response.lower().split()
        all_unigrams.extend(tokens)
        # Create pairs of consecutive words
        all_bigrams.extend(zip(tokens[:-1], tokens[1:]))
    
    # What fraction of words/pairs are unique?
    distinct_1 = len(set(all_unigrams)) / len(all_unigrams) if all_unigrams else 0
    distinct_2 = len(set(all_bigrams)) / len(all_bigrams) if all_bigrams else 0
    
    return distinct_1, distinct_2

# Generate a bunch of responses for diversity analysis
diversity_prompts = [
    "Tell me about machine learning.",
    "Explain artificial intelligence.",
    "What is deep learning?",
    "Describe natural language processing.",
    "Explain what data science is.",
]

print("\nGenerating responses for diversity analysis...")
responses = [generate_response(model, tokenizer, p) for p in diversity_prompts]
d1, d2 = compute_diversity(responses)

print(f"\n✓ Diversity analysis complete!")
print(f"  Distinct-1 (unique words): {d1:.2%}")
print(f"  Distinct-2 (unique word pairs): {d2:.2%}")
print(f"\nInterpretation:")
print(f"  - Distinct-1 > 40%: Good variety")
print(f"  - Distinct-1 20-40%: Okay")
print(f"  - Distinct-1 < 20%: Mode collapse (model stuck repeating itself)")
print(f"\nYour model: ", end="")
if d1 > 0.4:
    print("Good variety! 🎉")
elif d1 > 0.2:
    print("Okay diversity.")
else:
    print("Warning: Low diversity. Try different sampling parameters.")

# Save model and tokenizer to disk
save_path = "./my_finetuned_model"

print(f"Saving model to {save_path}...")
model.save_pretrained(save_path)
tokenizer.save_pretrained(save_path)
print("✓ Model saved!")

# Show what got saved
import os
print(f"\nSaved files:")
total_size = 0
for f in sorted(os.listdir(save_path)):
    size = os.path.getsize(os.path.join(save_path, f)) / 1e6
    total_size += size
    print(f"  {f}: {size:.1f} MB")

print(f"\nTotal size: {total_size:.1f} MB")
print(f"\nYou can now load this model anytime with:")
print(f"  model = AutoModelForCausalLM.from_pretrained('{save_path}')")
print(f"  tokenizer = AutoTokenizer.from_pretrained('{save_path}')")

# Let's verify the saved model actually works
print("Testing saved model (to make sure saving worked)...")

loaded_model = AutoModelForCausalLM.from_pretrained(save_path)
loaded_tokenizer = AutoTokenizer.from_pretrained(save_path)
loaded_model = loaded_model.to(device)
loaded_model.eval()

test_instruction = "What is the meaning of life?"
response = generate_response(loaded_model, loaded_tokenizer, test_instruction)

print(f"\n✓ Loaded model from disk successfully!")
print(f"\nTest question: {test_instruction}")
print(f"Answer: {response}")
print(f"\nLooks good! Your model is saved and ready to use.")