{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ANLP Assignment 2: Mixture of Experts, Pretraining Optimizers, and Decoding Strategies\n",
    "### Complete Self-Contained GPU Runner for Google Colab Pro (T4 / V100 / A100)\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Chiranjeevi-2020/ANLP_assignment2/blob/main/colab/ANLP_Assignment2_Colab_Runner.ipynb)\n",
    "\n",
    "Welcome! This notebook executes all three core sections of **Assignment 2 (Monsoon 2026, IIIT Hyderabad)**:\n",
    "1. **Part 1: Mixture of Experts (MoE) (40 marks)** \u2014 5 FFN variants, parameter parity matching, Switch auxiliary loss, and language specialization heatmaps on `belumind/en-vi-ja-curated-500k-triplets`.\n",
    "2. **Part 2: Custom Optimizers from Scratch (30 marks)** \u2014 Category 1 (AdamW), Category 2 (Adan), Category 3 (Lion), Category 4 (Muon), and Bonus (Sophia) pretraining dynamics on `browndw/human-ai-parallel-corpus`.\n",
    "3. **Part 3: Decoding Strategies from Scratch (30 marks)** \u2014 Greedy, Top-k, Top-p, and Beam Search ($B \\in \\{1, 2, 4\\}$) on `EleutherAI/Pythia-160M` evaluated on `hamishivi/ROCStories` (1,000 samples) using strictly `model.forward()`.\n",
    "\n",
    "---\n",
    "### \ud83d\udc76 Explain Like I'm 5 (ELI5): What Are We Doing in This Notebook?\n",
    "Imagine you are building a team of smart robots to read and write stories:\n",
    "- **Part 1 (Mixture of Experts)**: Instead of asking one big robot to know every language, we have a smart receptionist that sends Vietnamese words to one expert robot and Japanese words to another expert robot! They share the work so each robot does what it's best at.\n",
    "- **Part 2 (Optimizers)**: How do robots learn from their mistakes? An optimizer is like a coach telling the robot how big of a step to take. We test 5 different coaches from scratch (like Muon, which uses matrix geometry, and Lion, which saves half the memory) to see who trains the robot fastest!\n",
    "- **Part 3 (Decoding Strategies)**: When the robot writes the next word of a story, how does it choose? Greedy always takes the highest score; Top-k and Top-p pick from the best fun choices so the story isn't boring; Beam search looks ahead like a chess master to find the best storyline!\n",
    "\n",
    "---\n",
    "### \u23f1\ufe0f Expected Runtimes on Google Colab Pro\n",
    "| Pipeline Stage | Fast Verification / Smoke Test | Full Dataset Training (Colab Pro A100/V100) |\n",
    "| :--- | :--- | :--- |\n",
    "| **Environment & Git Setup** | ~30 seconds | ~30 seconds |\n",
    "| **Part 1: 5 MoE Ablations** | ~2 minutes (dummy tensors) | ~45-60 minutes (A100) / ~90 minutes (T4) |\n",
    "| **Part 2: 5 Optimizers Dynamics** | ~1 minute (warmup) | ~30-40 minutes (A100) / ~60 minutes (T4) |\n",
    "| **Part 3: Decoding Benchmark (1,000 samples)** | ~30 seconds (10 samples) | ~10-15 minutes (1,000 samples on A100) |\n",
    "| **Total End-to-End Execution** | **~5 minutes** | **~1.5 to 2.5 hours** |\n",
    "\n",
    "---\n",
    "### \ud83d\udcc1 Where Are Files & Checkpoints Generated?\n",
    "- Trained Model Checkpoints: `checkpoints/part1/` and `checkpoints/part2/`\n",
    "- Heatmaps & Plots: `assets/expert_heatmap.png`, `assets/optimizer_dynamics.png`, `assets/beam_search_runtime.png`\n",
    "- Numerical Results & Logs: `assets/part2/optimizer_benchmark_results.json`, `assets/part3/decoding_summary.json`\n",
    "- At the end of this notebook, a single cell zips all outputs into `submission_artifacts.zip` for instant local download."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Clone Repository & Set Working Directory (ELI5: Put our project on the computer)\n",
    "\n",
    "Google Colab starts with an empty `/content` folder. This cell automatically:\n",
    "1. Checks Colab Secrets for a `GH_TOKEN` (needed if your repo is **private**).\n",
    "2. Clones your GitHub repository `ANLP_assignment2` into `/content/ANLP_assignment2`.\n",
    "3. Changes the active working directory into the repository root.\n",
    "4. Pulls the latest commits from branch `main`.\n",
    "\n",
    "**If your repo is private**, add a secret named `GH_TOKEN` in the \ud83d\udd11 Secrets tab with a [GitHub Personal Access Token](https://github.com/settings/tokens) (Classic, with `repo` scope). Toggle **Notebook access** ON.\n",
    "\n",
    "**If your repo is public**, no token is needed \u2014 it just works!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, subprocess\n",
    "\n",
    "REPO_DIR = \"/content/ANLP_assignment2\"\n",
    "\n",
    "# --- Fetch GitHub token from Colab Secrets Vault for private repo access ---\n",
    "gh_token = None\n",
    "try:\n",
    "    from google.colab import userdata\n",
    "    for name in ['GH_TOKEN', 'GITHUB_TOKEN', 'github_token', 'gh_token']:\n",
    "        try:\n",
    "            gh_token = userdata.get(name)\n",
    "            if gh_token:\n",
    "                print(f'\\U0001f513 Found GitHub token in Colab Secrets (key: {name})')\n",
    "                break\n",
    "        except Exception:\n",
    "            pass\n",
    "except ImportError:\n",
    "    pass  # Not running in Colab\n",
    "\n",
    "# --- Build clone URL (with or without token) ---\n",
    "if gh_token:\n",
    "    clone_url = f'https://{gh_token}@github.com/Chiranjeevi-2020/ANLP_assignment2.git'\n",
    "else:\n",
    "    clone_url = 'https://github.com/Chiranjeevi-2020/ANLP_assignment2.git'\n",
    "    print('\\u26a0\\ufe0f No GitHub token found in Colab Secrets.')\n",
    "    print('   If your repo is PRIVATE, add a secret named GH_TOKEN with a GitHub Personal Access Token.')\n",
    "    print('   If your repo is PUBLIC, this will work without a token.')\n",
    "\n",
    "# --- Clone if not already present ---\n",
    "if not os.path.exists(REPO_DIR):\n",
    "    result = subprocess.run(['git', 'clone', clone_url, REPO_DIR], capture_output=True, text=True)\n",
    "    if result.returncode != 0:\n",
    "        print(f'\\u274c Clone failed: {result.stderr.strip()}')\n",
    "        print('\\n\\U0001f527 FIX: Your repo is likely PRIVATE. Do ONE of these:')\n",
    "        print('   Option A (Easiest): Go to https://github.com/Chiranjeevi-2020/ANLP_assignment2/settings')\n",
    "        print('                       Scroll to \"Danger Zone\" -> \"Change visibility\" -> Make it Public')\n",
    "        print('   Option B: Create a GitHub Personal Access Token at https://github.com/settings/tokens')\n",
    "        print('             Add it to Colab Secrets (\\U0001f511 icon) as GH_TOKEN with Notebook access ON')\n",
    "        raise RuntimeError('Repository clone failed. See instructions above.')\n",
    "    else:\n",
    "        print(f'\\u2705 Successfully cloned repository!')\n",
    "\n",
    "# --- Navigate into repository root ---\n",
    "%cd {REPO_DIR}\n",
    "\n",
    "# --- Pull latest changes ---\n",
    "!git pull origin main --ff-only 2>/dev/null || true\n",
    "\n",
    "# --- Add to Python path ---\n",
    "if REPO_DIR not in sys.path:\n",
    "    sys.path.insert(0, REPO_DIR)\n",
    "\n",
    "print(f'\\n\\u2705 Project Root: {os.getcwd()}')\n",
    "print(f'\\u2705 Files: {sorted([f for f in os.listdir(\".\") if not f.startswith(\".\")][:15])}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: GPU Diagnostics & Setup (ELI5: Check that the supercomputer is awake)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Verify GPU assignment\n",
    "!nvidia-smi\n",
    "import torch\n",
    "print(f\"CUDA Available: {torch.cuda.is_available()}\")\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"Allocated GPU: {torch.cuda.get_device_name(0)}\")\n",
    "    print(f\"Device Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Install Required Libraries (ELI5: Download the tools our robot needs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install core NLP, model evaluation, and dashboard packages\n",
    "!pip install -q pytest accelerate torch transformers datasets tokenizers sacrebleu evaluate wandb matplotlib seaborn einops huggingface_hub"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Authenticate WandB & Hugging Face via Colab Secrets Vault (ELI5: Safe Login)\n",
    "\n",
    "This cell securely fetches your API tokens from Google Colab's **Secrets Vault** without exposing them in code:\n",
    "\n",
    "### \ud83d\udd11 How to set up your secrets in Colab (One-Time Setup):\n",
    "1. Click the **Key icon (\ud83d\udd11)** on Google Colab's left sidebar (Secrets panel).\n",
    "2. Click **+ Add new secret**.\n",
    "3. Secret 1: Name = `WANDB_API_KEY`, Value = *(your WandB API key from [wandb.ai/authorize](https://wandb.ai/authorize))*\n",
    "4. Secret 2: Name = `HF_TOKEN`, Value = *(your Hugging Face token from [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens))*\n",
    "5. \u26a0\ufe0f **CRITICAL**: Toggle **Notebook access** to **ON** for both secrets!\n",
    "\n",
    "*If secrets are not added, the notebook will operate in safe offline mode with zero crashes.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# 1. Fetch WandB Secret from Colab Secrets Vault\n",
    "wandb_key = None\n",
    "try:\n",
    "    from google.colab import userdata\n",
    "    for name in ['WANDB_API_KEY', 'wandb_api_key', 'WANDB', 'wandb']:\n",
    "        try:\n",
    "            wandb_key = userdata.get(name)\n",
    "            if wandb_key:\n",
    "                print(f\"[OK] Found WandB secret under key: '{name}'\")\n",
    "                break\n",
    "        except Exception as e:\n",
    "            err_msg = str(e).lower()\n",
    "            if \"notebook access\" in err_msg or \"notebookaccesserror\" in type(e).__name__.lower():\n",
    "                print(f\"[WARNING] Secret '{name}' exists in Colab Secrets Vault, but 'Notebook access' toggle is OFF! Please turn ON the toggle switch next to '{name}' in the Colab Secrets (\ud83d\udd11) tab.\")\n",
    "except Exception:\n",
    "    pass\n",
    "\n",
    "if wandb_key:\n",
    "    os.environ['WANDB_API_KEY'] = wandb_key\n",
    "    try:\n",
    "        import wandb\n",
    "        wandb.login(key=wandb_key)\n",
    "        print(\"[SUCCESS] Successfully authenticated WandB via Colab Secrets!\")\n",
    "    except Exception as e:\n",
    "        print(f\"[NOTE] WandB login notice: {e}\")\n",
    "else:\n",
    "    print(\"[INFO] WandB secret not found in Secrets Vault. Operating in safe offline mode (wandb disabled).\")\n",
    "    os.environ['WANDB_MODE'] = 'offline'\n",
    "\n",
    "# 2. Fetch Hugging Face Secret from Colab Secrets Vault\n",
    "hf_token = None\n",
    "try:\n",
    "    from google.colab import userdata\n",
    "    for name in ['HF_TOKEN', 'hf_token', 'HUGGINGFACE_TOKEN', 'huggingface_token', 'HF', 'huggingface']:\n",
    "        try:\n",
    "            hf_token = userdata.get(name)\n",
    "            if hf_token:\n",
    "                print(f\"[OK] Found Hugging Face secret under key: '{name}'\")\n",
    "                break\n",
    "        except Exception as e:\n",
    "            err_msg = str(e).lower()\n",
    "            if \"notebook access\" in err_msg or \"notebookaccesserror\" in type(e).__name__.lower():\n",
    "                print(f\"[WARNING] Secret '{name}' exists in Colab Secrets Vault, but 'Notebook access' toggle is OFF! Please turn ON the toggle switch next to '{name}' in the Colab Secrets (\ud83d\udd11) tab.\")\n",
    "except Exception:\n",
    "    pass\n",
    "\n",
    "if hf_token:\n",
    "    os.environ['HF_TOKEN'] = hf_token\n",
    "    os.environ['HUGGING_FACE_HUB_TOKEN'] = hf_token\n",
    "    try:\n",
    "        from huggingface_hub import login\n",
    "        login(token=hf_token, add_to_git_credential=True)\n",
    "        print(\"[SUCCESS] Successfully authenticated Hugging Face via Colab Secrets!\")\n",
    "    except Exception as e:\n",
    "        print(f\"[NOTE] Hugging Face login notice: {e}\")\n",
    "else:\n",
    "    print(\"[INFO] HF secret not found in Secrets Vault. Public datasets and models will download normally.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Verify Mathematical Integrity & Test Suite (ELI5: Double check all formulas and code tests)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Runs comprehensive pytest suite across all 22 mathematical tests and main verify\n",
    "!python -m pytest tests/ -v\n",
    "!python main.py verify"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Part 1 \u2014 Mixture of Experts (MoE) Execution\n",
    "### ELI5:\n",
    "We test 5 designs:\n",
    "1. **Dense**: 1 big standard feedforward brain.\n",
    "2. **Top-1 MoE**: 4 experts, 1 chosen per token.\n",
    "3. **Top-2 MoE**: 4 experts, 2 chosen per token.\n",
    "4. **Shared + Routed**: 1 expert always stays on, plus 1 routed expert chosen from 3.\n",
    "5. **Matched Active**: 4 experts with doubled width so active parameters match the dense baseline exactly.\n",
    "\n",
    "We train on `belumind/en-vi-ja-curated-500k-triplets` to translate Vietnamese & Japanese to English, track routing with Switch auxiliary loss, and extract language heatmaps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run Part 1 training across all 5 architectural ablations\n",
    "# To run a quick smoke run, keep samples=2000, epochs=1.\n",
    "# For full training, increase samples=10000, epochs=3.\n",
    "\n",
    "ablations = [\"dense\", \"top1\", \"top2\", \"shared_routed\", \"matched_active\"]\n",
    "for abl in ablations:\n",
    "    print(f\"\\n{'='*25} RUNNING PART 1 ABLATION: {abl.upper()} {'='*25}\")\n",
    "    !python main.py part1 --ablation {abl} --train --samples 2000 --epochs 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: Part 2 \u2014 Custom Optimizers Pretraining Dynamics\n",
    "### ELI5:\n",
    "We pretrain our dense model on `browndw/human-ai-parallel-corpus` using 5 scratch optimizers:\n",
    "1. **AdamW**: Standard baseline with decoupled weight decay.\n",
    "2. **Adan**: Looks at gradient differences to smooth out bumpy steps.\n",
    "3. **Lion**: Uses only `+1` or `-1` signs, cutting memory in half!\n",
    "4. **Muon**: Uses Newton-Schulz polynomial iterations to rotate matrices without distortion.\n",
    "5. **Sophia**: Uses second-order diagonal curvature and clipping to handle steep hills.\n",
    "\n",
    "We log validation loss and test BLEU at $0.1\\times$ intervals to replicate Figure 2 of *Fantastic Pretraining Optimizers*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run Part 2 Pretraining Dynamics benchmark for all 5 optimizers\n",
    "# tokens=100000 provides quick validation; tokens=500000 gives full Chinchilla curve.\n",
    "!python main.py part2 --optimizer all --train --tokens 100000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 8: Part 3 \u2014 Decoding Strategies on ROCStories\n",
    "### ELI5:\n",
    "Using `EleutherAI/Pythia-160M` without calling `model.generate()`, we generate story endings for 1,000 ROCStories samples using:\n",
    "- **Greedy**: Always pick top token.\n",
    "- **Top-k ($k=40, T=0.8$)**: Pick from top 40 words.\n",
    "- **Top-p ($p=0.9, T=0.8$)**: Pick from the top 90% probability mass.\n",
    "- **Beam Search ($B=1, 2, 4$)**: Keep the best $B$ candidate story paths.\n",
    "\n",
    "We measure perplexity, accuracy, and confirm linear latency scaling with respect to beam width $B$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run Part 3 Decoding Benchmark on 1000 samples\n",
    "!python main.py part3 --model_name EleutherAI/Pythia-160M --samples 1000 --output_dir assets/part3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 9: View Visualizations (ELI5: Look at the pretty pictures our robots made!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from IPython.display import Image, display\n",
    "\n",
    "figs = [\n",
    "    (\"MoE Language Specialization Heatmap\", \"assets/expert_heatmap.png\"),\n",
    "    (\"Optimizer Pretraining Dynamics (Fig 2)\", \"assets/optimizer_dynamics.png\"),\n",
    "    (\"Beam Search Latency vs Beam Width\", \"assets/beam_search_runtime.png\")\n",
    "]\n",
    "\n",
    "for title, path in figs:\n",
    "    if os.path.exists(path):\n",
    "        print(f\"\\n=== {title} ({path}) ===\")\n",
    "        display(Image(path))\n",
    "    else:\n",
    "        print(f\"Note: {path} not yet generated.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 10: Download All Artifacts Locally (ELI5: Pack all homework into a neat zip bag to submit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Zips all assets, figures, configs, and json results for submission\n",
    "!zip -r submission_artifacts.zip assets/ report.pdf configs/\n",
    "\n",
    "try:\n",
    "    from google.colab import files\n",
    "    files.download('submission_artifacts.zip')\n",
    "    print(\"Downloaded submission_artifacts.zip!\")\n",
    "except Exception as e:\n",
    "    print(f\"Running outside Google Colab environment. File saved as submission_artifacts.zip in current directory ({e}).\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 11: Syncing Your Executed Notebook & Results to GitHub (ELI5: Save your work online)\n",
    "\n",
    "There are two easy ways to get your executed notebook and results on GitHub:\n",
    "\n",
    "### Option A: Save Executed Notebook Directly via Colab Menu (Recommended & Fastest \ud83d\ude80)\n",
    "To save this completed notebook with all its live execution outputs, logs, and plots:\n",
    "1. Click **File** (top left menu of Colab) $\\to$ **Save a copy in GitHub**.\n",
    "2. Select Repository: `Chiranjeevi-2020/ANLP_assignment2`.\n",
    "3. Select Branch: `main`.\n",
    "4. File path: `colab/ANLP_Assignment2_Colab_Runner.ipynb`.\n",
    "5. Click **OK**. Google Colab commits the executed notebook directly to your GitHub repo!\n",
    "\n",
    "### Option B: Push Generated Assets Directly from this Colab Cell (Optional)\n",
    "If you have a GitHub Personal Access Token saved in Colab Secrets as `GH_TOKEN` or `GITHUB_TOKEN`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional: Push generated figures and results directly from Colab VM\n",
    "try:\n",
    "    from google.colab import userdata\n",
    "    gh_token = None\n",
    "    for name in ['GH_TOKEN', 'GITHUB_TOKEN', 'github_token']:\n",
    "        try:\n",
    "            gh_token = userdata.get(name)\n",
    "            if gh_token: break\n",
    "        except Exception: pass\n",
    "    \n",
    "    if gh_token:\n",
    "        !git config user.email \"chiranjeevi.kumaravadivel@gmail.com\"\n",
    "        !git config user.name \"Chiranjeevi-2020\"\n",
    "        !git add assets/ configs/ report.pdf\n",
    "        !git commit -m \"Colab GPU execution: update high-res training dynamics and assets\"\n",
    "        !git push https://{gh_token}@github.com/Chiranjeevi-2020/ANLP_assignment2.git main\n",
    "        print(\"Successfully pushed Colab training assets to GitHub!\")\n",
    "    else:\n",
    "        print(\"No GH_TOKEN found in Colab Secrets. Use Option A (File -> Save a copy in GitHub) or download submission_artifacts.zip from Step 10.\")\n",
    "except Exception as e:\n",
    "    print(f\"Git push notice: {e}\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "language_info": {
   "name": "python",
   "version": "3.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}