Project 1 - Build an LLM Playground

# minimal_example.py
import torch
from models import GPTModel, ModelConfig
from generation import TextGenerator
from generation.strategies import GenerationConfig

# 1. Define a small model (runs on CPU)
config = ModelConfig(
    vocab_size=50257,
    hidden_size=256,
    num_layers=4,
    num_heads=4,
    intermediate_size=1024,
    max_position_embeddings=512,
    use_rope=True,
    use_rmsnorm=True,
    use_swiglu=True,
)

# 2. Build model
model = GPTModel(config)
print(f"Parameters: {model.num_parameters()/1e6:.1f}M")

# 3. Forward pass
input_ids = torch.randint(0, 50257, (1, 16))
output = model(input_ids=input_ids)
print(f"Logits shape: {output.logits.shape}")
# → torch.Size([1, 16, 50257])

# 4. Compute loss
labels = torch.randint(0, 50257, (1, 16))
output_with_loss = model(input_ids=input_ids, labels=labels)
print(f"Loss: {output_with_loss.loss.item():.4f}")
# → Loss: 10.8xxx (random init ≈ log(50257))

# 5. Train one step
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
output_with_loss.loss.backward()
optimizer.step()
optimizer.zero_grad()
print("✅ Forward + backward pass successful")

from data import WebCrawler, CrawlConfig
import asyncio

config = CrawlConfig(
    max_depth=3,            # BFS depth limit
    max_pages=50_000,       # Total page cap
    delay=1.0,              # Seconds between requests per domain
    max_concurrent=10,      # Simultaneous connections
    respect_robots=True,    # Always honour robots.txt
    allowed_domains=[       # Optional domain whitelist
        "wikipedia.org",
        "arxiv.org",
    ],
)

crawler = WebCrawler(config)
pages = asyncio.run(crawler.crawl([
    "https://en.wikipedia.org/wiki/Machine_learning",
    "https://arxiv.org/abs/2303.08774",
]))

crawler.save_results("./data/raw/crawl.jsonl")
print(f"Collected: {len(pages)} pages")

from data import CommonCrawlProcessor, CommonCrawlConfig

config = CommonCrawlConfig(
    crawl_id="CC-MAIN-2023-50",    # Specific crawl snapshot
    num_segments=10,                # WET file segments to process
    max_pages_per_segment=100_000,  # Pages per segment
    languages=["en"],               # Language filter
    output_dir="./data/raw/cc",
)

processor = CommonCrawlProcessor(config)
total = processor.process_and_save("./data/raw/cc/output.jsonl")
print(f"Extracted: {total:,} documents")

from data import RefinedWebCleaner, RefinedWebConfig
from data import DolmaTextCleaner, DolmaConfig
from data import FineWebCleaner, FineWebConfig

# RefinedWeb
rw = RefinedWebCleaner(RefinedWebConfig())
stats = rw.process_file("./raw.jsonl", "./refined.jsonl")
print(f"Passed: {stats['passed']}/{stats['total']}")

# Dolma
dolma = DolmaTextCleaner(DolmaConfig())
stats = dolma.process_file("./raw.jsonl", "./dolma.jsonl")

# FineWeb (educational filter)
fw = FineWebCleaner(FineWebConfig(min_educational_score=0.3))
stats = fw.process_file("./raw.jsonl", "./fineweb.jsonl")

from data.tokenization.bpe_tokenizer import BPEConfig, ByteLevelBPETokenizer

config = BPEConfig(
    vocab_size=50257,
    min_frequency=2,
    special_tokens={
        "bos_token":  "<|startoftext|>",
        "eos_token":  "<|endoftext|>",
        "unk_token":  "<|unk|>",
        "pad_token":  "<|pad|>",
    }
)

# Train on corpus
tokenizer = ByteLevelBPETokenizer(config)
with open("./data/processed/clean.jsonl") as f:
    texts = [json.loads(l)["text"] for l in f]

tokenizer.train(texts)
tokenizer.save("./tokenizers/bpe_50k")

# Use
ids = tokenizer.encode("Hello, world!")
# → [1, 31373, 11, 995, 0, 2]  (with BOS/EOS)

