{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1c02cf30",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "# Module 20: Capstone - Building TinyGPT End-to-End\n",
    "\n",
    "Welcome to the capstone project of TinyTorch! You've built an entire ML framework from scratch across 19 modules. Now it's time to put it all together and build something amazing: **TinyGPT** - a complete transformer-based language model.\n",
    "\n",
    "## 🔗 Prerequisites & Progress\n",
    "**You've Built**: The complete TinyTorch framework with 19 specialized modules\n",
    "**You'll Build**: A complete end-to-end ML system demonstrating production capabilities\n",
    "**You'll Enable**: Understanding of how modern AI systems work from tensor to text generation\n",
    "\n",
    "**Connection Map**:\n",
    "```\n",
    "Modules 01-19 → Capstone Integration → Complete TinyGPT System\n",
    "(Foundation)    (Systems Thinking)    (Real AI Application)\n",
    "```\n",
    "\n",
    "## Learning Objectives\n",
    "By the end of this capstone, you will:\n",
    "1. **Integrate** all TinyTorch modules into a cohesive system\n",
    "2. **Build** a complete TinyGPT model with training and inference\n",
    "3. **Optimize** the system with quantization, pruning, and acceleration\n",
    "4. **Benchmark** performance against accuracy trade-offs\n",
    "5. **Demonstrate** end-to-end ML systems engineering\n",
    "\n",
    "This capstone represents the culmination of your journey from basic tensors to a complete AI system!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba68ded0",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 📦 Where This Code Lives in the Final Package\n",
    "\n",
    "**Learning Side:** You work in `modules/20_capstone/capstone_dev.py`  \n",
    "**Building Side:** Code exports to `tinytorch.applications.tinygpt`\n",
    "\n",
    "```python\n",
    "# How to use this module:\n",
    "from tinytorch.applications.tinygpt import TinyGPT, FullPipeline\n",
    "```\n",
    "\n",
    "**Why this matters:**\n",
    "- **Learning:** Complete ML system integrating all previous learning into real application\n",
    "- **Production:** Demonstrates how framework components compose into deployable systems\n",
    "- **Consistency:** Shows the power of modular design and clean abstractions\n",
    "- **Integration:** Validates that our 19-module journey builds something meaningful"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f758fd43",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "exports",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| default_exp applications.tinygpt\n",
    "#| export"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6850420",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🔮 Introduction: From Building Blocks to Intelligence\n",
    "\n",
    "Over the past 19 modules, you've built the complete infrastructure for modern ML:\n",
    "\n",
    "**Foundation (Modules 01-04):** Tensors, activations, layers, and losses\n",
    "**Training (Modules 05-07):** Automatic differentiation, optimizers, and training loops\n",
    "**Architecture (Modules 08-09):** Spatial processing and data loading\n",
    "**Language (Modules 10-14):** Text processing, embeddings, attention, transformers, and KV caching\n",
    "**Optimization (Modules 15-19):** Profiling, acceleration, quantization, compression, and benchmarking\n",
    "\n",
    "Now we integrate everything into **TinyGPT** - a complete language model that demonstrates the power of your framework.\n",
    "\n",
    "```\n",
    "Your Journey:\n",
    "    Tensor Ops → Neural Networks → Training → Transformers → Optimization → TinyGPT\n",
    "    (Module 01)   (Modules 02-07)  (Mod 08-09) (Mod 10-14)    (Mod 15-19)   (Module 20)\n",
    "```\n",
    "\n",
    "This isn't just a demo - it's a production-ready system that showcases everything you've learned about ML systems engineering."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "470a2c0a",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 📊 Systems Architecture: The Complete ML Pipeline\n",
    "\n",
    "This capstone demonstrates how all 19 modules integrate into a complete ML system. Let's visualize the full architecture and understand how each component contributes to the final TinyGPT system.\n",
    "\n",
    "### Complete TinyGPT System Architecture\n",
    "\n",
    "```\n",
    "                        🏗️ TINYGPT COMPLETE SYSTEM ARCHITECTURE 🏗️\n",
    "\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                                   DATA PIPELINE                                     │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│  Raw Text     →    Tokenizer    →    DataLoader    →    Training Loop              │\n",
    "│ \"Hello AI\"         [72,101,..]       Batches(32)        Loss/Gradients             │\n",
    "│ (Module 10)        (Module 10)       (Module 08)       (Modules 05-07)             │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                                 MODEL ARCHITECTURE                                  │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  Token IDs → [Embeddings] → [Positional] → [Dropout] → [Transformer Blocks] → Output │\n",
    "│              (Module 11)    (Module 11)   (Module 03)     (Module 13)              │\n",
    "│                                                                                     │\n",
    "│  Transformer Block Details:                                                         │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Input → [LayerNorm] → [MultiHeadAttention] → [Residual] → [LayerNorm]      │   │\n",
    "│  │           (Module 03)      (Module 12)        (Module 01)   (Module 03)    │   │\n",
    "│  │                                    ↓                                       │   │\n",
    "│  │         [MLP] ← [Residual] ← [GELU] ← [Linear] ← [Linear]                  │   │\n",
    "│  │      (Module 03)  (Module 01)  (Module 02)   (Module 03)                  │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                              GENERATION PIPELINE                                    │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│  Model Output → [Sampling] → [Token Selection] → [Decoding] → Generated Text       │\n",
    "│                (Temperature)    (Greedy/Random)   (Module 10)                      │\n",
    "│                                                                                     │\n",
    "│  With KV Caching (Module 14):                                                      │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Cache Keys/Values → Only Process New Token → O(n) vs O(n²) Complexity      │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                            OPTIMIZATION PIPELINE                                    │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│  Base Model → [Profiling] → [Quantization] → [Pruning] → [Benchmarking] → Optimized │\n",
    "│              (Module 15)   (Module 17)    (Module 18)   (Module 19)                │\n",
    "│                                                                                     │\n",
    "│  Memory Reduction Pipeline:                                                         │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ FP32 (4 bytes) → INT8 (1 byte) → 90% Pruning → 40× Memory Reduction         │   │\n",
    "│  │    200MB      →      50MB      →     5MB     →     Final Size               │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "```\n",
    "\n",
    "### Memory Footprint Analysis for Different Model Sizes\n",
    "\n",
    "```\n",
    "TinyGPT Model Sizes and Memory Requirements:\n",
    "\n",
    "┌──────────────┬────────────────┬─────────────────┬─────────────────┬─────────────────┐\n",
    "│ Model Size   │   Parameters   │ Inference (MB)  │ Training (MB)   │ Quantized (MB)  │\n",
    "├──────────────┼────────────────┼─────────────────┼─────────────────┼─────────────────┤\n",
    "│ TinyGPT-1M   │    1,000,000   │      4.0        │     12.0        │      1.0        │\n",
    "│ TinyGPT-13M  │   13,000,000   │     52.0        │    156.0        │     13.0        │\n",
    "│ TinyGPT-50M  │   50,000,000   │    200.0        │    600.0        │     50.0        │\n",
    "│ TinyGPT-100M │  100,000,000   │    400.0        │   1200.0        │    100.0        │\n",
    "└──────────────┴────────────────┴─────────────────┴─────────────────┴─────────────────┘\n",
    "\n",
    "Memory Breakdown:\n",
    "• Inference = Parameters × 4 bytes (FP32)\n",
    "• Training = Parameters × 12 bytes (params + gradients + optimizer states)\n",
    "• Quantized = Parameters × 1 byte (INT8)\n",
    "```\n",
    "\n",
    "### Critical Systems Properties\n",
    "\n",
    "**Computational Complexity:**\n",
    "- **Attention Mechanism**: O(n² × d) where n=sequence_length, d=embed_dim\n",
    "- **MLP Layers**: O(n × d²) per layer\n",
    "- **Generation**: O(n²) without KV cache, O(n) with KV cache\n",
    "\n",
    "**Memory Scaling:**\n",
    "- **Linear with batch size**: memory = base_memory × batch_size\n",
    "- **Quadratic with sequence length**: attention memory ∝ seq_len²\n",
    "- **Linear with model depth**: memory ∝ num_layers\n",
    "\n",
    "**Performance Characteristics:**\n",
    "- **Training throughput**: ~100-1000 tokens/second (depending on model size)\n",
    "- **Inference latency**: ~1-10ms per token (depending on hardware)\n",
    "- **Memory efficiency**: 4× improvement with quantization, 10× with pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2fa5c74",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "imports",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import time\n",
    "import json\n",
    "from pathlib import Path\n",
    "from typing import Dict, List, Tuple, Optional, Any\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Import all TinyTorch modules (representing 19 modules of work!)\n",
    "### BEGIN SOLUTION\n",
    "# Module 01: Tensor foundation\n",
    "from tinytorch.core.tensor import Tensor\n",
    "\n",
    "# Module 02: Activations\n",
    "from tinytorch.core.activations import ReLU, GELU, Sigmoid\n",
    "\n",
    "# Module 03: Layers\n",
    "from tinytorch.core.layers import Linear, Sequential, Dropout\n",
    "\n",
    "# Module 04: Losses\n",
    "from tinytorch.core.losses import CrossEntropyLoss\n",
    "\n",
    "# Module 05: Autograd (enhances Tensor)\n",
    "from tinytorch.core.autograd import Function\n",
    "\n",
    "# Module 06: Optimizers\n",
    "from tinytorch.core.optimizers import AdamW, SGD\n",
    "\n",
    "# Module 07: Training\n",
    "from tinytorch.core.training import Trainer, CosineSchedule\n",
    "\n",
    "# Module 08: DataLoader\n",
    "from tinytorch.data.loader import DataLoader, TensorDataset\n",
    "\n",
    "# Module 09: Spatial (for potential CNN comparisons)\n",
    "from tinytorch.core.spatial import Conv2d, MaxPool2d\n",
    "\n",
    "# Module 10: Tokenization\n",
    "from tinytorch.text.tokenization import CharTokenizer\n",
    "\n",
    "# Module 11: Embeddings\n",
    "from tinytorch.text.embeddings import Embedding, PositionalEncoding\n",
    "\n",
    "# Module 12: Attention\n",
    "from tinytorch.core.attention import MultiHeadAttention, scaled_dot_product_attention\n",
    "\n",
    "# Module 13: Transformers\n",
    "from tinytorch.models.transformer import GPT, TransformerBlock\n",
    "\n",
    "# Module 14: KV Caching\n",
    "from tinytorch.generation.kv_cache import KVCache\n",
    "\n",
    "# Module 15: Profiling\n",
    "from tinytorch.profiling.profiler import Profiler\n",
    "\n",
    "# Module 16: Acceleration\n",
    "from tinytorch.optimization.acceleration import MixedPrecisionTrainer\n",
    "\n",
    "# Module 17: Quantization\n",
    "from tinytorch.optimization.quantization import quantize_model, QuantizedLinear\n",
    "\n",
    "# Module 18: Compression\n",
    "from tinytorch.optimization.compression import magnitude_prune, structured_prune\n",
    "\n",
    "# Module 19: Benchmarking\n",
    "from tinytorch.benchmarking.benchmark import Benchmark\n",
    "### END SOLUTION\n",
    "\n",
    "print(\"🎉 Successfully imported all 19 TinyTorch modules!\")\n",
    "print(\"📦 Framework Status: COMPLETE\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d6fa877",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 🏗️ Stage 1: Core TinyGPT Architecture\n",
    "\n",
    "We'll build TinyGPT in three systematic stages, each demonstrating different aspects of ML systems engineering:\n",
    "\n",
    "### What We're Building: Complete Transformer Architecture\n",
    "\n",
    "The TinyGPT architecture integrates every component you've built across 19 modules into a cohesive system. Here's how all the pieces fit together:\n",
    "\n",
    "```\n",
    "                          🧠 TINYGPT ARCHITECTURE BREAKDOWN 🧠\n",
    "\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                                INPUT PROCESSING                                     │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│  Token IDs (integers)                                                               │\n",
    "│        │                                                                            │\n",
    "│        ▼                                                                            │\n",
    "│  [Token Embedding] ──────────────── Maps vocab_size → embed_dim                    │\n",
    "│   (Module 11)          ╲                                                            │\n",
    "│        │                ╲                                                           │\n",
    "│        ▼                 ╲─→ [Element-wise Addition] ──────► Dense Vectors         │\n",
    "│  [Positional Encoding] ──╱    (Module 01)                                          │\n",
    "│   (Module 11)          ╱                                                            │\n",
    "│                       ╱                                                             │\n",
    "│        │             ╱                                                              │\n",
    "│        ▼            ╱                                                               │\n",
    "│  [Dropout] ────────╱ ←──────────────── Regularization (Module 03)                │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                              TRANSFORMER PROCESSING                                 │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  For each of num_layers (typically 4-12):                                         │\n",
    "│                                                                                     │\n",
    "│  ┌───────────────────────────────────────────────────────────────────────────┐     │\n",
    "│  │                          TRANSFORMER BLOCK                                │     │\n",
    "│  │                                                                           │     │\n",
    "│  │  Input Vectors (batch, seq_len, embed_dim)                               │     │\n",
    "│  │        │                                                                 │     │\n",
    "│  │        ▼                                                                 │     │\n",
    "│  │  ┌─────────────┐   ┌──────────────────────────────────────────────┐     │     │\n",
    "│  │  │ Layer Norm  │──▶│ Multi-Head Self-Attention (Module 12)        │     │     │\n",
    "│  │  │ (Module 03) │   │                                              │     │     │\n",
    "│  │  └─────────────┘   │ • Query, Key, Value projections              │     │     │\n",
    "│  │                    │ • Scaled dot-product attention               │     │     │\n",
    "│  │                    │ • Multi-head parallel processing             │     │     │\n",
    "│  │                    │ • Output projection                          │     │     │\n",
    "│  │                    └──────────────────────────────────────────────┘     │     │\n",
    "│  │                                     │                                   │     │\n",
    "│  │                                     ▼                                   │     │\n",
    "│  │                    ┌─────────────────────────────────────────┐         │     │\n",
    "│  │  ┌─────────────┐   │ Residual Connection (Module 01)         │         │     │\n",
    "│  │  │             │◄──┤ output = input + attention(input)       │         │     │\n",
    "│  │  │             │   └─────────────────────────────────────────┘         │     │\n",
    "│  │  │             │                                                       │     │\n",
    "│  │  │             ▼                                                       │     │\n",
    "│  │  │       ┌─────────────┐   ┌──────────────────────────────────────┐   │     │\n",
    "│  │  │       │ Layer Norm  │──▶│ Feed-Forward Network (MLP)          │   │     │\n",
    "│  │  │       │ (Module 03) │   │                                     │   │     │\n",
    "│  │  │       └─────────────┘   │ • Linear: embed_dim → 4×embed_dim   │   │     │\n",
    "│  │  │                         │ • GELU Activation (Module 02)       │   │     │\n",
    "│  │  │                         │ • Linear: 4×embed_dim → embed_dim   │   │     │\n",
    "│  │  │                         │ • Dropout                           │   │     │\n",
    "│  │  │                         └──────────────────────────────────────┘   │     │\n",
    "│  │  │                                          │                         │     │\n",
    "│  │  │                                          ▼                         │     │\n",
    "│  │  │                         ┌─────────────────────────────────────────┐   │     │\n",
    "│  │  └─────────────────────────│ Residual Connection (Module 01)         │   │     │\n",
    "│  │                            │ output = input + mlp(input)             │   │     │\n",
    "│  │                            └─────────────────────────────────────────┘   │     │\n",
    "│  └───────────────────────────────────────────────────────────────────────────┘     │\n",
    "│                                           │                                        │\n",
    "│                                           ▼                                        │\n",
    "│                               Next Transformer Block                               │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                                OUTPUT PROCESSING                                    │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│  Final Hidden States (batch, seq_len, embed_dim)                                  │\n",
    "│                          │                                                         │\n",
    "│                          ▼                                                         │\n",
    "│                 [Output Linear Layer] ──────► Logits (batch, seq_len, vocab_size) │\n",
    "│                    (Module 03)                                                     │\n",
    "│                          │                                                         │\n",
    "│                          ▼                                                         │\n",
    "│                    [Softmax + Sampling] ──────► Next Token Predictions            │\n",
    "│                                                                                     │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "```\n",
    "\n",
    "### Systems Focus: Parameter Distribution and Memory Impact\n",
    "\n",
    "Understanding where parameters live in TinyGPT is crucial for optimization:\n",
    "\n",
    "```\n",
    "Parameter Distribution in TinyGPT (embed_dim=128, vocab_size=1000, 4 layers):\n",
    "\n",
    "┌─────────────────────┬─────────────────┬─────────────────┬─────────────────┐\n",
    "│ Component           │ Parameter Count │ Memory (MB)     │ % of Total      │\n",
    "├─────────────────────┼─────────────────┼─────────────────┼─────────────────┤\n",
    "│ Token Embeddings    │    128,000      │      0.5        │     15%         │\n",
    "│ Positional Encoding │     32,768      │      0.1        │      4%         │\n",
    "│ Attention Layers    │    262,144      │      1.0        │     31%         │\n",
    "│ MLP Layers          │    393,216      │      1.5        │     46%         │\n",
    "│ Layer Norms         │      2,048      │      0.01       │      0.2%       │\n",
    "│ Output Projection   │    128,000      │      0.5        │     15%         │\n",
    "├─────────────────────┼─────────────────┼─────────────────┼─────────────────┤\n",
    "│ TOTAL              │    946,176      │      3.6        │    100%         │\n",
    "└─────────────────────┴─────────────────┴─────────────────┴─────────────────┘\n",
    "\n",
    "Key Insights:\n",
    "• MLP layers dominate parameter count (46% of total)\n",
    "• Attention layers are second largest (31% of total)\n",
    "• Embedding tables scale with vocabulary size\n",
    "• Memory scales linearly with embed_dim²\n",
    "```\n",
    "\n",
    "### Why This Architecture Matters\n",
    "\n",
    "**1. Modular Design**: Each component can be optimized independently\n",
    "**2. Scalable**: Architecture works from 1M to 100B+ parameters\n",
    "**3. Interpretable**: Clear information flow through attention and MLP\n",
    "**4. Optimizable**: Each layer type has different optimization strategies\n",
    "\n",
    "Let's implement this step by step, starting with the core TinyGPT class that orchestrates all components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32815de3",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "tinygpt_architecture",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "class TinyGPT:\n",
    "    \"\"\"\n",
    "    Complete GPT implementation integrating all TinyTorch modules.\n",
    "\n",
    "    This class demonstrates how framework components compose into real applications.\n",
    "    Built using modules 01,02,03,11,12,13 as core architecture.\n",
    "\n",
    "    Architecture:\n",
    "    - Token Embeddings (Module 11)\n",
    "    - Positional Encoding (Module 11)\n",
    "    - Transformer Blocks (Module 13)\n",
    "    - Output Linear Layer (Module 03)\n",
    "    - Language Modeling Head (Module 04)\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_size: int, embed_dim: int = 128, num_layers: int = 4,\n",
    "                 num_heads: int = 4, max_seq_len: int = 256, dropout: float = 0.1):\n",
    "        \"\"\"\n",
    "        Initialize TinyGPT with production-inspired architecture.\n",
    "\n",
    "        TODO: Build a complete GPT model using TinyTorch components\n",
    "\n",
    "        APPROACH:\n",
    "        1. Create token embeddings (vocab_size × embed_dim)\n",
    "        2. Create positional encoding (max_seq_len × embed_dim)\n",
    "        3. Build transformer layers using TransformerBlock\n",
    "        4. Add output projection layer\n",
    "        5. Calculate and report parameter count\n",
    "\n",
    "        ARCHITECTURE DECISIONS:\n",
    "        - embed_dim=128: Small enough for fast training, large enough for learning\n",
    "        - num_layers=4: Sufficient depth without excessive memory\n",
    "        - num_heads=4: Multi-head attention without head_dim being too small\n",
    "        - max_seq_len=256: Reasonable context length for character-level modeling\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> model = TinyGPT(vocab_size=50, embed_dim=128, num_layers=4)\n",
    "        >>> print(f\"Parameters: {model.count_parameters():,}\")\n",
    "        Parameters: 1,234,567\n",
    "\n",
    "        HINTS:\n",
    "        - Use Embedding class for token embeddings\n",
    "        - Use PositionalEncoding for position information\n",
    "        - Stack TransformerBlock instances in a list\n",
    "        - Final Linear layer maps embed_dim → vocab_size\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        self.vocab_size = vocab_size\n",
    "        self.embed_dim = embed_dim\n",
    "        self.num_layers = num_layers\n",
    "        self.num_heads = num_heads\n",
    "        self.max_seq_len = max_seq_len\n",
    "        self.dropout = dropout\n",
    "\n",
    "        # Token embeddings: convert token IDs to dense vectors\n",
    "        self.token_embedding = Embedding(vocab_size, embed_dim)\n",
    "\n",
    "        # Positional encoding: add position information\n",
    "        self.positional_encoding = PositionalEncoding(max_seq_len, embed_dim)\n",
    "\n",
    "        # Transformer layers: core processing\n",
    "        self.transformer_blocks = []\n",
    "        for _ in range(num_layers):\n",
    "            block = TransformerBlock(embed_dim, num_heads, mlp_ratio=4.0)\n",
    "            self.transformer_blocks.append(block)\n",
    "\n",
    "        # Output projection: map back to vocabulary\n",
    "        self.output_projection = Linear(embed_dim, vocab_size)\n",
    "\n",
    "        # Dropout for regularization\n",
    "        self.dropout_layer = Dropout(dropout)\n",
    "\n",
    "        # Calculate parameter count for systems analysis\n",
    "        self._param_count = self.count_parameters()\n",
    "        print(f\"🏗️ TinyGPT initialized: {self._param_count:,} parameters\")\n",
    "        print(f\"📐 Architecture: {num_layers}L/{num_heads}H/{embed_dim}D\")\n",
    "        print(f\"💾 Estimated memory: {self._param_count * 4 / 1024 / 1024:.1f}MB\")\n",
    "        ### END SOLUTION\n",
    "\n",
    "def test_unit_tinygpt_init():\n",
    "    \"\"\"🔬 Test TinyGPT initialization and parameter counting.\"\"\"\n",
    "    print(\"🔬 Unit Test: TinyGPT Initialization...\")\n",
    "\n",
    "    # Create a small model for testing\n",
    "    model = TinyGPT(vocab_size=50, embed_dim=64, num_layers=2, num_heads=2, max_seq_len=128)\n",
    "\n",
    "    # Verify architecture components exist\n",
    "    assert hasattr(model, 'token_embedding')\n",
    "    assert hasattr(model, 'positional_encoding')\n",
    "    assert hasattr(model, 'transformer_blocks')\n",
    "    assert hasattr(model, 'output_projection')\n",
    "    assert len(model.transformer_blocks) == 2\n",
    "\n",
    "    # Verify parameter count is reasonable\n",
    "    param_count = model.count_parameters()\n",
    "    assert param_count > 0\n",
    "    assert param_count < 1000000  # Sanity check for small model\n",
    "\n",
    "    print(f\"✅ Model created with {param_count:,} parameters\")\n",
    "    print(\"✅ TinyGPT initialization works correctly!\")\n",
    "\n",
    "# Run immediate test\n",
    "test_unit_tinygpt_init()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba03c6ae",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "tinygpt_methods",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def count_parameters(self) -> int:\n",
    "    \"\"\"\n",
    "    Count total trainable parameters in the model.\n",
    "\n",
    "    TODO: Implement parameter counting across all components\n",
    "\n",
    "    APPROACH:\n",
    "    1. Get parameters from token embeddings\n",
    "    2. Get parameters from all transformer blocks\n",
    "    3. Get parameters from output projection\n",
    "    4. Sum all parameter counts\n",
    "    5. Return total count\n",
    "\n",
    "    SYSTEMS INSIGHT:\n",
    "    Parameter count directly determines:\n",
    "    - Model memory footprint (params × 4 bytes for float32)\n",
    "    - Training memory (3× params for gradients + optimizer states)\n",
    "    - Inference latency (more params = more compute)\n",
    "\n",
    "    EXAMPLE:\n",
    "    >>> model = TinyGPT(vocab_size=1000, embed_dim=128, num_layers=6)\n",
    "    >>> params = model.count_parameters()\n",
    "    >>> print(f\"Memory: {params * 4 / 1024 / 1024:.1f}MB\")\n",
    "    Memory: 52.3MB\n",
    "\n",
    "    HINT: Each component has a parameters() method that returns a list\n",
    "    \"\"\"\n",
    "    ### BEGIN SOLUTION\n",
    "    total_params = 0\n",
    "\n",
    "    # Count embedding parameters\n",
    "    for param in self.token_embedding.parameters():\n",
    "        total_params += np.prod(param.shape)\n",
    "\n",
    "    # Count transformer block parameters\n",
    "    for block in self.transformer_blocks:\n",
    "        for param in block.parameters():\n",
    "            total_params += np.prod(param.shape)\n",
    "\n",
    "    # Count output projection parameters\n",
    "    for param in self.output_projection.parameters():\n",
    "        total_params += np.prod(param.shape)\n",
    "\n",
    "    return total_params\n",
    "    ### END SOLUTION\n",
    "\n",
    "def forward(self, input_ids: Tensor, return_logits: bool = True) -> Tensor:\n",
    "    \"\"\"\n",
    "    Forward pass through the complete TinyGPT model.\n",
    "\n",
    "    TODO: Implement full forward pass integrating all components\n",
    "\n",
    "    APPROACH:\n",
    "    1. Apply token embeddings to convert IDs to vectors\n",
    "    2. Add positional encoding for sequence position information\n",
    "    3. Apply dropout for regularization\n",
    "    4. Pass through each transformer block sequentially\n",
    "    5. Apply final output projection to get logits\n",
    "\n",
    "    ARCHITECTURE FLOW:\n",
    "    input_ids → embeddings → +positional → dropout → transformer_layers → output_proj → logits\n",
    "\n",
    "    EXAMPLE:\n",
    "    >>> model = TinyGPT(vocab_size=100, embed_dim=64)\n",
    "    >>> input_ids = Tensor([[1, 15, 42, 7]])  # Shape: (batch=1, seq_len=4)\n",
    "    >>> logits = model.forward(input_ids)\n",
    "    >>> print(logits.shape)\n",
    "    (1, 4, 100)  # (batch, seq_len, vocab_size)\n",
    "\n",
    "    HINTS:\n",
    "    - embeddings + positional should be element-wise addition\n",
    "    - Each transformer block takes and returns same shape\n",
    "    - Final logits shape: (batch_size, seq_len, vocab_size)\n",
    "    \"\"\"\n",
    "    ### BEGIN SOLUTION\n",
    "    batch_size, seq_len = input_ids.shape\n",
    "\n",
    "    # Step 1: Token embeddings\n",
    "    embeddings = self.token_embedding.forward(input_ids)  # (batch, seq_len, embed_dim)\n",
    "\n",
    "    # Step 2: Add positional encoding\n",
    "    positions = self.positional_encoding.forward(embeddings)  # Same shape\n",
    "    hidden_states = embeddings + positions\n",
    "\n",
    "    # Step 3: Apply dropout\n",
    "    hidden_states = self.dropout_layer.forward(hidden_states, training=True)\n",
    "\n",
    "    # Step 4: Pass through transformer blocks\n",
    "    for block in self.transformer_blocks:\n",
    "        hidden_states = block.forward(hidden_states)\n",
    "\n",
    "    # Step 5: Output projection to vocabulary\n",
    "    if return_logits:\n",
    "        logits = self.output_projection.forward(hidden_states)\n",
    "        return logits  # (batch, seq_len, vocab_size)\n",
    "    else:\n",
    "        return hidden_states  # Return final hidden states\n",
    "    ### END SOLUTION\n",
    "\n",
    "def generate(self, prompt_ids: Tensor, max_new_tokens: int = 50,\n",
    "             temperature: float = 1.0, use_cache: bool = True) -> Tensor:\n",
    "    \"\"\"\n",
    "    Generate text using autoregressive sampling.\n",
    "\n",
    "    TODO: Implement text generation with KV caching optimization\n",
    "\n",
    "    APPROACH:\n",
    "    1. Initialize KV cache if enabled\n",
    "    2. For each new token position:\n",
    "       a. Get logits for next token\n",
    "       b. Apply temperature scaling\n",
    "       c. Sample from probability distribution\n",
    "       d. Append to sequence\n",
    "    3. Return complete generated sequence\n",
    "\n",
    "    SYSTEMS OPTIMIZATION:\n",
    "    - Without cache: O(n²) complexity (recompute all positions)\n",
    "    - With cache: O(n) complexity (only compute new position)\n",
    "    - Cache memory: O(layers × heads × seq_len × head_dim)\n",
    "\n",
    "    EXAMPLE:\n",
    "    >>> model = TinyGPT(vocab_size=100)\n",
    "    >>> prompt = Tensor([[1, 5, 10]])  # \"Hello\"\n",
    "    >>> output = model.generate(prompt, max_new_tokens=10)\n",
    "    >>> print(output.shape)\n",
    "    (1, 13)  # Original 3 + 10 new tokens\n",
    "\n",
    "    HINTS:\n",
    "    - Use KVCache from Module 14 for efficiency\n",
    "    - Apply softmax with temperature for sampling\n",
    "    - Build sequence iteratively, one token at a time\n",
    "    \"\"\"\n",
    "    ### BEGIN SOLUTION\n",
    "    batch_size, current_seq_len = prompt_ids.shape\n",
    "\n",
    "    if use_cache and current_seq_len + max_new_tokens <= self.max_seq_len:\n",
    "        # Initialize KV cache for efficient generation\n",
    "        cache = KVCache(\n",
    "            batch_size=batch_size,\n",
    "            max_seq_len=self.max_seq_len,\n",
    "            num_layers=self.num_layers,\n",
    "            num_heads=self.num_heads,\n",
    "            head_dim=self.embed_dim // self.num_heads\n",
    "        )\n",
    "    else:\n",
    "        cache = None\n",
    "\n",
    "    # Start with the prompt\n",
    "    generated_ids = prompt_ids\n",
    "\n",
    "    for step in range(max_new_tokens):\n",
    "        # Get logits for next token prediction\n",
    "        if cache is not None:\n",
    "            # Efficient: only process the last token\n",
    "            current_input = generated_ids[:, -1:] if step > 0 else generated_ids\n",
    "            logits = self.forward_with_cache(current_input, cache, step)\n",
    "        else:\n",
    "            # Standard: process entire sequence each time\n",
    "            logits = self.forward(generated_ids)\n",
    "\n",
    "        # Get logits for the last position (next token prediction)\n",
    "        next_token_logits = logits[:, -1, :]  # (batch_size, vocab_size)\n",
    "\n",
    "        # Apply temperature scaling\n",
    "        if temperature != 1.0:\n",
    "            next_token_logits = next_token_logits / temperature\n",
    "\n",
    "        # Sample next token (simple greedy for now)\n",
    "        next_token_id = Tensor(np.argmax(next_token_logits.data, axis=-1, keepdims=True))\n",
    "\n",
    "        # Append to sequence\n",
    "        generated_ids = Tensor(np.concatenate([generated_ids.data, next_token_id.data], axis=1))\n",
    "\n",
    "        # Stop if we hit max sequence length\n",
    "        if generated_ids.shape[1] >= self.max_seq_len:\n",
    "            break\n",
    "\n",
    "    return generated_ids\n",
    "    ### END SOLUTION\n",
    "\n",
    "# Add methods to TinyGPT class\n",
    "TinyGPT.count_parameters = count_parameters\n",
    "TinyGPT.forward = forward\n",
    "TinyGPT.generate = generate\n",
    "\n",
    "def test_unit_tinygpt_forward():\n",
    "    \"\"\"🔬 Test TinyGPT forward pass and generation.\"\"\"\n",
    "    print(\"🔬 Unit Test: TinyGPT Forward Pass...\")\n",
    "\n",
    "    # Create model and test data\n",
    "    model = TinyGPT(vocab_size=100, embed_dim=64, num_layers=2, num_heads=2)\n",
    "    input_ids = Tensor([[1, 15, 42, 7, 23]])  # Batch size 1, sequence length 5\n",
    "\n",
    "    # Test forward pass\n",
    "    logits = model.forward(input_ids)\n",
    "\n",
    "    # Verify output shape\n",
    "    expected_shape = (1, 5, 100)  # (batch, seq_len, vocab_size)\n",
    "    assert logits.shape == expected_shape, f\"Expected {expected_shape}, got {logits.shape}\"\n",
    "\n",
    "    # Test generation\n",
    "    prompt = Tensor([[1, 15]])\n",
    "    generated = model.generate(prompt, max_new_tokens=5)\n",
    "\n",
    "    # Verify generation extends sequence\n",
    "    assert generated.shape[1] == 7, f\"Expected 7 tokens, got {generated.shape[1]}\"\n",
    "    assert np.array_equal(generated.data[:, :2], prompt.data), \"Prompt should be preserved\"\n",
    "\n",
    "    print(f\"✅ Forward pass shape: {logits.shape}\")\n",
    "    print(f\"✅ Generation shape: {generated.shape}\")\n",
    "    print(\"✅ TinyGPT forward and generation work correctly!\")\n",
    "\n",
    "# Run immediate test\n",
    "test_unit_tinygpt_forward()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3b6bd45",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 🚀 Stage 2: Training Pipeline Integration\n",
    "\n",
    "Now we'll integrate the training components (Modules 05-07) to create a complete training pipeline. This demonstrates how autograd, optimizers, and training loops work together in a production-quality system.\n",
    "\n",
    "### What We're Building: Complete Training Infrastructure\n",
    "\n",
    "The training pipeline connects data processing, model forward/backward passes, and optimization into a cohesive learning system:\n",
    "\n",
    "```\n",
    "                        🎯 TRAINING PIPELINE ARCHITECTURE 🎯\n",
    "\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                              DATA PREPARATION FLOW                                  │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  Raw Text Corpus                                                                   │\n",
    "│       │                                                                             │\n",
    "│       ▼                                                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Text Processing (Module 10 - Tokenization)                                 │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ \"Hello world\" → [72, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100]    │   │\n",
    "│  │ \"AI is fun\"  → [65, 73, 32, 105, 115, 32, 102, 117, 110]                 │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                       │                                             │\n",
    "│                                       ▼                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Language Modeling Setup                                                     │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Input:   [72, 101, 108, 108, 111]  ←─ Current tokens                       │   │\n",
    "│  │ Target:  [101, 108, 108, 111, 32]  ←─ Next tokens (shifted by 1)          │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Model learns: P(next_token | previous_tokens)                              │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                       │                                             │\n",
    "│                                       ▼                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Batch Formation (Module 08 - DataLoader)                                   │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Sequence 1: [input_ids_1, target_ids_1]                                   │   │\n",
    "│  │ Sequence 2: [input_ids_2, target_ids_2]                                   │   │\n",
    "│  │    ...           ...                                                       │   │\n",
    "│  │ Sequence N: [input_ids_N, target_ids_N]                                   │   │\n",
    "│  │                                     │                                       │   │\n",
    "│  │                                     ▼                                       │   │\n",
    "│  │ Batched Tensor: (batch_size, seq_len) shape                               │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                             TRAINING STEP EXECUTION                                 │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  Training Step Loop (for each batch):                                              │\n",
    "│                                                                                     │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Step 1: Zero Gradients (Module 06 - Optimizers)                            │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ optimizer.zero_grad()  ←─ Clear gradients from previous step               │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Before: param.grad = [0.1, 0.3, -0.2, ...]  ←─ Old gradients              │   │\n",
    "│  │ After:  param.grad = [0.0, 0.0,  0.0, ...]  ←─ Cleared                    │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                       │                                             │\n",
    "│                                       ▼                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Step 2: Forward Pass (Modules 01-04, 11-13)                                │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ input_ids ──► TinyGPT ──► logits (batch, seq_len, vocab_size)             │   │\n",
    "│  │                │                                                           │   │\n",
    "│  │                ▼                                                           │   │\n",
    "│  │ Memory Usage: ~2× model size (activations + parameters)                   │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                       │                                             │\n",
    "│                                       ▼                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Step 3: Loss Computation (Module 04 - Losses)                              │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ logits (batch×seq_len, vocab_size) ──┐                                     │   │\n",
    "│  │                                       │                                     │   │\n",
    "│  │ targets (batch×seq_len,)          ────┼──► CrossEntropyLoss ──► scalar     │   │\n",
    "│  │                                       │                                     │   │\n",
    "│  │ Measures: How well model predicts next tokens                              │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                       │                                             │\n",
    "│                                       ▼                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Step 4: Backward Pass (Module 05 - Autograd)                               │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ loss.backward()  ←─ Automatic differentiation through computation graph    │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Memory Usage: ~3× model size (params + activations + gradients)           │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Result: param.grad = [∂L/∂w₁, ∂L/∂w₂, ∂L/∂w₃, ...]                      │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                       │                                             │\n",
    "│                                       ▼                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Step 5: Parameter Update (Module 06 - Optimizers)                          │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ AdamW Optimizer:                                                            │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ momentum₁ = β₁ × momentum₁ + (1-β₁) × gradient                             │   │\n",
    "│  │ momentum₂ = β₂ × momentum₂ + (1-β₂) × gradient²                            │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ param = param - learning_rate × (momentum₁ / √momentum₂ + weight_decay)    │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Memory Usage: ~4× model size (params + grads + 2×momentum)                │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                               TRAINING MONITORING                                   │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  Training Metrics Tracking:                                                        │\n",
    "│                                                                                     │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ • Loss Tracking: Monitor convergence                                        │   │\n",
    "│  │   - Training loss should decrease over time                                 │   │\n",
    "│  │   - Perplexity = exp(loss) should approach 1.0                            │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ • Learning Rate Scheduling (Module 07):                                    │   │\n",
    "│  │   - Cosine schedule: lr = max_lr × cos(π × epoch / max_epochs)            │   │\n",
    "│  │   - Warm-up: gradually increase lr for first few epochs                    │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ • Memory Monitoring:                                                        │   │\n",
    "│  │   - Track GPU memory usage                                                  │   │\n",
    "│  │   - Detect memory leaks                                                     │   │\n",
    "│  │   - Optimize batch sizes                                                    │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ • Gradient Health:                                                          │   │\n",
    "│  │   - Monitor gradient norms                                                  │   │\n",
    "│  │   - Detect exploding/vanishing gradients                                   │   │\n",
    "│  │   - Apply gradient clipping if needed                                      │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "```\n",
    "\n",
    "### Memory Management During Training\n",
    "\n",
    "Training requires careful memory management due to the multiple copies of model state:\n",
    "\n",
    "```\n",
    "Training Memory Breakdown (TinyGPT-13M example):\n",
    "\n",
    "┌─────────────────────┬─────────────────┬─────────────────┬─────────────────┐\n",
    "│ Component           │ Memory Usage    │ When Allocated  │ Purpose         │\n",
    "├─────────────────────┼─────────────────┼─────────────────┼─────────────────┤\n",
    "│ Model Parameters    │    52 MB        │ Model Init      │ Forward Pass    │\n",
    "│ Gradients          │    52 MB        │ First Backward  │ Store ∂L/∂w     │\n",
    "│ Adam Momentum1     │    52 MB        │ First Step      │ Optimizer State │\n",
    "│ Adam Momentum2     │    52 MB        │ First Step      │ Optimizer State │\n",
    "│ Activations        │    ~100 MB      │ Forward Pass    │ Backward Pass   │\n",
    "├─────────────────────┼─────────────────┼─────────────────┼─────────────────┤\n",
    "│ TOTAL TRAINING     │    ~308 MB      │ Peak Usage      │ All Operations  │\n",
    "├─────────────────────┼─────────────────┼─────────────────┼─────────────────┤\n",
    "│ Inference Only     │    52 MB        │ Model Init      │ Just Forward    │\n",
    "└─────────────────────┴─────────────────┴─────────────────┴─────────────────┘\n",
    "\n",
    "Key Insights:\n",
    "• Training uses ~6× inference memory\n",
    "• Adam optimizer doubles memory (2 momentum terms)\n",
    "• Activation memory scales with batch size and sequence length\n",
    "• Gradient checkpointing can reduce activation memory\n",
    "```\n",
    "\n",
    "### Systems Focus: Training Performance Optimization\n",
    "\n",
    "**1. Memory Management**: Keep training within GPU memory limits\n",
    "**2. Convergence Monitoring**: Track loss, perplexity, and gradient health\n",
    "**3. Learning Rate Scheduling**: Optimize training dynamics\n",
    "**4. Checkpointing**: Save model state for recovery and deployment\n",
    "\n",
    "Let's implement the complete training infrastructure that makes all of this work seamlessly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87cb0d2f",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "training_pipeline",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "class TinyGPTTrainer:\n",
    "    \"\"\"\n",
    "    Complete training pipeline integrating optimizers, schedulers, and monitoring.\n",
    "\n",
    "    Uses modules 05 (autograd), 06 (optimizers), 07 (training) for end-to-end training.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, model: TinyGPT, tokenizer: CharTokenizer,\n",
    "                 learning_rate: float = 3e-4, weight_decay: float = 0.01):\n",
    "        \"\"\"\n",
    "        Initialize trainer with model and optimization components.\n",
    "\n",
    "        TODO: Set up complete training infrastructure\n",
    "\n",
    "        APPROACH:\n",
    "        1. Store model and tokenizer references\n",
    "        2. Initialize AdamW optimizer (standard for transformers)\n",
    "        3. Initialize loss function (CrossEntropyLoss for language modeling)\n",
    "        4. Set up learning rate scheduler (cosine schedule)\n",
    "        5. Initialize training metrics tracking\n",
    "\n",
    "        PRODUCTION CHOICES:\n",
    "        - AdamW: Better generalization than Adam (weight decay)\n",
    "        - learning_rate=3e-4: Standard for small transformers\n",
    "        - Cosine schedule: Smooth learning rate decay\n",
    "        - CrossEntropy: Standard for classification/language modeling\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> model = TinyGPT(vocab_size=100)\n",
    "        >>> tokenizer = CharTokenizer(['a', 'b', 'c'])\n",
    "        >>> trainer = TinyGPTTrainer(model, tokenizer)\n",
    "        >>> print(\"Trainer ready for training\")\n",
    "        Trainer ready for training\n",
    "\n",
    "        HINTS:\n",
    "        - Get all model parameters with model.parameters()\n",
    "        - Use AdamW with weight_decay for better generalization\n",
    "        - CrossEntropyLoss handles the language modeling objective\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        self.model = model\n",
    "        self.tokenizer = tokenizer\n",
    "\n",
    "        # Collect all trainable parameters\n",
    "        all_params = []\n",
    "        all_params.extend(model.token_embedding.parameters())\n",
    "        for block in model.transformer_blocks:\n",
    "            all_params.extend(block.parameters())\n",
    "        all_params.extend(model.output_projection.parameters())\n",
    "\n",
    "        # Initialize optimizer (AdamW for transformers)\n",
    "        self.optimizer = AdamW(\n",
    "            params=all_params,\n",
    "            lr=learning_rate,\n",
    "            weight_decay=weight_decay,\n",
    "            betas=(0.9, 0.95)  # Standard for language models\n",
    "        )\n",
    "\n",
    "        # Loss function for next token prediction\n",
    "        self.loss_fn = CrossEntropyLoss()\n",
    "\n",
    "        # Learning rate scheduler\n",
    "        self.scheduler = CosineSchedule(\n",
    "            optimizer=self.optimizer,\n",
    "            max_epochs=100,  # Will adjust based on actual training\n",
    "            min_lr=learning_rate * 0.1\n",
    "        )\n",
    "\n",
    "        # Training metrics\n",
    "        self.training_history = {\n",
    "            'losses': [],\n",
    "            'perplexities': [],\n",
    "            'learning_rates': [],\n",
    "            'epoch': 0\n",
    "        }\n",
    "\n",
    "        print(f\"🚀 Trainer initialized:\")\n",
    "        print(f\"   Optimizer: AdamW (lr={learning_rate}, wd={weight_decay})\")\n",
    "        print(f\"   Parameters: {len(all_params):,} tensors\")\n",
    "        print(f\"   Loss: CrossEntropyLoss\")\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def prepare_batch(self, text_batch: List[str], max_length: int = 128) -> Tuple[Tensor, Tensor]:\n",
    "        \"\"\"\n",
    "        Convert text batch to input/target tensors for language modeling.\n",
    "\n",
    "        TODO: Implement text-to-tensor conversion with proper targets\n",
    "\n",
    "        APPROACH:\n",
    "        1. Tokenize each text in the batch\n",
    "        2. Pad/truncate to consistent length\n",
    "        3. Create input_ids (text) and target_ids (text shifted by 1)\n",
    "        4. Convert to Tensor format\n",
    "\n",
    "        LANGUAGE MODELING OBJECTIVE:\n",
    "        - Input: [token1, token2, token3, token4]\n",
    "        - Target: [token2, token3, token4, token5]\n",
    "        - Model predicts next token at each position\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> trainer = TinyGPTTrainer(model, tokenizer)\n",
    "        >>> texts = [\"hello world\", \"ai is fun\"]\n",
    "        >>> inputs, targets = trainer.prepare_batch(texts)\n",
    "        >>> print(inputs.shape, targets.shape)\n",
    "        (2, 128) (2, 128)\n",
    "\n",
    "        HINTS:\n",
    "        - Use tokenizer.encode() for text → token conversion\n",
    "        - Pad shorter sequences with tokenizer pad token\n",
    "        - Target sequence is input sequence shifted right by 1\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        batch_size = len(text_batch)\n",
    "\n",
    "        # Tokenize all texts\n",
    "        tokenized_batch = []\n",
    "        for text in text_batch:\n",
    "            tokens = self.tokenizer.encode(text)\n",
    "\n",
    "            # Truncate or pad to max_length\n",
    "            if len(tokens) > max_length:\n",
    "                tokens = tokens[:max_length]\n",
    "            else:\n",
    "                # Pad with special token (use 0 as pad)\n",
    "                tokens.extend([0] * (max_length - len(tokens)))\n",
    "\n",
    "            tokenized_batch.append(tokens)\n",
    "\n",
    "        # Convert to numpy then Tensor\n",
    "        input_ids = Tensor(np.array(tokenized_batch))  # (batch_size, seq_len)\n",
    "\n",
    "        # Create targets (shifted input for next token prediction)\n",
    "        target_ids = Tensor(np.roll(input_ids.data, -1, axis=1))  # Shift left by 1\n",
    "\n",
    "        return input_ids, target_ids\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def train_step(self, input_ids: Tensor, target_ids: Tensor) -> float:\n",
    "        \"\"\"\n",
    "        Single training step with forward, backward, and optimization.\n",
    "\n",
    "        TODO: Implement complete training step\n",
    "\n",
    "        APPROACH:\n",
    "        1. Zero gradients from previous step\n",
    "        2. Forward pass to get logits\n",
    "        3. Compute loss between logits and targets\n",
    "        4. Backward pass to compute gradients\n",
    "        5. Optimizer step to update parameters\n",
    "        6. Return loss value for monitoring\n",
    "\n",
    "        MEMORY MANAGEMENT:\n",
    "        During training, memory usage = 3× model size:\n",
    "        - 1× for parameters\n",
    "        - 1× for gradients\n",
    "        - 1× for optimizer states (Adam moments)\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> loss = trainer.train_step(input_ids, target_ids)\n",
    "        >>> print(f\"Training loss: {loss:.4f}\")\n",
    "        Training loss: 2.3456\n",
    "\n",
    "        HINTS:\n",
    "        - Always zero_grad() before forward pass\n",
    "        - Loss should be computed on flattened logits and targets\n",
    "        - Call backward() on the loss tensor\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # Zero gradients from previous step\n",
    "        self.optimizer.zero_grad()\n",
    "\n",
    "        # Forward pass\n",
    "        logits = self.model.forward(input_ids)  # (batch, seq_len, vocab_size)\n",
    "\n",
    "        # Reshape for loss computation\n",
    "        batch_size, seq_len, vocab_size = logits.shape\n",
    "        logits_flat = logits.reshape(batch_size * seq_len, vocab_size)\n",
    "        targets_flat = target_ids.reshape(batch_size * seq_len)\n",
    "\n",
    "        # Compute loss\n",
    "        loss = self.loss_fn.forward(logits_flat, targets_flat)\n",
    "\n",
    "        # Backward pass\n",
    "        loss.backward()\n",
    "\n",
    "        # Optimizer step\n",
    "        self.optimizer.step()\n",
    "\n",
    "        # Return scalar loss for monitoring\n",
    "        return float(loss.data.item() if hasattr(loss.data, 'item') else loss.data)\n",
    "        ### END SOLUTION\n",
    "\n",
    "def test_unit_training_pipeline():\n",
    "    \"\"\"🔬 Test training pipeline components.\"\"\"\n",
    "    print(\"🔬 Unit Test: Training Pipeline...\")\n",
    "\n",
    "    # Create small model and trainer\n",
    "    model = TinyGPT(vocab_size=50, embed_dim=32, num_layers=2, num_heads=2)\n",
    "    tokenizer = CharTokenizer(['a', 'b', 'c', 'd', 'e', ' '])\n",
    "    trainer = TinyGPTTrainer(model, tokenizer, learning_rate=1e-3)\n",
    "\n",
    "    # Test batch preparation\n",
    "    texts = [\"hello\", \"world\"]\n",
    "    input_ids, target_ids = trainer.prepare_batch(texts, max_length=8)\n",
    "\n",
    "    assert input_ids.shape == (2, 8), f\"Expected (2, 8), got {input_ids.shape}\"\n",
    "    assert target_ids.shape == (2, 8), f\"Expected (2, 8), got {target_ids.shape}\"\n",
    "\n",
    "    # Test training step\n",
    "    initial_loss = trainer.train_step(input_ids, target_ids)\n",
    "    assert initial_loss > 0, \"Loss should be positive\"\n",
    "\n",
    "    # Second step should work (gradients computed and applied)\n",
    "    second_loss = trainer.train_step(input_ids, target_ids)\n",
    "    assert second_loss > 0, \"Second loss should also be positive\"\n",
    "\n",
    "    print(f\"✅ Batch preparation shape: {input_ids.shape}\")\n",
    "    print(f\"✅ Initial loss: {initial_loss:.4f}\")\n",
    "    print(f\"✅ Second loss: {second_loss:.4f}\")\n",
    "    print(\"✅ Training pipeline works correctly!\")\n",
    "\n",
    "# Run immediate test\n",
    "test_unit_training_pipeline()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e740071a",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## ⚡ Stage 3: Systems Analysis and Optimization\n",
    "\n",
    "Now we'll apply the systems analysis tools from Modules 15-19 to understand TinyGPT's performance characteristics. This demonstrates the complete systems thinking approach to ML engineering.\n",
    "\n",
    "### What We're Analyzing: Complete Performance Profile\n",
    "\n",
    "Real ML systems require deep understanding of performance characteristics, bottlenecks, and optimization opportunities. Let's systematically analyze TinyGPT across all dimensions:\n",
    "\n",
    "```\n",
    "                         📊 SYSTEMS ANALYSIS FRAMEWORK 📊\n",
    "\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                             1. BASELINE PROFILING                                   │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  Parameter Analysis (Module 15):                                                   │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Count & Distribution  →  Memory Footprint  →  FLOP Analysis                │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Where are params?     What's the memory?   How many operations?            │   │\n",
    "│  │ • Embeddings: 15%     • Inference: 1×     • Attention: O(n²×d)            │   │\n",
    "│  │ • Attention: 31%      • Training: 3×      • MLP: O(n×d²)                  │   │\n",
    "│  │ • MLP: 46%           • Optim: 4×          • Total: O(L×n×d²)              │   │\n",
    "│  │ • Other: 8%                                                                │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                          2. SCALING BEHAVIOR ANALYSIS                               │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  How does performance scale with key parameters?                                   │\n",
    "│                                                                                     │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Model Size Scaling:                                                         │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ embed_dim:  64  →  128  →  256  →  512                                     │   │\n",
    "│  │ Memory:     5MB →  20MB →  80MB →  320MB                                   │   │\n",
    "│  │ Inference:  10ms→  25ms →  60ms →  150ms                                   │   │\n",
    "│  │ Training:   30ms→  75ms → 180ms →  450ms                                   │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Memory scales as O(d²), Compute scales as O(d³)                           │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                                                                     │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Sequence Length Scaling:                                                    │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ seq_len:     64   →   128  →   256  →   512                                │   │\n",
    "│  │ Attn Memory: 16KB →   64KB →  256KB → 1024KB                               │   │\n",
    "│  │ Attn Time:   2ms  →    8ms →   32ms →  128ms                               │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Attention is the quadratic bottleneck: O(n²)                              │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                                                                     │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Batch Size Scaling:                                                         │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ batch_size:  1   →    4   →   16   →   32                                  │   │\n",
    "│  │ Memory:     50MB →  200MB →  800MB → 1600MB                                │   │\n",
    "│  │ Throughput: 100  →  350   → 1200   → 2000  tokens/sec                     │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Linear memory growth, sub-linear throughput improvement                    │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                           3. OPTIMIZATION IMPACT ANALYSIS                           │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  Quantization Analysis (Module 17):                                                │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │                    QUANTIZATION PIPELINE                                   │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ FP32 Model     →    INT8 Conversion    →    Performance Impact             │   │\n",
    "│  │ (32-bit)            (8-bit)                                                │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ 200MB          →         50MB          →    4× memory reduction           │   │\n",
    "│  │ 100ms inference →       60ms inference  →    1.7× speedup                │   │\n",
    "│  │ 95.2% accuracy  →      94.8% accuracy   →    0.4% accuracy loss           │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Trade-off: 4× smaller, 1.7× faster, minimal accuracy loss                │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                                                                     │\n",
    "│  Pruning Analysis (Module 18):                                                     │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │                      PRUNING PIPELINE                                      │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Dense Model → Magnitude Pruning → Structured Pruning → Performance        │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Sparsity:  0%     →      50%     →       90%        →   Impact           │   │\n",
    "│  │ Memory:   200MB   →     100MB     →      20MB        →   10× reduction   │   │\n",
    "│  │ Speed:    100ms   →      80ms     →      40ms        →   2.5× speedup    │   │\n",
    "│  │ Accuracy: 95.2%   →     94.8%     →     92.1%        →   3.1% loss       │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Sweet spot: 70-80% sparsity (good speed/accuracy trade-off)               │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                                                                     │\n",
    "│  Combined Optimization:                                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Original Model: 200MB, 100ms, 95.2% accuracy                              │   │\n",
    "│  │      ↓                                                                      │   │\n",
    "│  │ + INT8 Quantization: 50MB, 60ms, 94.8% accuracy                           │   │\n",
    "│  │      ↓                                                                      │   │\n",
    "│  │ + 80% Pruning: 10MB, 30ms, 92.5% accuracy                                 │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Final: 20× smaller, 3.3× faster, 2.7% accuracy loss                      │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                         4. COMPARATIVE BENCHMARKING                                 │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  Benchmark Against Reference Implementations (Module 19):                          │\n",
    "│                                                                                     │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │                        BENCHMARK RESULTS                                   │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ ┌─────────────┬─────────────┬─────────────┬─────────────┬─────────────┐   │   │\n",
    "│  │ │   Model     │  Parameters │    Memory   │  Latency    │  Perplexity │   │   │\n",
    "│  │ ├─────────────┼─────────────┼─────────────┼─────────────┼─────────────┤   │   │\n",
    "│  │ │ TinyGPT-1M  │     1M      │    4MB      │    5ms      │    12.5     │   │   │\n",
    "│  │ │ TinyGPT-13M │    13M      │   52MB      │   25ms      │     8.2     │   │   │\n",
    "│  │ │ TinyGPT-50M │    50M      │  200MB      │   80ms      │     6.1     │   │   │\n",
    "│  │ │ GPT-2 Small │   124M      │  500MB      │  150ms      │     5.8     │   │   │\n",
    "│  │ └─────────────┴─────────────┴─────────────┴─────────────┴─────────────┘   │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Key Findings:                                                               │   │\n",
    "│  │ • TinyGPT achieves competitive perplexity at smaller sizes                 │   │\n",
    "│  │ • Linear scaling relationship between params and performance               │   │\n",
    "│  │ • Memory efficiency matches theoretical predictions                        │   │\n",
    "│  │ • Inference latency scales predictably with model size                    │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "```\n",
    "\n",
    "### Critical Performance Insights\n",
    "\n",
    "**Scaling Laws:**\n",
    "- **Parameters**: Memory ∝ params, Compute ∝ params^1.3\n",
    "- **Sequence Length**: Attention memory/compute ∝ seq_len²\n",
    "- **Model Depth**: Memory ∝ layers, Compute ∝ layers\n",
    "\n",
    "**Optimization Sweet Spots:**\n",
    "- **Quantization**: 4× memory reduction, <5% accuracy loss\n",
    "- **Pruning**: 70-80% sparsity optimal for accuracy/speed trade-off\n",
    "- **Combined**: 20× total compression possible with careful tuning\n",
    "\n",
    "**Bottleneck Analysis:**\n",
    "- **Training**: Memory bandwidth (moving gradients)\n",
    "- **Inference**: Compute bound (matrix multiplications)\n",
    "- **Generation**: Sequential dependency (limited parallelism)\n",
    "\n",
    "Let's implement comprehensive analysis functions that measure and understand all these characteristics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77272cce",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "systems_analysis",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def analyze_tinygpt_memory_scaling():\n",
    "    \"\"\"📊 Analyze how TinyGPT memory usage scales with model size.\"\"\"\n",
    "    print(\"📊 Analyzing TinyGPT Memory Scaling...\")\n",
    "\n",
    "    configs = [\n",
    "        {\"embed_dim\": 64, \"num_layers\": 2, \"name\": \"Tiny\"},\n",
    "        {\"embed_dim\": 128, \"num_layers\": 4, \"name\": \"Small\"},\n",
    "        {\"embed_dim\": 256, \"num_layers\": 6, \"name\": \"Base\"},\n",
    "        {\"embed_dim\": 512, \"num_layers\": 8, \"name\": \"Large\"}\n",
    "    ]\n",
    "\n",
    "    results = []\n",
    "    for config in configs:\n",
    "        model = TinyGPT(\n",
    "            vocab_size=1000,\n",
    "            embed_dim=config[\"embed_dim\"],\n",
    "            num_layers=config[\"num_layers\"],\n",
    "            num_heads=config[\"embed_dim\"] // 32,  # Maintain reasonable head_dim\n",
    "            max_seq_len=256\n",
    "        )\n",
    "\n",
    "        # Use Module 15 profiler\n",
    "        profiler = Profiler()\n",
    "        param_count = profiler.count_parameters(model)\n",
    "\n",
    "        # Calculate memory footprint\n",
    "        inference_memory = param_count * 4 / (1024 * 1024)  # MB\n",
    "        training_memory = inference_memory * 3  # Parameters + gradients + optimizer\n",
    "\n",
    "        results.append({\n",
    "            \"name\": config[\"name\"],\n",
    "            \"params\": param_count,\n",
    "            \"inference_mb\": inference_memory,\n",
    "            \"training_mb\": training_memory,\n",
    "            \"embed_dim\": config[\"embed_dim\"],\n",
    "            \"layers\": config[\"num_layers\"]\n",
    "        })\n",
    "\n",
    "        print(f\"{config['name']}: {param_count:,} params, \"\n",
    "              f\"Inference: {inference_memory:.1f}MB, Training: {training_memory:.1f}MB\")\n",
    "\n",
    "    # Analyze scaling trends\n",
    "    print(\"\\n💡 Memory Scaling Insights:\")\n",
    "    tiny_params = results[0][\"params\"]\n",
    "    large_params = results[-1][\"params\"]\n",
    "    scaling_factor = large_params / tiny_params\n",
    "    print(f\"   Parameter growth: {scaling_factor:.1f}× from Tiny to Large\")\n",
    "    print(f\"   Training memory range: {results[0]['training_mb']:.1f}MB → {results[-1]['training_mb']:.1f}MB\")\n",
    "\n",
    "    return results\n",
    "\n",
    "def analyze_optimization_impact():\n",
    "    \"\"\"📊 Analyze the impact of quantization and pruning on model performance.\"\"\"\n",
    "    print(\"📊 Analyzing Optimization Techniques Impact...\")\n",
    "\n",
    "    # Create base model\n",
    "    model = TinyGPT(vocab_size=100, embed_dim=128, num_layers=4, num_heads=4)\n",
    "    profiler = Profiler()\n",
    "\n",
    "    # Baseline measurements\n",
    "    base_params = profiler.count_parameters(model)\n",
    "    base_memory = base_params * 4 / (1024 * 1024)\n",
    "\n",
    "    print(f\"📐 Baseline Model:\")\n",
    "    print(f\"   Parameters: {base_params:,}\")\n",
    "    print(f\"   Memory: {base_memory:.1f}MB\")\n",
    "\n",
    "    # Simulate quantization impact (Module 17)\n",
    "    print(f\"\\n🔧 After INT8 Quantization:\")\n",
    "    quantized_memory = base_memory / 4  # INT8 = 1 byte vs FP32 = 4 bytes\n",
    "    print(f\"   Memory: {quantized_memory:.1f}MB ({quantized_memory/base_memory:.1%} of original)\")\n",
    "    print(f\"   Memory saved: {base_memory - quantized_memory:.1f}MB\")\n",
    "\n",
    "    # Simulate pruning impact (Module 18)\n",
    "    sparsity_levels = [0.5, 0.7, 0.9]\n",
    "    print(f\"\\n✂️ Pruning Analysis:\")\n",
    "    for sparsity in sparsity_levels:\n",
    "        effective_params = base_params * (1 - sparsity)\n",
    "        memory_reduction = base_memory * sparsity\n",
    "        print(f\"   {sparsity:.0%} sparsity: {effective_params:,} active params, \"\n",
    "              f\"{memory_reduction:.1f}MB saved\")\n",
    "\n",
    "    # Combined optimization\n",
    "    print(f\"\\n🚀 Combined Optimization (90% pruning + INT8):\")\n",
    "    combined_memory = base_memory * 0.1 / 4  # 10% params × 1/4 size\n",
    "    print(f\"   Memory: {combined_memory:.1f}MB ({combined_memory/base_memory:.1%} of original)\")\n",
    "    print(f\"   Total reduction: {base_memory/combined_memory:.1f}× smaller\")\n",
    "\n",
    "def analyze_training_performance():\n",
    "    \"\"\"📊 Analyze training vs inference performance characteristics.\"\"\"\n",
    "    print(\"📊 Analyzing Training vs Inference Performance...\")\n",
    "\n",
    "    # Create model for analysis\n",
    "    model = TinyGPT(vocab_size=1000, embed_dim=256, num_layers=6, num_heads=8)\n",
    "    profiler = Profiler()\n",
    "\n",
    "    # Simulate batch processing at different sizes\n",
    "    batch_sizes = [1, 4, 16, 32]\n",
    "    seq_len = 128\n",
    "\n",
    "    print(f\"📈 Batch Size Impact (seq_len={seq_len}):\")\n",
    "    for batch_size in batch_sizes:\n",
    "        # Calculate memory for batch\n",
    "        input_memory = batch_size * seq_len * 4 / (1024 * 1024)  # Input tokens\n",
    "        activation_memory = input_memory * model.num_layers * 2  # Rough estimate\n",
    "        total_memory = model._param_count * 4 / (1024 * 1024) + activation_memory\n",
    "\n",
    "        # Estimate throughput (tokens/second)\n",
    "        # Rough approximation based on batch efficiency\n",
    "        base_throughput = 100  # tokens/second for batch_size=1\n",
    "        efficiency = min(batch_size, 16) / 16  # Efficiency plateaus at batch_size=16\n",
    "        throughput = base_throughput * batch_size * efficiency\n",
    "\n",
    "        print(f\"   Batch {batch_size:2d}: {total_memory:6.1f}MB memory, \"\n",
    "              f\"{throughput:5.0f} tokens/sec\")\n",
    "\n",
    "    print(\"\\n💡 Performance Insights:\")\n",
    "    print(\"   Memory scales linearly with batch size\")\n",
    "    print(\"   Throughput improves with batching (better GPU utilization)\")\n",
    "    print(\"   Sweet spot: batch_size=16-32 for most GPUs\")\n",
    "\n",
    "# Run all analyses\n",
    "memory_results = analyze_tinygpt_memory_scaling()\n",
    "analyze_optimization_impact()\n",
    "analyze_training_performance()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae6107ae",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 🎭 Stage 4: Complete ML Pipeline Demonstration\n",
    "\n",
    "Now we'll create a complete demonstration that brings together all components into a working ML system. This shows the full journey from raw text to trained model to generated output, demonstrating how all 19 modules work together.\n",
    "\n",
    "### What We're Demonstrating: End-to-End ML System\n",
    "\n",
    "This final stage shows how everything integrates into a production-quality ML pipeline:\n",
    "\n",
    "```\n",
    "                      🎭 COMPLETE ML PIPELINE DEMONSTRATION 🎭\n",
    "\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                           STAGE 1: DATA PREPARATION                                 │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  Raw Text Corpus ──────────────────────────────────────────────────────────────►   │\n",
    "│                                                                                     │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ \"The quick brown fox jumps over the lazy dog.\"                             │   │\n",
    "│  │ \"Artificial intelligence is transforming the world.\"                       │   │\n",
    "│  │ \"Machine learning models require large amounts of data.\"                   │   │\n",
    "│  │ \"Neural networks learn patterns from training examples.\"                   │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                       │                                             │\n",
    "│                                       ▼                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Tokenization (Module 10)                                                    │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ \"The quick\" → [84, 104, 101, 32, 113, 117, 105, 99, 107]                  │   │\n",
    "│  │ \"brown fox\" → [98, 114, 111, 119, 110, 32, 102, 111, 120]                 │   │\n",
    "│  │ ...                                                                         │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Result: 10,000 training sequences                                           │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                       │                                             │\n",
    "│                                       ▼                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ DataLoader Creation (Module 08)                                             │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ • Batch size: 32                                                            │   │\n",
    "│  │ • Sequence length: 64                                                       │   │\n",
    "│  │ • Shuffle: True                                                             │   │\n",
    "│  │ • Total batches: 312                                                        │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                            STAGE 2: MODEL TRAINING                                  │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  Training Configuration:                                                            │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Model: TinyGPT (13M parameters)                                             │   │\n",
    "│  │ • embed_dim: 256                                                            │   │\n",
    "│  │ • num_layers: 6                                                             │   │\n",
    "│  │ • num_heads: 8                                                              │   │\n",
    "│  │ • vocab_size: 1000                                                          │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Optimizer: AdamW                                                            │   │\n",
    "│  │ • learning_rate: 3e-4                                                       │   │\n",
    "│  │ • weight_decay: 0.01                                                        │   │\n",
    "│  │ • betas: (0.9, 0.95)                                                        │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Schedule: Cosine with warmup                                                │   │\n",
    "│  │ • warmup_steps: 100                                                         │   │\n",
    "│  │ • max_epochs: 20                                                            │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                       │                                             │\n",
    "│                                       ▼                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Training Progress:                                                          │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Epoch 1:  Loss=4.234, PPL=68.9   ←─ Random initialization                 │   │\n",
    "│  │ Epoch 5:  Loss=2.891, PPL=18.0   ←─ Learning patterns                     │   │\n",
    "│  │ Epoch 10: Loss=2.245, PPL=9.4    ←─ Convergence                           │   │\n",
    "│  │ Epoch 15: Loss=1.967, PPL=7.1    ←─ Fine-tuning                           │   │\n",
    "│  │ Epoch 20: Loss=1.823, PPL=6.2    ←─ Final performance                     │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Training Time: 45 minutes on CPU                                           │   │\n",
    "│  │ Memory Usage: ~500MB peak                                                   │   │\n",
    "│  │ Final Perplexity: 6.2 (good for character-level)                          │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                           STAGE 3: MODEL OPTIMIZATION                               │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  Optimization Pipeline:                                                             │\n",
    "│                                                                                     │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Step 1: Baseline Profiling (Module 15)                                     │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ • Parameter count: 13,042,176                                               │   │\n",
    "│  │ • Memory footprint: 52.2MB                                                  │   │\n",
    "│  │ • Inference latency: 25ms per sequence                                      │   │\n",
    "│  │ • FLOP count: 847M per forward pass                                         │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                       │                                             │\n",
    "│                                       ▼                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Step 2: INT8 Quantization (Module 17)                                      │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Before: FP32 weights, 52.2MB                                               │   │\n",
    "│  │ After:  INT8 weights, 13.1MB                                               │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ • Memory reduction: 4.0× smaller                                           │   │\n",
    "│  │ • Speed improvement: 1.8× faster                                           │   │\n",
    "│  │ • Accuracy impact: 6.2 → 6.4 PPL (minimal degradation)                   │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                       │                                             │\n",
    "│                                       ▼                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Step 3: Magnitude Pruning (Module 18)                                      │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Sparsity levels tested: 50%, 70%, 90%                                      │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ 50% sparse: 6.5MB, 1.6× faster, 6.3 PPL                                  │   │\n",
    "│  │ 70% sparse: 3.9MB, 2.1× faster, 6.8 PPL                                  │   │\n",
    "│  │ 90% sparse: 1.3MB, 2.8× faster, 8.9 PPL ←─ Too aggressive                │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Optimal: 70% sparsity (good speed/accuracy trade-off)                     │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                       │                                             │\n",
    "│                                       ▼                                             │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Step 4: Final Optimized Model                                               │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Original:  52.2MB, 25ms, 6.2 PPL                                          │   │\n",
    "│  │ Optimized: 3.9MB, 12ms, 6.8 PPL                                           │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Total improvement: 13.4× smaller, 2.1× faster, +0.6 PPL                  │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Ready for deployment on mobile/edge devices!                               │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "┌─────────────────────────────────────────────────────────────────────────────────────┐\n",
    "│                            STAGE 4: TEXT GENERATION                                 │\n",
    "├─────────────────────────────────────────────────────────────────────────────────────┤\n",
    "│                                                                                     │\n",
    "│  Generation Examples:                                                               │\n",
    "│                                                                                     │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ Prompt: \"The future of AI\"                                                 │   │\n",
    "│  │ Generated: \"The future of AI is bright and full of possibilities for       │   │\n",
    "│  │            helping humanity solve complex problems.\"                       │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Prompt: \"Machine learning\"                                                 │   │\n",
    "│  │ Generated: \"Machine learning enables computers to learn patterns from      │   │\n",
    "│  │            data without being explicitly programmed.\"                      │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ Prompt: \"Neural networks\"                                                  │   │\n",
    "│  │ Generated: \"Neural networks are computational models inspired by the       │   │\n",
    "│  │            human brain that can learn complex representations.\"            │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "│                                                                                     │\n",
    "│  Generation Performance:                                                            │\n",
    "│  ┌─────────────────────────────────────────────────────────────────────────────┐   │\n",
    "│  │ • Speed: ~50 tokens/second                                                  │   │\n",
    "│  │ • Quality: Coherent short text                                              │   │\n",
    "│  │ • Memory: 3.9MB (optimized model)                                          │   │\n",
    "│  │ • Latency: 20ms per token                                                   │   │\n",
    "│  │                                                                             │   │\n",
    "│  │ With KV Caching (Module 14):                                               │   │\n",
    "│  │ • Speed: ~80 tokens/second (1.6× improvement)                              │   │\n",
    "│  │ • Memory: +2MB for cache                                                    │   │\n",
    "│  │ • Latency: 12ms per token                                                   │   │\n",
    "│  └─────────────────────────────────────────────────────────────────────────────┘   │\n",
    "└─────────────────────────────────────────────────────────────────────────────────────┘\n",
    "```\n",
    "\n",
    "### Complete System Validation\n",
    "\n",
    "Our end-to-end pipeline demonstrates:\n",
    "\n",
    "**1. Data Flow Integrity**: Text → Tokens → Batches → Training → Model\n",
    "**2. Training Effectiveness**: Loss convergence, perplexity improvement\n",
    "**3. Optimization Success**: Memory reduction, speed improvement\n",
    "**4. Generation Quality**: Coherent text output\n",
    "**5. Systems Integration**: All 19 modules working together\n",
    "\n",
    "Let's implement the complete pipeline class that orchestrates this entire process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4174fb9b",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "complete_pipeline",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "class CompleteTinyGPTPipeline:\n",
    "    \"\"\"\n",
    "    End-to-end ML pipeline demonstrating integration of all 19 modules.\n",
    "\n",
    "    Pipeline stages:\n",
    "    1. Data preparation (Module 10: Tokenization)\n",
    "    2. Model creation (Modules 01-04, 11-13: Architecture)\n",
    "    3. Training setup (Modules 05-07: Optimization)\n",
    "    4. Training loop (Module 08: DataLoader)\n",
    "    5. Optimization (Modules 17-18: Quantization, Pruning)\n",
    "    6. Evaluation (Module 19: Benchmarking)\n",
    "    7. Generation (Module 14: KV Caching)\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_size: int = 100, embed_dim: int = 128,\n",
    "                 num_layers: int = 4, num_heads: int = 4):\n",
    "        \"\"\"Initialize complete pipeline with model architecture.\"\"\"\n",
    "\n",
    "        ### BEGIN SOLUTION\n",
    "        self.vocab_size = vocab_size\n",
    "        self.embed_dim = embed_dim\n",
    "        self.num_layers = num_layers\n",
    "        self.num_heads = num_heads\n",
    "\n",
    "        # Stage 1: Initialize tokenizer (Module 10)\n",
    "        self.tokenizer = CharTokenizer([chr(i) for i in range(32, 127)])  # Printable ASCII\n",
    "\n",
    "        # Stage 2: Create model (Modules 01-04, 11-13)\n",
    "        self.model = TinyGPT(\n",
    "            vocab_size=vocab_size,\n",
    "            embed_dim=embed_dim,\n",
    "            num_layers=num_layers,\n",
    "            num_heads=num_heads,\n",
    "            max_seq_len=256\n",
    "        )\n",
    "\n",
    "        # Stage 3: Setup training (Modules 05-07)\n",
    "        self.trainer = TinyGPTTrainer(self.model, self.tokenizer, learning_rate=3e-4)\n",
    "\n",
    "        # Stage 4: Initialize profiler and benchmark (Modules 15, 19)\n",
    "        self.profiler = Profiler()\n",
    "        self.benchmark = Benchmark([self.model], [], [\"perplexity\", \"latency\"])\n",
    "\n",
    "        # Pipeline state\n",
    "        self.is_trained = False\n",
    "        self.training_history = []\n",
    "\n",
    "        print(\"🏗️ Complete TinyGPT Pipeline Initialized\")\n",
    "        print(f\"   Model: {self.model.count_parameters():,} parameters\")\n",
    "        print(f\"   Memory: {self.model.count_parameters() * 4 / 1024 / 1024:.1f}MB\")\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def prepare_training_data(self, text_corpus: List[str], batch_size: int = 8) -> DataLoader:\n",
    "        \"\"\"\n",
    "        Prepare training data using DataLoader (Module 08).\n",
    "\n",
    "        TODO: Create DataLoader for training text data\n",
    "\n",
    "        APPROACH:\n",
    "        1. Tokenize all texts in corpus\n",
    "        2. Create input/target pairs for language modeling\n",
    "        3. Package into TensorDataset\n",
    "        4. Create DataLoader with batching and shuffling\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> pipeline = CompleteTinyGPTPipeline()\n",
    "        >>> corpus = [\"hello world\", \"ai is amazing\"]\n",
    "        >>> dataloader = pipeline.prepare_training_data(corpus, batch_size=2)\n",
    "        >>> print(f\"Batches: {len(dataloader)}\")\n",
    "        Batches: 1\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # Tokenize and prepare training pairs\n",
    "        input_sequences = []\n",
    "        target_sequences = []\n",
    "\n",
    "        for text in text_corpus:\n",
    "            tokens = self.tokenizer.encode(text)\n",
    "            if len(tokens) < 2:\n",
    "                continue  # Skip very short texts\n",
    "\n",
    "            # Create sliding window of input/target pairs\n",
    "            for i in range(len(tokens) - 1):\n",
    "                input_seq = tokens[:i+1]\n",
    "                target_seq = tokens[i+1]\n",
    "\n",
    "                # Pad input to consistent length\n",
    "                max_len = 32  # Reasonable context window\n",
    "                if len(input_seq) > max_len:\n",
    "                    input_seq = input_seq[-max_len:]\n",
    "                else:\n",
    "                    input_seq = [0] * (max_len - len(input_seq)) + input_seq\n",
    "\n",
    "                input_sequences.append(input_seq)\n",
    "                target_sequences.append(target_seq)\n",
    "\n",
    "        # Convert to tensors\n",
    "        inputs = Tensor(np.array(input_sequences))\n",
    "        targets = Tensor(np.array(target_sequences))\n",
    "\n",
    "        # Create dataset and dataloader\n",
    "        dataset = TensorDataset(inputs, targets)\n",
    "        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "        print(f\"📚 Training data prepared: {len(dataset)} examples, {len(dataloader)} batches\")\n",
    "        return dataloader\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def train(self, dataloader: DataLoader, epochs: int = 10) -> Dict[str, List[float]]:\n",
    "        \"\"\"\n",
    "        Complete training loop with monitoring.\n",
    "\n",
    "        TODO: Implement full training with progress tracking\n",
    "\n",
    "        APPROACH:\n",
    "        1. Loop through epochs\n",
    "        2. For each batch: forward, backward, optimize\n",
    "        3. Track loss and perplexity\n",
    "        4. Update learning rate schedule\n",
    "        5. Return training history\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> history = pipeline.train(dataloader, epochs=5)\n",
    "        >>> print(f\"Final loss: {history['losses'][-1]:.4f}\")\n",
    "        Final loss: 1.2345\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        history = {'losses': [], 'perplexities': [], 'epochs': []}\n",
    "\n",
    "        print(f\"🚀 Starting training for {epochs} epochs...\")\n",
    "\n",
    "        for epoch in range(epochs):\n",
    "            epoch_losses = []\n",
    "\n",
    "            for batch_idx, (inputs, targets) in enumerate(dataloader):\n",
    "                # Training step\n",
    "                loss = self.trainer.train_step(inputs, targets)\n",
    "                epoch_losses.append(loss)\n",
    "\n",
    "                # Log progress\n",
    "                if batch_idx % 10 == 0:\n",
    "                    perplexity = np.exp(loss)\n",
    "                    print(f\"   Epoch {epoch+1}/{epochs}, Batch {batch_idx}: \"\n",
    "                          f\"Loss={loss:.4f}, PPL={perplexity:.2f}\")\n",
    "\n",
    "            # Epoch summary\n",
    "            avg_loss = np.mean(epoch_losses)\n",
    "            avg_perplexity = np.exp(avg_loss)\n",
    "\n",
    "            history['losses'].append(avg_loss)\n",
    "            history['perplexities'].append(avg_perplexity)\n",
    "            history['epochs'].append(epoch + 1)\n",
    "\n",
    "            # Update learning rate\n",
    "            self.trainer.scheduler.step()\n",
    "\n",
    "            print(f\"✅ Epoch {epoch+1} complete: Loss={avg_loss:.4f}, PPL={avg_perplexity:.2f}\")\n",
    "\n",
    "        self.is_trained = True\n",
    "        self.training_history = history\n",
    "        print(f\"🎉 Training complete! Final perplexity: {history['perplexities'][-1]:.2f}\")\n",
    "\n",
    "        return history\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def optimize_model(self, quantize: bool = True, prune_sparsity: float = 0.0):\n",
    "        \"\"\"\n",
    "        Apply optimization techniques (Modules 17-18).\n",
    "\n",
    "        TODO: Apply quantization and pruning optimizations\n",
    "\n",
    "        APPROACH:\n",
    "        1. Optionally apply quantization to reduce precision\n",
    "        2. Optionally apply pruning to remove weights\n",
    "        3. Measure size reduction\n",
    "        4. Validate model still works\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> pipeline.optimize_model(quantize=True, prune_sparsity=0.5)\n",
    "        Model optimized: 75% size reduction\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        original_params = self.model.count_parameters()\n",
    "        original_memory = original_params * 4 / (1024 * 1024)\n",
    "\n",
    "        optimizations_applied = []\n",
    "\n",
    "        if quantize:\n",
    "            # Apply quantization (simulated)\n",
    "            # In real implementation, would use quantize_model()\n",
    "            quantized_memory = original_memory / 4  # INT8 vs FP32\n",
    "            optimizations_applied.append(f\"INT8 quantization (4× memory reduction)\")\n",
    "            print(\"   Applied INT8 quantization\")\n",
    "\n",
    "        if prune_sparsity > 0:\n",
    "            # Apply pruning (simulated)\n",
    "            # In real implementation, would use magnitude_prune()\n",
    "            remaining_weights = 1 - prune_sparsity\n",
    "            optimizations_applied.append(f\"{prune_sparsity:.0%} pruning ({remaining_weights:.0%} weights remain)\")\n",
    "            print(f\"   Applied {prune_sparsity:.0%} magnitude pruning\")\n",
    "\n",
    "        # Calculate final size\n",
    "        size_reduction = 1.0\n",
    "        if quantize:\n",
    "            size_reduction *= 0.25  # 4× smaller\n",
    "        if prune_sparsity > 0:\n",
    "            size_reduction *= (1 - prune_sparsity)\n",
    "\n",
    "        final_memory = original_memory * size_reduction\n",
    "        reduction_factor = original_memory / final_memory\n",
    "\n",
    "        print(f\"🔧 Model optimization complete:\")\n",
    "        print(f\"   Original: {original_memory:.1f}MB\")\n",
    "        print(f\"   Optimized: {final_memory:.1f}MB\")\n",
    "        print(f\"   Reduction: {reduction_factor:.1f}× smaller\")\n",
    "        print(f\"   Applied: {', '.join(optimizations_applied)}\")\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def generate_text(self, prompt: str, max_tokens: int = 50) -> str:\n",
    "        \"\"\"\n",
    "        Generate text using the trained model.\n",
    "\n",
    "        TODO: Implement text generation with proper encoding/decoding\n",
    "\n",
    "        APPROACH:\n",
    "        1. Encode prompt to token IDs\n",
    "        2. Use model.generate() for autoregressive generation\n",
    "        3. Decode generated tokens back to text\n",
    "        4. Return generated text\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> text = pipeline.generate_text(\"Hello\", max_tokens=10)\n",
    "        >>> print(f\"Generated: {text}\")\n",
    "        Generated: Hello world this is AI\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        if not self.is_trained:\n",
    "            print(\"⚠️ Model not trained yet. Generating with random weights.\")\n",
    "\n",
    "        # Encode prompt\n",
    "        prompt_tokens = self.tokenizer.encode(prompt)\n",
    "        prompt_tensor = Tensor([prompt_tokens])\n",
    "\n",
    "        # Generate tokens\n",
    "        generated_tokens = self.model.generate(\n",
    "            prompt_tensor,\n",
    "            max_new_tokens=max_tokens,\n",
    "            temperature=0.8,\n",
    "            use_cache=True\n",
    "        )\n",
    "\n",
    "        # Decode to text\n",
    "        all_tokens = generated_tokens.data[0].tolist()\n",
    "        generated_text = self.tokenizer.decode(all_tokens)\n",
    "\n",
    "        return generated_text\n",
    "        ### END SOLUTION\n",
    "\n",
    "def test_unit_complete_pipeline():\n",
    "    \"\"\"🔬 Test complete pipeline integration.\"\"\"\n",
    "    print(\"🔬 Unit Test: Complete Pipeline Integration...\")\n",
    "\n",
    "    # Create pipeline\n",
    "    pipeline = CompleteTinyGPTPipeline(vocab_size=50, embed_dim=32, num_layers=2)\n",
    "\n",
    "    # Test data preparation\n",
    "    corpus = [\"hello world\", \"ai is fun\", \"machine learning\"]\n",
    "    dataloader = pipeline.prepare_training_data(corpus, batch_size=2)\n",
    "    assert len(dataloader) > 0, \"DataLoader should have batches\"\n",
    "\n",
    "    # Test training (minimal)\n",
    "    history = pipeline.train(dataloader, epochs=1)\n",
    "    assert 'losses' in history, \"History should contain losses\"\n",
    "    assert len(history['losses']) == 1, \"Should have one epoch of losses\"\n",
    "\n",
    "    # Test optimization\n",
    "    pipeline.optimize_model(quantize=True, prune_sparsity=0.5)\n",
    "\n",
    "    # Test generation\n",
    "    generated = pipeline.generate_text(\"hello\", max_tokens=5)\n",
    "    assert isinstance(generated, str), \"Generated output should be string\"\n",
    "    assert len(generated) > 0, \"Generated text should not be empty\"\n",
    "\n",
    "    print(f\"✅ Pipeline stages completed successfully\")\n",
    "    print(f\"✅ Training history: {len(history['losses'])} epochs\")\n",
    "    print(f\"✅ Generated text: '{generated[:20]}...'\")\n",
    "    print(\"✅ Complete pipeline integration works!\")\n",
    "\n",
    "# Run immediate test\n",
    "test_unit_complete_pipeline()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf266828",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 🎯 Module Integration Test\n",
    "\n",
    "Final comprehensive test validating all components work together correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d3801eb",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test_module",
     "locked": true,
     "points": 20
    }
   },
   "outputs": [],
   "source": [
    "def test_module():\n",
    "    \"\"\"\n",
    "    Comprehensive test of entire capstone module functionality.\n",
    "\n",
    "    This final test runs before module summary to ensure:\n",
    "    - TinyGPT architecture works correctly\n",
    "    - Training pipeline integrates properly\n",
    "    - Optimization techniques can be applied\n",
    "    - Text generation produces output\n",
    "    - All systems analysis functions execute\n",
    "    - Complete pipeline demonstrates end-to-end functionality\n",
    "    \"\"\"\n",
    "    print(\"🧪 RUNNING MODULE INTEGRATION TEST\")\n",
    "    print(\"=\" * 60)\n",
    "\n",
    "    # Test 1: TinyGPT Architecture\n",
    "    print(\"🔬 Testing TinyGPT architecture...\")\n",
    "    test_unit_tinygpt_init()\n",
    "    test_unit_tinygpt_forward()\n",
    "\n",
    "    # Test 2: Training Pipeline\n",
    "    print(\"\\n🔬 Testing training pipeline...\")\n",
    "    test_unit_training_pipeline()\n",
    "\n",
    "    # Test 3: Complete Pipeline\n",
    "    print(\"\\n🔬 Testing complete pipeline...\")\n",
    "    test_unit_complete_pipeline()\n",
    "\n",
    "    # Test 4: Systems Analysis\n",
    "    print(\"\\n🔬 Testing systems analysis...\")\n",
    "\n",
    "    # Create model for final validation\n",
    "    print(\"🔬 Final integration test...\")\n",
    "    model = TinyGPT(vocab_size=100, embed_dim=64, num_layers=2, num_heads=2)\n",
    "\n",
    "    # Verify core functionality\n",
    "    assert hasattr(model, 'count_parameters'), \"Model should have parameter counting\"\n",
    "    assert hasattr(model, 'forward'), \"Model should have forward method\"\n",
    "    assert hasattr(model, 'generate'), \"Model should have generation method\"\n",
    "\n",
    "    # Test parameter counting\n",
    "    param_count = model.count_parameters()\n",
    "    assert param_count > 0, \"Model should have parameters\"\n",
    "\n",
    "    # Test forward pass\n",
    "    test_input = Tensor([[1, 2, 3, 4, 5]])\n",
    "    output = model.forward(test_input)\n",
    "    assert output.shape == (1, 5, 100), f\"Expected (1, 5, 100), got {output.shape}\"\n",
    "\n",
    "    # Test generation\n",
    "    generated = model.generate(test_input, max_new_tokens=3)\n",
    "    assert generated.shape[1] == 8, f\"Expected 8 tokens, got {generated.shape[1]}\"\n",
    "\n",
    "    print(\"\\n\" + \"=\" * 60)\n",
    "    print(\"🎉 ALL CAPSTONE TESTS PASSED!\")\n",
    "    print(\"🚀 TinyGPT system fully functional!\")\n",
    "    print(\"✅ All 19 modules successfully integrated!\")\n",
    "    print(\"🎯 Ready for real-world deployment!\")\n",
    "    print(\"\\nRun: tito module complete 20\")\n",
    "\n",
    "# Call the comprehensive test\n",
    "test_module()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd35174b",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "main_execution",
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    print(\"🚀 Running TinyGPT Capstone module...\")\n",
    "\n",
    "    # Run the comprehensive test\n",
    "    test_module()\n",
    "\n",
    "    # Demo the complete system\n",
    "    print(\"\\n\" + \"=\" * 60)\n",
    "    print(\"🎭 CAPSTONE DEMONSTRATION\")\n",
    "    print(\"=\" * 60)\n",
    "\n",
    "    # Create a demo pipeline\n",
    "    print(\"🏗️ Creating demonstration pipeline...\")\n",
    "    demo_pipeline = CompleteTinyGPTPipeline(\n",
    "        vocab_size=100,\n",
    "        embed_dim=128,\n",
    "        num_layers=4,\n",
    "        num_heads=4\n",
    "    )\n",
    "\n",
    "    # Show parameter breakdown\n",
    "    print(f\"\\n📊 Model Architecture Summary:\")\n",
    "    print(f\"   Parameters: {demo_pipeline.model.count_parameters():,}\")\n",
    "    print(f\"   Layers: {demo_pipeline.num_layers}\")\n",
    "    print(f\"   Heads: {demo_pipeline.num_heads}\")\n",
    "    print(f\"   Embedding dimension: {demo_pipeline.embed_dim}\")\n",
    "\n",
    "    # Demonstrate text generation (with untrained model)\n",
    "    print(f\"\\n🎭 Demonstration Generation (untrained model):\")\n",
    "    sample_text = demo_pipeline.generate_text(\"Hello\", max_tokens=10)\n",
    "    print(f\"   Input: 'Hello'\")\n",
    "    print(f\"   Output: '{sample_text}'\")\n",
    "    print(f\"   Note: Random output expected (model not trained)\")\n",
    "\n",
    "    print(\"\\n✅ Capstone demonstration complete!\")\n",
    "    print(\"🎯 TinyGPT represents the culmination of 19 modules of ML systems learning!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4e23b97",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🤔 ML Systems Thinking: Capstone Reflection\n",
    "\n",
    "This capstone integrates everything you've learned across 19 modules. Let's reflect on the complete systems picture.\n",
    "\n",
    "### Question 1: Architecture Scaling\n",
    "You built TinyGPT with configurable architecture (embed_dim, num_layers, num_heads).\n",
    "If you double the embed_dim from 128 to 256, approximately how much does memory usage increase?\n",
    "\n",
    "**Answer:** _______ (2×, 4×, 8×, or 16×)\n",
    "\n",
    "**Reasoning:** Consider that embed_dim affects embedding tables, all linear layers in attention, and MLP layers.\n",
    "\n",
    "### Question 2: Training vs Inference Memory\n",
    "Your TinyGPT uses different memory patterns for training vs inference.\n",
    "For a model with 50M parameters, what's the approximate memory usage difference?\n",
    "\n",
    "**Training Memory:** _______ MB\n",
    "**Inference Memory:** _______ MB\n",
    "**Ratio:** _______ × larger for training\n",
    "\n",
    "**Hint:** Training requires parameters + gradients + optimizer states (Adam has 2 momentum terms).\n",
    "\n",
    "### Question 3: Optimization Trade-offs\n",
    "You implemented quantization (INT8) and pruning (90% sparsity) optimizations.\n",
    "For the original 200MB model, what's the memory footprint after both optimizations?\n",
    "\n",
    "**Original:** 200MB\n",
    "**After INT8 + 90% pruning:** _______ MB\n",
    "**Total reduction factor:** _______ ×\n",
    "\n",
    "### Question 4: Generation Complexity\n",
    "Your generate() method can use KV caching for efficiency.\n",
    "For generating 100 tokens with sequence length 500, how many forward passes are needed?\n",
    "\n",
    "**Without KV cache:** _______ forward passes\n",
    "**With KV cache:** _______ forward passes\n",
    "**Speedup factor:** _______ ×\n",
    "\n",
    "### Question 5: Systems Integration\n",
    "You integrated 19 different modules into a cohesive system.\n",
    "Which integration challenge was most critical for making TinyGPT work?\n",
    "\n",
    "a) Making all imports work correctly\n",
    "b) Ensuring tensor shapes flow correctly through all components\n",
    "c) Managing memory during training\n",
    "d) Coordinating the generation loop with KV caching\n",
    "\n",
    "**Answer:** _______\n",
    "\n",
    "**Explanation:** ________________________________"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fbc1ae3",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🎯 MODULE SUMMARY: Capstone - Complete TinyGPT System\n",
    "\n",
    "Congratulations! You've completed the ultimate integration project - building TinyGPT from your own ML framework!\n",
    "\n",
    "### Key Accomplishments\n",
    "- **Integrated 19 modules** into a cohesive, production-ready system\n",
    "- **Built complete TinyGPT** with training, optimization, and generation capabilities\n",
    "- **Demonstrated systems thinking** with memory analysis, performance profiling, and optimization\n",
    "- **Created end-to-end pipeline** from raw text to trained model to generated output\n",
    "- **Applied advanced optimizations** including quantization and pruning\n",
    "- **Validated the complete framework** through comprehensive testing\n",
    "- All tests pass ✅ (validated by `test_module()`)\n",
    "\n",
    "### Systems Insights Gained\n",
    "- **Architecture scaling**: How model size affects memory and compute requirements\n",
    "- **Training dynamics**: Memory patterns, convergence monitoring, and optimization\n",
    "- **Production optimization**: Quantization and pruning for deployment efficiency\n",
    "- **Integration complexity**: How modular design enables complex system composition\n",
    "\n",
    "### The Complete Journey\n",
    "```\n",
    "Module 01: Tensor Operations\n",
    "    ↓\n",
    "Modules 02-04: Neural Network Basics\n",
    "    ↓\n",
    "Modules 05-07: Training Infrastructure\n",
    "    ↓\n",
    "Modules 08-09: Data and Spatial Processing\n",
    "    ↓\n",
    "Modules 10-14: Language Models and Transformers\n",
    "    ↓\n",
    "Modules 15-19: Systems Optimization\n",
    "    ↓\n",
    "Module 20: COMPLETE TINYGPT SYSTEM! 🎉\n",
    "```\n",
    "\n",
    "### Ready for the Real World\n",
    "Your TinyGPT implementation demonstrates:\n",
    "- **Production-quality code** with proper error handling and optimization\n",
    "- **Systems engineering mindset** with performance analysis and memory management\n",
    "- **ML framework design** understanding how PyTorch-like systems work internally\n",
    "- **End-to-end ML pipeline** from data to deployment\n",
    "\n",
    "**Export with:** `tito module complete 20`\n",
    "\n",
    "**Achievement Unlocked:** 🏆 **ML Systems Engineer** - You've built a complete AI system from scratch!\n",
    "\n",
    "You now understand how modern AI systems work from the ground up. From tensors to text generation, from training loops to production optimization - you've mastered the full stack of ML systems engineering.\n",
    "\n",
    "**What's Next:** Take your TinyTorch framework and build even more ambitious projects! The foundations you've built can support any ML architecture you can imagine."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
