{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "588c9c40",
   "metadata": {},
   "source": [
    "# EmpathRAG — RoBERTa Emotion Classifier (Day 6)\n",
    "**MSML641 · NLP Course Project**  \n",
    "Fine-tune RoBERTa-base + LoRA on GoEmotions → 5-class emotion taxonomy  \n",
    "Target: Weighted F1 > 0.55 on 5-class test set\n",
    "\n",
    "---\n",
    "**Before running:** Runtime → Change runtime type → **A100** → Save  \n",
    "Cell 1 will hard-stop if the wrong GPU is attached.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cce46322",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── STEP 1: Verify A100 is attached ──────────────────────────────────────────\n",
    "import torch\n",
    "\n",
    "gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"NO GPU\"\n",
    "print(f\"GPU     : {gpu_name}\")\n",
    "print(f\"CUDA    : {torch.cuda.is_available()}\")\n",
    "if torch.cuda.is_available():\n",
    "    vram_gb = torch.cuda.get_device_properties(0).total_memory / 1e9\n",
    "    print(f\"VRAM    : {vram_gb:.1f} GB\")\n",
    "    print(f\"PyTorch : {torch.__version__}\")\n",
    "\n",
    "assert torch.cuda.is_available(), \"No GPU found — switch to A100 runtime before proceeding.\"\n",
    "print(\"\\n✅ GPU check passed.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fdcc2e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── STEP 2: Install missing packages ─────────────────────────────────────────\n",
    "# Strategy: Colab already has numpy 2.x, scikit-learn 1.6+, transformers 4.57+,\n",
    "# torch 2.8+, datasets 2.x. We install ONLY what's missing.\n",
    "# DO NOT pin numpy or scikit-learn — that caused the previous conflicts.\n",
    "\n",
    "!pip install -q --upgrade \\\n",
    "    peft>=0.18.0 \\\n",
    "    evaluate==0.4.6 \\\n",
    "    accelerate>=1.0.0\n",
    "\n",
    "# Verify no numpy conflict\n",
    "import numpy as np\n",
    "import sklearn\n",
    "print(f\"numpy       : {np.__version__}\")\n",
    "print(f\"scikit-learn: {sklearn.__version__}\")\n",
    "\n",
    "import transformers, peft, evaluate as ev, accelerate\n",
    "print(f\"transformers: {transformers.__version__}\")\n",
    "print(f\"peft        : {peft.__version__}\")\n",
    "print(f\"evaluate    : {ev.__version__}\")\n",
    "print(f\"accelerate  : {accelerate.__version__}\")\n",
    "print(\"\\n✅ All packages ready.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0e0abe9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── STEP 3: Mount Drive, set checkpoint dir, detect existing checkpoints ──────\n",
    "from google.colab import drive\n",
    "import os, glob\n",
    "\n",
    "drive.mount(\"/content/drive\")\n",
    "\n",
    "SAVE_DIR = \"/content/drive/MyDrive/empathrag/emotion_classifier\"\n",
    "os.makedirs(SAVE_DIR, exist_ok=True)\n",
    "print(f\"Checkpoint dir: {SAVE_DIR}\")\n",
    "\n",
    "existing = sorted(glob.glob(os.path.join(SAVE_DIR, \"checkpoint-*\")))\n",
    "if existing:\n",
    "    print(f\"\\nFound {len(existing)} checkpoint(s):\")\n",
    "    for c in existing:\n",
    "        print(f\"  {c}\")\n",
    "    RESUME_FROM = existing[-1]\n",
    "    print(f\"\\nWill resume from: {RESUME_FROM}\")\n",
    "else:\n",
    "    RESUME_FROM = None\n",
    "    print(\"\\nNo existing checkpoints — starting fresh.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed2782fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── STEP 4: Load GoEmotions dataset, remap 27 → 5 classes ────────────────────\n",
    "from datasets import load_dataset\n",
    "from collections import Counter\n",
    "\n",
    "# 5-class taxonomy:\n",
    "#   0 = distress   (grief, remorse, fear, sadness)\n",
    "#   1 = anxiety    (nervousness, confusion, embarrassment)\n",
    "#   2 = frustration(anger, annoyance, disappointment, disgust)\n",
    "#   3 = neutral    (neutral)\n",
    "#   4 = hopeful    (optimism, relief, gratitude, joy, love, admiration,\n",
    "#                   amusement, approval, caring, curiosity, desire,\n",
    "#                   excitement, pride, realization, surprise)\n",
    "\n",
    "LABEL_MAP = {\n",
    "    \"grief\": 0, \"remorse\": 0, \"fear\": 0, \"sadness\": 0,\n",
    "    \"nervousness\": 1, \"confusion\": 1, \"embarrassment\": 1,\n",
    "    \"anger\": 2, \"annoyance\": 2, \"disappointment\": 2, \"disgust\": 2,\n",
    "    \"neutral\": 3,\n",
    "    \"optimism\": 4, \"relief\": 4, \"gratitude\": 4, \"joy\": 4,\n",
    "    \"love\": 4, \"admiration\": 4, \"amusement\": 4, \"approval\": 4,\n",
    "    \"caring\": 4, \"curiosity\": 4, \"desire\": 4, \"excitement\": 4,\n",
    "    \"pride\": 4, \"realization\": 4, \"surprise\": 4,\n",
    "}\n",
    "LABEL_NAMES = [\"distress\", \"anxiety\", \"frustration\", \"neutral\", \"hopeful\"]\n",
    "\n",
    "print(\"Loading GoEmotions (simplified split)...\")\n",
    "raw = load_dataset(\"google-research-datasets/go_emotions\", \"simplified\")\n",
    "print(f\"Train: {len(raw['train']):,} | Val: {len(raw['validation']):,} | Test: {len(raw['test']):,}\")\n",
    "\n",
    "# Inspect label format — GoEmotions simplified stores labels as list of ints\n",
    "# that are indices into feature_names\n",
    "feature_names = raw[\"train\"].features[\"labels\"].feature.names\n",
    "print(f\"\\nDataset has {len(feature_names)} emotion classes.\")\n",
    "print(f\"Sample row labels: {raw['train'][0]['labels']} → \"\n",
    "      f\"{[feature_names[i] for i in raw['train'][0]['labels']]}\")\n",
    "\n",
    "def remap(example):\n",
    "    \"\"\"Collapse multi-hot 27-class labels to single 5-class coarse label.\n",
    "    Takes the first label that appears in LABEL_MAP; falls back to neutral (3).\n",
    "    \"\"\"\n",
    "    for lid in example[\"labels\"]:\n",
    "        name = feature_names[lid]\n",
    "        if name in LABEL_MAP:\n",
    "            return {\"label\": LABEL_MAP[name]}\n",
    "    return {\"label\": 3}  # neutral fallback\n",
    "\n",
    "print(\"\\nRemapping labels...\")\n",
    "dataset = raw.map(remap, desc=\"27→5 remap\")\n",
    "\n",
    "# Distribution check\n",
    "dist = Counter(dataset[\"train\"][\"label\"])\n",
    "print(\"\\nTrain class distribution:\")\n",
    "for i, name in enumerate(LABEL_NAMES):\n",
    "    pct = 100 * dist[i] / len(dataset[\"train\"])\n",
    "    print(f\"  {i} {name:<15} {dist[i]:>6,}  ({pct:.1f}%)\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52ddb3a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── STEP 5: Tokenize ─────────────────────────────────────────────────────────\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "MODEL_NAME = \"roberta-base\"\n",
    "print(f\"Loading tokenizer: {MODEL_NAME}\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
    "\n",
    "def tokenize(batch):\n",
    "    return tokenizer(\n",
    "        batch[\"text\"],\n",
    "        truncation=True,\n",
    "        max_length=128,\n",
    "        padding=\"max_length\",\n",
    "    )\n",
    "\n",
    "print(\"Tokenizing (this takes ~2 min)...\")\n",
    "tokenized = dataset.map(tokenize, batched=True, desc=\"Tokenizing\")\n",
    "tokenized = tokenized.rename_column(\"label\", \"labels\")\n",
    "tokenized.set_format(type=\"torch\", columns=[\"input_ids\", \"attention_mask\", \"labels\"])\n",
    "\n",
    "print(f\"\\nTokenization complete.\")\n",
    "print(f\"  Train : {len(tokenized['train']):,}\")\n",
    "print(f\"  Val   : {len(tokenized['validation']):,}\")\n",
    "print(f\"  Test  : {len(tokenized['test']):,}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c985e81",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── STEP 6: RoBERTa + LoRA model ─────────────────────────────────────────────\n",
    "import torch\n",
    "from transformers import AutoModelForSequenceClassification\n",
    "from peft import get_peft_model, LoraConfig, TaskType\n",
    "\n",
    "id2label = {i: n for i, n in enumerate(LABEL_NAMES)}\n",
    "label2id = {n: i for i, n in enumerate(LABEL_NAMES)}\n",
    "\n",
    "# LoRA config: targets query + value projections in RoBERTa attention\n",
    "# r=16 → ~3M trainable params vs ~125M full fine-tune\n",
    "lora_cfg = LoraConfig(\n",
    "    task_type=TaskType.SEQ_CLS,\n",
    "    r=16,\n",
    "    lora_alpha=32,\n",
    "    lora_dropout=0.1,\n",
    "    target_modules=[\"query\", \"value\"],\n",
    "    bias=\"none\",\n",
    ")\n",
    "\n",
    "base_model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    MODEL_NAME,\n",
    "    num_labels=5,\n",
    "    id2label=id2label,\n",
    "    label2id=label2id,\n",
    ")\n",
    "model = get_peft_model(base_model, lora_cfg)\n",
    "model.print_trainable_parameters()\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"\\nTraining on: {device}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c43de57",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── STEP 7: Evaluation metrics ───────────────────────────────────────────────\n",
    "import evaluate\n",
    "import numpy as np\n",
    "\n",
    "f1_metric = evaluate.load(\"f1\")\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    logits, labels = eval_pred\n",
    "    preds = np.argmax(logits, axis=-1)\n",
    "\n",
    "    # Weighted F1 — primary metric (handles class imbalance)\n",
    "    weighted_f1 = f1_metric.compute(\n",
    "        predictions=preds, references=labels, average=\"weighted\"\n",
    "    )[\"f1\"]\n",
    "\n",
    "    # Per-class F1 — required for paper's breakdown table\n",
    "    per_class = f1_metric.compute(\n",
    "        predictions=preds, references=labels,\n",
    "        average=None, labels=list(range(5))\n",
    "    )[\"f1\"]\n",
    "\n",
    "    result = {\"f1_weighted\": weighted_f1}\n",
    "    for i, name in enumerate(LABEL_NAMES):\n",
    "        result[f\"f1_{name}\"] = float(per_class[i])\n",
    "    return result\n",
    "\n",
    "print(\"Metrics function ready (weighted F1 + per-class F1 for all 5 classes).\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70116ba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── STEP 8: Training configuration ───────────────────────────────────────────\n",
    "from transformers import TrainingArguments\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=SAVE_DIR,\n",
    "\n",
    "    # Schedule\n",
    "    num_train_epochs=5,\n",
    "    per_device_train_batch_size=64,    # A100 40GB handles this fine at max_len=128\n",
    "    per_device_eval_batch_size=128,\n",
    "    learning_rate=2e-4,                # Higher than full fine-tune — correct for LoRA\n",
    "    warmup_ratio=0.1,\n",
    "    weight_decay=0.01,\n",
    "\n",
    "    # Checkpointing\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"f1_weighted\",\n",
    "    greater_is_better=True,\n",
    "    save_total_limit=3,                # Keep last 3 to save Drive space\n",
    "\n",
    "    # Speed\n",
    "    fp16=True,\n",
    "    dataloader_num_workers=4,\n",
    "\n",
    "    # Logging\n",
    "    logging_dir=f\"{SAVE_DIR}/logs\",\n",
    "    logging_steps=200,\n",
    "    report_to=\"none\",\n",
    ")\n",
    "\n",
    "print(\"TrainingArguments set:\")\n",
    "print(f\"  Epochs        : {training_args.num_train_epochs}\")\n",
    "print(f\"  Train batch   : {training_args.per_device_train_batch_size}\")\n",
    "print(f\"  Learning rate : {training_args.learning_rate}\")\n",
    "print(f\"  Best metric   : {training_args.metric_for_best_model}\")\n",
    "print(f\"  Save dir      : {SAVE_DIR}\")\n",
    "print(f\"  Resume from   : {RESUME_FROM}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d9a3ee0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── STEP 9: Train ────────────────────────────────────────────────────────────\n",
    "# Expected time on A100: ~75–90 minutes\n",
    "# Checkpoints saved to Drive after every epoch — safe to disconnect and resume.\n",
    "from transformers import Trainer\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized[\"train\"],\n",
    "    eval_dataset=tokenized[\"validation\"],\n",
    "    compute_metrics=compute_metrics,\n",
    ")\n",
    "\n",
    "print(\"Starting training...\")\n",
    "print(f\"Resuming from: {RESUME_FROM}\" if RESUME_FROM else \"Fresh training run.\")\n",
    "print(\"-\" * 60)\n",
    "\n",
    "trainer.train(resume_from_checkpoint=RESUME_FROM)\n",
    "\n",
    "print(\"-\" * 60)\n",
    "print(\"Training complete.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "173bf559",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── STEP 10: Save final model + tokenizer to Drive ───────────────────────────\n",
    "import os\n",
    "\n",
    "trainer.save_model(SAVE_DIR)\n",
    "tokenizer.save_pretrained(SAVE_DIR)\n",
    "\n",
    "print(f\"Model saved to: {SAVE_DIR}\")\n",
    "print(\"\\nContents:\")\n",
    "for f in sorted(os.listdir(SAVE_DIR)):\n",
    "    fpath = os.path.join(SAVE_DIR, f)\n",
    "    size  = os.path.getsize(fpath) if os.path.isfile(fpath) else 0\n",
    "    label = f\"{size/1e6:.1f} MB\" if size > 0 else \"dir\"\n",
    "    print(f\"  {f:<45} {label}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0258fa90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── STEP 11: Evaluate on held-out test set ───────────────────────────────────\n",
    "import json\n",
    "import numpy as np\n",
    "\n",
    "print(\"Evaluating on test set (5,427 examples)...\")\n",
    "test_results = trainer.evaluate(tokenized[\"test\"])\n",
    "\n",
    "print(\"\\n=== TEST SET RESULTS ===\")\n",
    "weighted_f1 = test_results[\"eval_f1_weighted\"]\n",
    "print(f\"Weighted F1 : {weighted_f1:.4f}  (target: > 0.55)\")\n",
    "print()\n",
    "print(\"Per-class F1:\")\n",
    "for name in LABEL_NAMES:\n",
    "    val = test_results.get(f\"eval_f1_{name}\", float(\"nan\"))\n",
    "    print(f\"  {name:<15} {val:.4f}\")\n",
    "\n",
    "print()\n",
    "if weighted_f1 >= 0.55:\n",
    "    print(\"✅ PASS — weighted F1 exceeds target of 0.55\")\n",
    "else:\n",
    "    print(\"⚠️  BELOW TARGET\")\n",
    "    print(\"   Fallback option: lower lr to 1e-4, increase epochs to 8, re-run from checkpoint.\")\n",
    "\n",
    "# Save to Drive for the paper\n",
    "results_path = f\"{SAVE_DIR}/test_results.json\"\n",
    "with open(results_path, \"w\") as f:\n",
    "    json.dump({k: float(v) for k, v in test_results.items()}, f, indent=2)\n",
    "print(f\"\\nResults saved: {results_path}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c47dff9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── STEP 12: Confusion matrix + classification report ────────────────────────\n",
    "from sklearn.metrics import confusion_matrix, classification_report\n",
    "import numpy as np\n",
    "\n",
    "print(\"Computing confusion matrix on test set...\")\n",
    "pred_output = trainer.predict(tokenized[\"test\"])\n",
    "preds  = np.argmax(pred_output.predictions, axis=-1)\n",
    "labels = pred_output.label_ids\n",
    "\n",
    "print(\"\\n=== CLASSIFICATION REPORT ===\")\n",
    "print(classification_report(labels, preds, target_names=LABEL_NAMES))\n",
    "\n",
    "cm = confusion_matrix(labels, preds)\n",
    "print(\"=== CONFUSION MATRIX (rows=true, cols=pred) ===\")\n",
    "header = f\"{'':>15}\" + \"\".join(f\"{n:>13}\" for n in LABEL_NAMES)\n",
    "print(header)\n",
    "for i, row in enumerate(cm):\n",
    "    print(f\"{LABEL_NAMES[i]:>15}\" + \"\".join(f\"{v:>13}\" for v in row))\n",
    "\n",
    "np.save(f\"{SAVE_DIR}/confusion_matrix.npy\", cm)\n",
    "print(f\"\\nConfusion matrix saved to Drive.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc6f3e30",
   "metadata": {},
   "source": [
    "## ✅ Training Complete — Day 10 Checklist\n",
    "\n",
    "### 1. Download checkpoint from Drive\n",
    "Navigate in Google Drive to:  \n",
    "`My Drive/empathrag/emotion_classifier/`\n",
    "\n",
    "Download the entire folder and place it at:  \n",
    "`models/emotion_classifier/` in your local repo.\n",
    "\n",
    "Required files:\n",
    "- `adapter_config.json`\n",
    "- `adapter_model.safetensors` (or `.bin`) — ~12 MB LoRA weights\n",
    "- `config.json`\n",
    "- `tokenizer.json`, `tokenizer_config.json`, `vocab.json`, `merges.txt`\n",
    "- `test_results.json` — paste weighted F1 into paper\n",
    "- `confusion_matrix.npy` — per-class breakdown for paper table\n",
    "\n",
    "### 2. Run corpus annotation (Day 10, locally)\n",
    "```bash\n",
    "python src/models/annotate_corpus.py\n",
    "```\n",
    "Updates all 1,674,369 SQLite rows with predicted emotion labels (~45 min on CPU).\n",
    "\n",
    "### 3. Verify the checkpoint loads locally\n",
    "```python\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "from peft import PeftModel\n",
    "\n",
    "tok = AutoTokenizer.from_pretrained(\"models/emotion_classifier\")\n",
    "base = AutoModelForSequenceClassification.from_pretrained(\"roberta-base\", num_labels=5)\n",
    "model = PeftModel.from_pretrained(base, \"models/emotion_classifier\").eval()\n",
    "print(\"✅ Checkpoint loads correctly.\")\n",
    "```\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "name": "colab_emotion_classifier.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}