text = tokenizer.decode(ids, skip_special_tokens=True)
# → "Hello, world!"

# Batch encoding with padding
batch = tokenizer.encode_batch(
    ["Short text", "A longer sentence here"],
    padding=True,
    max_length=128,
    truncation=True,
)
# → {"input_ids": [[...], [...]], "attention_mask": [[...], [...]]}

from data.tokenization.utils import (
    tokenize_corpus_to_bin,
    split_dataset,
    estimate_token_count,
)

# Estimate before committing to full tokenization
estimate = estimate_token_count(
    "./data/processed/",
    tokenizer,
    sample_size=1000,
)
print(f"Estimated: {estimate['estimated_total_tokens_billions']:.2f}B tokens")

# Tokenize → binary shards
stats = tokenize_corpus_to_bin(
    input_path="./data/processed/",   # Directory of JSONL files
    output_dir="./data/tokenized/",
    tokenizer=tokenizer,
    max_tokens_per_file=100_000_000,  # 100M tokens per shard
    add_eos=True,                      # Separate documents with EOS
)
print(f"Shards: {stats['num_shards']}, Tokens: {stats['total_tokens']:,}")

# Split train/val/test
splits = split_dataset(
    data_dir="./data/tokenized/",
    output_dir="./data/split/",
    val_ratio=0.005,     # 0.5% validation
    test_ratio=0.001,    # 0.1% test
)
print(f"Train: {splits.train_tokens:,} tokens")

from models import GPTModel, ModelConfig

# GPT-2 small equivalent
config = ModelConfig(
    vocab_size=50257,
    hidden_size=768,
    num_layers=12,
    num_heads=12,
    num_kv_heads=12,         # = num_heads → standard MHA
    intermediate_size=3072,
    max_position_embeddings=2048,
    use_rope=False,          # GPT-2 uses absolute positions
    use_rmsnorm=False,       # GPT-2 uses LayerNorm
    use_swiglu=False,        # GPT-2 uses GeLU
    use_gqa=False,
    tie_word_embeddings=True,
)

# LLaMA-style
config = ModelConfig(
    vocab_size=32000,
    hidden_size=4096,
    num_layers=32,
    num_heads=32,
    num_kv_heads=8,          # GQA: 4 queries share 1 KV head
    intermediate_size=11008,
    max_position_embeddings=4096,
    use_rope=True,           # Rotary embeddings
    use_rmsnorm=True,        # RMSNorm
    use_swiglu=True,         # SwiGLU
    use_gqa=True,            # Grouped Query Attention
    rope_theta=10000.0,
)

model = GPTModel(config)
print(f"Parameters: {model.num_parameters()/1e9:.2f}B")

from models.architectures.deepseek import DeepSeekConfig, DeepSeekModel

config = DeepSeekConfig(
    vocab_size=102400,
    hidden_size=5120,
    num_layers=60,
    num_heads=128,
    num_kv_heads=128,
    intermediate_size=12288,
    # MLA
    use_mla=True,
    kv_lora_rank=512,
    q_lora_rank=1536,
    qk_rope_head_dim=64,
    qk_nope_head_dim=128,
    v_head_dim=128,
    # MoE
    use_moe=True,
    num_experts=160,
    num_experts_per_token=6,
    num_shared_experts=2,
    moe_intermediate_size=1536,
    first_k_dense_layers=1,
    # Load balancing
    expert_load_balance_coef=0.01,
    router_z_loss_coef=0.001,
)
model = DeepSeekModel(config)

from models.architectures.gemma import GemmaConfig, GemmaModel

config = GemmaConfig(
    vocab_size=256000,
    hidden_size=3072,
    num_layers=28,
    num_heads=16,
    num_kv_heads=16,
    intermediate_size=24576,
    # Gemma-specific
    logit_soft_cap=30.0,        # Final logit tanh cap
    attn_logit_softcap=50.0,    # Attention logit tanh cap
    use_logit_softcap=True,
    hidden_act="gelu",          # "gelu" or "relu2"
    # Embedding scaling: multiply by sqrt(hidden_size)
    # GemmaRMSNorm: uses (1 + weight) formulation
)
model = GemmaModel(config)

