{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0eb387df",
   "metadata": {},
   "source": [
    "# 🧠 UMBRA — Google Colab Training Notebook\n",
    "\n",
    "**Model:** `HuggingFaceTB/SmolLM-135M` (4-bit QLoRA)  \n",
    "**Trainer:** GRPO via TRL  \n",
    "**Datasets:** hh-rlhf · pii-masking-300k · truthful_qa · toxic-chat · daily_dialog  \n",
    "\n",
    "**Runtime:** T4 GPU (16 GB) — set via *Runtime → Change runtime type → T4*\n",
    "\n",
    "---\n",
    "### Steps\n",
    "1. Run **Cell 1** to install dependencies  \n",
    "2. Run **Cell 2** to mount Google Drive and upload your project  \n",
    "3. Run **Cell 3** to verify the GPU  \n",
    "4. Run **Cell 4** to start training  \n",
    "5. Run **Cell 5** to evaluate after training  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4e4290f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Cell 1: Install all dependencies ──────────────────────────────────────────\n",
    "!pip install -q \\\n",
    "    \"trl>=0.9.0\" \\\n",
    "    \"datasets>=2.18.0\" \\\n",
    "    \"transformers>=4.40.0\" \\\n",
    "    \"bitsandbytes>=0.43.0\" \\\n",
    "    \"peft>=0.10.0\" \\\n",
    "    \"accelerate>=0.29.0\" \\\n",
    "    \"gymnasium>=0.29.0\" \\\n",
    "    \"matplotlib>=3.8\" \\\n",
    "    \"numpy>=1.26\" \\\n",
    "    \"fastapi\" \\\n",
    "    \"uvicorn\"\n",
    "\n",
    "print('✅ All dependencies installed.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63e49dee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Cell 2: Mount Google Drive and set project path ───────────────────────────\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "\n",
    "import os\n",
    "\n",
    "# ⚠️  Change this path to wherever you uploaded the Umbra folder in your Drive\n",
    "PROJECT_PATH = '/content/drive/MyDrive/Umbra'\n",
    "\n",
    "assert os.path.isdir(PROJECT_PATH), (\n",
    "    f'❌ Project folder not found at {PROJECT_PATH}. '\n",
    "    'Upload your Umbra folder to Google Drive and update PROJECT_PATH above.'\n",
    ")\n",
    "\n",
    "os.chdir(PROJECT_PATH)\n",
    "print(f'✅ Working directory: {os.getcwd()}')\n",
    "print('Files:', os.listdir('.'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0478e990",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Cell 3: Verify GPU + fix Python path + smoke test ─────────────────────────\n",
    "import sys, os, torch\n",
    "\n",
    "# ── CRITICAL: add project root to sys.path so all UMBRA modules resolve ───────\n",
    "# This is the fix for: ModuleNotFoundError: No module named 'reward'\n",
    "PROJECT_PATH = os.getcwd()   # already set to /content/drive/MyDrive/Umbra by Cell 2\n",
    "if PROJECT_PATH not in sys.path:\n",
    "    sys.path.insert(0, PROJECT_PATH)\n",
    "print(f'✅ sys.path[0] = {sys.path[0]}')\n",
    "\n",
    "# ── GPU check ─────────────────────────────────────────────────────────────────\n",
    "if not torch.cuda.is_available():\n",
    "    raise RuntimeError(\n",
    "        '❌ No GPU detected! Go to Runtime → Change runtime type → T4 GPU and restart.'\n",
    "    )\n",
    "gpu_name = torch.cuda.get_device_name(0)\n",
    "vram_gb  = torch.cuda.get_device_properties(0).total_memory / 1e9\n",
    "print(f'✅ GPU: {gpu_name}  |  VRAM: {vram_gb:.1f} GB')\n",
    "\n",
    "# ── Full import sanity check ───────────────────────────────────────────────────\n",
    "from reward.reward_model      import ShapedRewardModel\n",
    "from sentrix.cialdini_stress  import run_cialdini_stress\n",
    "from shadow.shadow_agent      import ShadowAgent\n",
    "from shadow.arms_race_trainer import ArmsRaceTrainer\n",
    "from demo.graph_generator     import generate_all_graphs\n",
    "from data.dataset_loader      import load_hh_rlhf, build_grpo_prompts\n",
    "print('✅ All UMBRA modules imported successfully.')\n",
    "\n",
    "# ── Dataset loader quick smoke test ───────────────────────────────────────────\n",
    "sample_ds      = load_hh_rlhf(split='train', max_samples=5)\n",
    "sample_prompts = build_grpo_prompts(sample_ds, max_prompts=3)\n",
    "print(f'✅ Dataset loader OK — sample prompt: {sample_prompts[0][:80]}…')\n",
    "print('\\n🚀 Ready to train! Run Cell 4.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9ec67e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Cell 4: Run full UMBRA training ───────────────────────────────────────────\n",
    "# Full pipeline (~2h 30min on T4 GPU):\n",
    "#   1. Load 5 HuggingFace datasets            (cached after first run)\n",
    "#   2. Benchmark Sentrix PII guard             → logs/sentrix_benchmark.json\n",
    "#   3. Capture BEFORE metrics (baseline)       → logs/before_metrics.json\n",
    "#   4. GRPO fine-tune SmolLM-135M             (~45 min) → umbra_grpo_ckpt/\n",
    "#   5. Shadow Arms Race (2 rounds × 50 eps)   (~15 min) → logs/arms_race_log.jsonl\n",
    "#   6. RL Loop 500 eps ShapedRewardModel      (~30 min) → logs/rollout_samples.jsonl\n",
    "#   7. Cialdini Stress Test (6 × 10 eps)      (~10 min) → logs/cialdini_results.json\n",
    "#   8. AFTER metrics + 4 graphs generated              → logs/reward_graphs/\n",
    "#   9. Save final model                                → umbra_final/\n",
    "#\n",
    "# ── To run in FAST MODE (~5 min), uncomment these lines BEFORE running: ───────\n",
    "# import re, pathlib\n",
    "# src = pathlib.Path('train.py').read_text()\n",
    "# src = re.sub(r'USE_GRPO\\s*=\\s*True',  'USE_GRPO  = False', src)\n",
    "# src = re.sub(r'USE_SHADOW\\s*=\\s*True','USE_SHADOW = False', src)\n",
    "# src = re.sub(r'EPISODES\\s*=\\s*500',   'EPISODES   = 50',   src)\n",
    "# src = re.sub(r'CIALDINI_EPS\\s*=\\s*10','CIALDINI_EPS = 3',  src)\n",
    "# pathlib.Path('train_fast.py').write_text(src)\n",
    "# %run train_fast.py\n",
    "# ─────────────────────────────────────────────────────────────────────────────\n",
    "\n",
    "%run train.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2d0148a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Cell 5: Evaluate trained model ────────────────────────────────────────────\n",
    "%run evaluate.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf4ab64f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Cell 6: View generated training graphs ────────────────────────────────────\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "from pathlib import Path\n",
    "\n",
    "graph_dir = Path('logs/reward_graphs')\n",
    "graphs    = sorted(graph_dir.glob('*.png'))\n",
    "\n",
    "if not graphs:\n",
    "    print('No graphs found yet — run Cell 4 (train.py) first.')\n",
    "else:\n",
    "    fig, axes = plt.subplots(2, 2, figsize=(18, 12))\n",
    "    fig.patch.set_facecolor('#0d1117')\n",
    "    for ax, path in zip(axes.flat, graphs):\n",
    "        img = mpimg.imread(str(path))\n",
    "        ax.imshow(img)\n",
    "        ax.set_title(path.stem.replace('_', ' ').title(),\n",
    "                     color='white', fontsize=13, pad=8)\n",
    "        ax.axis('off')\n",
    "    # Hide unused subplot if fewer than 4 graphs\n",
    "    for ax in axes.flat[len(graphs):]:\n",
    "        ax.set_visible(False)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "    print(f'✅ Displayed {len(graphs)} graphs from {graph_dir}/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f08ee1f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Cell 7 (optional): Download model + logs to local machine ─────────────────\n",
    "import shutil, os\n",
    "from google.colab import files\n",
    "\n",
    "# Bundle the trained model\n",
    "if os.path.isdir('umbra_final'):\n",
    "    shutil.make_archive('umbra_final', 'zip', 'umbra_final')\n",
    "    files.download('umbra_final.zip')\n",
    "    print('✅ Model download started.')\n",
    "else:\n",
    "    print('⚠ umbra_final/ not found — run Cell 4 first.')\n",
    "\n",
    "# Bundle logs (graphs + metrics + Cialdini results)\n",
    "if os.path.isdir('logs'):\n",
    "    shutil.make_archive('umbra_logs', 'zip', 'logs')\n",
    "    files.download('umbra_logs.zip')\n",
    "    print('✅ Logs download started.')\n",
    "else:\n",
    "    print('⚠ logs/ not found.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b6d67f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Cell 7 (optional): View Sentrix benchmark report ─────────────────────────\n",
    "import json\n",
    "\n",
    "with open('logs/sentrix_benchmark.json') as f:\n",
    "    report = json.load(f)\n",
    "\n",
    "print('=== Sentrix PII Guard Benchmark (vs ai4privacy/pii-masking-300k) ===')\n",
    "print(f\"  True Positives  : {report['true_positives']}\")\n",
    "print(f\"  False Positives : {report['false_positives']}\")\n",
    "print(f\"  False Negatives : {report['false_negatives']}\")\n",
    "print(f\"  Precision       : {report['precision']}\")\n",
    "print(f\"  Recall          : {report['recall']}\")\n",
    "print(f\"  F1 Score        : {report['f1']}\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}