from models.architectures.qwen import QwenConfig, QwenModel

config = QwenConfig(
    vocab_size=151936,
    hidden_size=4096,
    num_layers=32,
    num_heads=32,
    num_kv_heads=32,
    intermediate_size=22016,
    kv_channels=128,
    # Qwen-specific
    seq_length=8192,            # Training context length
    use_dynamic_ntk=True,       # Auto-scale RoPE beyond training len
    use_logn_attn=True,         # Log-N attention scaling
    # Dynamic NTK: if seq_len > seq_length:
    #   alpha = (seq_len/seq_length)^(dim/(dim-2))
    #   base = original_base * alpha
)
model = QwenModel(config)

from generation import TextGenerator
from generation.strategies import GenerationConfig

gen_config = GenerationConfig(
    max_new_tokens=512,
    temperature=0.7,
    top_k=50,
    top_p=0.9,
    repetition_penalty=1.1,
    no_repeat_ngram_size=3,
)

generator = TextGenerator(model, tokenizer, gen_config)

# Standard generation
text = generator.generate(
    "The theory of neural scaling laws suggests",
    strategy="sampling",       # "greedy" | "sampling" | "beam_search"
)
print(text)

# Streaming (token by token)
for token in generator.generate_stream("Explain transformers:"):
    print(token, end="", flush=True)

# Batch generation
results = generator.batch_generate(
    prompts=["Question 1:", "Question 2:", "Question 3:"],
    strategy="greedy",
)

from models import GPTModel, ModelConfig
from training import PretrainTrainer, PretrainConfig, PretrainDataset

# Model
model = GPTModel(ModelConfig(
    vocab_size=50257,
    hidden_size=768,
    num_layers=12,
    num_heads=12,
    intermediate_size=3072,
    use_rope=True,
    use_rmsnorm=True,
    use_swiglu=True,
))

# Config
config = PretrainConfig(
    data_path="./data/split",
    max_seq_length=2048,
    batch_size=32,
    gradient_accumulation_steps=8,   # Effective batch = 256
    max_steps=100_000,
    warmup_steps=2_000,
    learning_rate=3e-4,
    min_lr=3e-5,                     # 10× decay at end
    weight_decay=0.1,
    bf16=True,
    output_dir="./checkpoints/pretrain",
    save_steps=5_000,
    eval_steps=1_000,
    wandb_project="llm_playground",
)

# Data
train_data = PretrainDataset(config.data_path, config.max_seq_length, "train")
val_data   = PretrainDataset(config.data_path, config.max_seq_length, "val")

# Train
trainer = PretrainTrainer(model, config)
trainer.train(train_data, val_data)

# LoRA: W' = W + (B @ A) × (α/r)
# A: (r, in_features)  - kaiming init
# B: (out_features, r) - zeros init → starts identical to base
# Trainable params: r×(in+out) vs in×out baseline

from models import GPTModel
from training import SFTTrainer, SFTConfig, SFTDataset

model = GPTModel.from_pretrained("./checkpoints/pretrain/best")

config = SFTConfig(
    data_path="./data/sft",
    max_length=2048,
    batch_size=8,
    gradient_accumulation_steps=4,
    num_epochs=3,
    learning_rate=2e-5,
    use_lora=True,
    lora_r=16,               # Rank (higher = more capacity)
    lora_alpha=32,           # Scaling = alpha/r = 2.0
    lora_dropout=0.05,
    lora_target_modules=[
        "q_proj", "k_proj", "v_proj", "out_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    bf16=True,
    output_dir="./checkpoints/sft",
)

# Data format options (all supported):
# 1. ChatML messages format
# 2. Alpaca instruction format
# 3. Simple prompt/response format
dataset = SFTDataset("./data/sft", tokenizer, config)
trainer = SFTTrainer(model, tokenizer, config)
trainer.train(dataset)

from models import GPTModel
from training.rlhf import RewardModel, RewardModelConfig, RewardModelTrainer
from training.rlhf import PreferenceDataset

# Build reward model from SFT checkpoint
backbone = GPTModel.from_pretrained("./checkpoints/sft/final")
rm = RewardModel(backbone, RewardModelConfig(hidden_size=768))

# Preference data format:
# {"prompt": "...", "chosen": "...", "rejected": "..."}
dataset = PreferenceDataset("./data/preference", tokenizer)
trainer = RewardModelTrainer(rm, RewardModelConfig())
trainer.train(dataset)

# Evaluate
chosen_ids  = tokenizer.encode("Good response", ...)
rejected_ids = tokenizer.encode("Bad response", ...)
reward_good = rm(chosen_ids)
reward_bad  = rm(rejected_ids)
print(f"Good: {reward_good.item():.3f}, Bad: {reward_bad.item():.3f}")
# Good: 2.341, Bad: -1.872

from models import GPTModel
from training.rlhf import PPOTrainer, PPOConfig

actor   = GPTModel.from_pretrained("./checkpoints/sft/final")
ref     = GPTModel.from_pretrained("./checkpoints/sft/final")  # frozen
critic  = GPTModel.from_pretrained("./checkpoints/sft/final")
rm      = RewardModel(GPTModel.from_pretrained("./checkpoints/rm/final"), ...)

config = PPOConfig(
    total_steps=1000,
    batch_size=64,
    ppo_epochs=4,            # Inner optimization epochs
    gamma=1.0,               # Discount (1.0 for language tasks)
    lam=0.95,                # GAE lambda
    clip_range=0.2,          # PPO epsilon
    vf_coef=0.1,             # Value loss weight
    kl_coef=0.1,             # Initial KL penalty coefficient
    target_kl=6.0,           # Target KL for adaptive control
    max_new_tokens=256,
    normalize_rewards=True,  # Normalize RM scores per batch
    reward_clip=10.0,        # Clip extreme rewards
)

trainer = PPOTrainer(actor, critic, ref, rm, tokenizer, config)

# Load prompts
with open("./data/ppo_prompts.txt") as f:
    prompts = f.readlines()

trainer.train(prompts)

from training.rlhf import MathVerifier, CodeVerifier, FormatVerifier

# Math: compare extracted numbers
math_v = MathVerifier()
result = math_v.compute_reward(
    prompt="Solve: 15 × 23 = ?",
    response="Let me calculate: 15 × 23 = 15 × 20 + 15 × 3 "
             "= 300 + 45 = 345. The answer is 345.",
    ground_truth="345"
)
print(f"Reward: {result.reward}, Passed: {result.passed}")
# Reward: 1.1 (bonus for showing work), Passed: True

# Code: execute and compare output
code_v = CodeVerifier(timeout=10)
result = code_v.compute_reward(
    prompt="Write code to print 'Hello, World!'",
    response="```python\nprint('Hello, World!')\n```",
    ground_truth="Hello, World!"
)
print(f"Reward: {result.reward}")  # 1.0 if output matches

# Format: check structural requirements
fmt_v = FormatVerifier()
result = fmt_v.compute_reward(
    prompt="Return a JSON object",
    response='{"name": "Alice", "age": 30}',
    ground_truth="json"
)
print(f"Reward: {result.reward}")  # 1.0 for valid JSON

from evaluation import EvaluationSuite

suite = EvaluationSuite(model=model, tokenizer=tokenizer)

# Language modeling quality
ppl_result = suite.evaluate_lm(test_texts)
print(f"Perplexity: {ppl_result['perplexity']:.2f}")

# Generation quality
results = suite.evaluate_generation(
    hypotheses=generated_texts,
    references=reference_texts,
)
print(f"BLEU-4:  {results['bleu']:.1f}")
print(f"ROUGE-1: {results['rouge1']:.1f}")
print(f"ROUGE-L: {results['rougeL']:.1f}")

from evaluation import MMLUBenchmark, GSM8KBenchmark
from generation import TextGenerator

# MMLU — 57 academic subjects, 5-shot
mmlu = MMLUBenchmark(model, tokenizer)
results = mmlu.evaluate(
    data_dir="./data/benchmarks/mmlu",
    num_few_shot=5,
)
print(f"MMLU Average: {results['mmlu_avg']*100:.1f}%")
print(f"Best subject: {max(results['per_subject'], key=results['per_subject'].get)}")

# GSM8K — grade school math, chain-of-thought
generator = TextGenerator(model, tokenizer, GenerationConfig())
gsm8k = GSM8KBenchmark(model, tokenizer, generator)
result = gsm8k.evaluate(gsm8k_examples, num_few_shot=8)
print(f"GSM8K: {result.accuracy*100:.1f}%")

from evaluation import HumanEvalCollector

collector = HumanEvalCollector(storage_path="./data/human_eval")

# Submit comparisons
collector.submit_comparison(
    prompt="Explain quantum entanglement simply.",
    model_a="llm-playground-v1",
    response_a="Quantum entanglement is when two particles...",
    model_b="llm-playground-v2",
    response_b="Imagine two coins that always show opposite faces...",
    winner="b",              # Human judge preferred model B
    judge_id="evaluator_01",
)

# Get leaderboard
board = collector.get_leaderboard()
for rank in board:
    print(
        f"#{rank['rank']} {rank['model']:30s} "
        f"ELO: {rank['elo']:6.1f}  "
        f"Win Rate: {rank['win_rate']:5.1f}%  "
        f"95% CI: [{rank['ci_95_low']:.1f}, {rank['ci_95_high']:.1f}]"
    )

import requests

class LLMPlaygroundClient:
    def __init__(self, base_url="http://localhost:8000", api_key="dev-key-123"):
        self.base = base_url
        self.headers = {"X-API-Key": api_key}

    def chat(self, messages, model="demo", **kwargs):
        resp = requests.post(
            f"{self.base}/api/v1/chat/completions",
            json={"model": model, "messages": messages, **kwargs},
            headers=self.headers,
        )
        return resp.json()["choices"][0]["message"]["content"]

    def stream(self, messages, model="demo"):
        import json
        resp = requests.post(
            f"{self.base}/api/v1/chat/completions",
            json={"model": model, "messages": messages, "stream": True},
            headers=self.headers,
            stream=True,
        )
        for line in resp.iter_lines():
            if line and line.startswith(b"data: "):
                data = line[6:]
                if data.strip() == b"[DONE]":
                    break
                chunk = json.loads(data)
                token = chunk["choices"][0]["delta"].get("content", "")
                if token:
                    yield token

# Usage
client = LLMPlaygroundClient()

# Standard call
answer = client.chat([{"role": "user", "content": "Hello!"}])
print(answer)

# Streaming
for token in client.stream([{"role": "user", "content": "Tell me a story"}]):
    print(token, end="", flush=True)

# 1. Flash Attention (2-4× faster attention)
pip install flash-attn --no-build-isolation
# Automatically used when available (checked in MultiHeadAttention)

# 2. Gradient checkpointing (reduce VRAM by ~40%)
model.enable_gradient_checkpointing()

# 3. Mixed precision
config = PretrainConfig(bf16=True)  # BF16 preferred for stability

# 4. Gradient accumulation (simulate larger batch)
config = PretrainConfig(
    batch_size=4,
    gradient_accumulation_steps=32,  # Effective batch = 128
)

# 5. Compile (PyTorch 2.0+)
model = torch.compile(model)

# Loss is NaN
# → Reduce learning rate (3e-4 → 1e-4)
# → Enable gradient clipping (already in our trainer: max_grad_norm=1.0)
# → Check for bad data (empty strings, corrupted tokens)

# Loss not decreasing
# → Check data is shuffled
# → Verify labels are not all -100 (SFT masking bug)
# → Check learning rate schedule (print scheduler.get_lr())

# VRAM OOM during training
config = PretrainConfig(
    batch_size=4,                      # Reduce batch
    gradient_accumulation_steps=64,    # Compensate
    bf16=True,                         # Half precision
)
model.enable_gradient_checkpointing()  # Trade compute for memory

# SFT: model not following instructions
# → Check loss masking is working (some labels should be real, not -100)
# → Increase training data diversity
# → Reduce learning rate (avoid catastrophic forgetting)