Add Colab notebooks: emotion classifier + corpus annotation (Day 12)

notebooks/colab_annotate_corpus.ipynb

@@ -0,0 +1,307 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "e11600f1",
+   "metadata": {},
+   "source": [
+    "# EmpathRAG — Corpus Annotation (Day 10)\n",
+    "Annotate all 1,674,369 FAISS chunks with emotion labels using the trained RoBERTa checkpoint.  \n",
+    "**Runtime: A100 GPU. Expected time: ~25 minutes.**\n",
+    "\n",
+    "This notebook:\n",
+    "1. Copies `metadata.db` from your local machine (uploaded here) to Colab\n",
+    "2. Loads the RoBERTa+LoRA checkpoint from Drive\n",
+    "3. Runs inference in batches of 512 on GPU\n",
+    "4. Writes emotion labels back to the SQLite db\n",
+    "5. Saves the annotated db back to Drive for you to download\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e85f0c1f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"NO GPU\"\n",
+    "print(f\"GPU  : {gpu}\")\n",
+    "print(f\"VRAM : {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB\")\n",
+    "assert torch.cuda.is_available(), \"Switch to A100 runtime.\"\n",
+    "print(\"✅ GPU ready.\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "90dcada6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install -q peft>=0.18.0 accelerate>=1.0.0\n",
+    "import peft, accelerate\n",
+    "print(f\"peft: {peft.__version__}  |  accelerate: {accelerate.__version__}\")\n",
+    "print(\"✅ Packages ready.\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "51d42310",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from google.colab import drive\n",
+    "import os, shutil\n",
+    "\n",
+    "drive.mount(\"/content/drive\")\n",
+    "\n",
+    "CHECKPOINT_DIR = \"/content/drive/MyDrive/empathrag/emotion_classifier\"\n",
+    "DB_DRIVE_PATH  = \"/content/drive/MyDrive/empathrag/metadata.db\"\n",
+    "DB_LOCAL_PATH  = \"/content/metadata.db\"\n",
+    "\n",
+    "# Verify checkpoint exists\n",
+    "required = [\"adapter_config.json\", \"adapter_model.safetensors\", \"tokenizer.json\"]\n",
+    "for f in required:\n",
+    "    path = os.path.join(CHECKPOINT_DIR, f)\n",
+    "    assert os.path.exists(path), f\"Missing: {path}\"\n",
+    "print(f\"✅ Checkpoint found at: {CHECKPOINT_DIR}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f6533428",
+   "metadata": {},
+   "source": [
+    "## ⬆️ Upload metadata.db\n",
+    "\n",
+    "Run the cell below. It will prompt you to upload a file.  \n",
+    "Upload `data/indexes/metadata.db` from your local machine.\n",
+    "\n",
+    "The file is ~1.26 GB — upload will take 2-5 minutes depending on your connection.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8f3928e3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from google.colab import files\n",
+    "import shutil, os\n",
+    "\n",
+    "print(\"Click 'Choose Files' and select data/indexes/metadata.db from your local machine.\")\n",
+    "print(\"File size: ~1.26 GB. Wait for upload to complete before proceeding.\")\n",
+    "\n",
+    "uploaded = files.upload()\n",
+    "\n",
+    "# Move to working path\n",
+    "for fname in uploaded:\n",
+    "    dest = \"/content/metadata.db\"\n",
+    "    if fname != \"metadata.db\":\n",
+    "        os.rename(fname, dest)\n",
+    "    else:\n",
+    "        pass  # already at /content/metadata.db if uploaded directly\n",
+    "\n",
+    "# Verify\n",
+    "import sqlite3\n",
+    "conn = sqlite3.connect(DB_LOCAL_PATH)\n",
+    "total      = conn.execute(\"SELECT COUNT(*) FROM chunks\").fetchone()[0]\n",
+    "annotated  = conn.execute(\"SELECT COUNT(*) FROM chunks WHERE emotion_label != -1\").fetchone()[0]\n",
+    "unannotated= conn.execute(\"SELECT COUNT(*) FROM chunks WHERE emotion_label = -1\").fetchone()[0]\n",
+    "conn.close()\n",
+    "\n",
+    "print(f\"\\nDB loaded:\")\n",
+    "print(f\"  Total rows   : {total:,}\")\n",
+    "print(f\"  Annotated    : {annotated:,}  (already done — will skip)\")\n",
+    "print(f\"  Unannotated  : {unannotated:,}  (will process these)\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "444f0fea",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
+    "from peft import PeftModel\n",
+    "\n",
+    "LABEL_NAMES = [\"distress\", \"anxiety\", \"frustration\", \"neutral\", \"hopeful\"]\n",
+    "SAFETY_SCORE_MAP = {0: 0.0, 1: 0.0, 2: 0.3, 3: 0.7, 4: 1.0}\n",
+    "DEVICE = torch.device(\"cuda\")\n",
+    "\n",
+    "print(\"Loading RoBERTa + LoRA checkpoint onto GPU...\")\n",
+    "tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR)\n",
+    "base = AutoModelForSequenceClassification.from_pretrained(\"roberta-base\", num_labels=5)\n",
+    "model = PeftModel.from_pretrained(base, CHECKPOINT_DIR).to(DEVICE).eval()\n",
+    "\n",
+    "# Quick sanity check\n",
+    "enc = tokenizer(\"I feel completely overwhelmed and hopeless\",\n",
+    "                return_tensors=\"pt\", truncation=True, max_length=128).to(DEVICE)\n",
+    "with torch.no_grad():\n",
+    "    pred = model(**enc).logits.argmax(-1).item()\n",
+    "print(f\"Sanity check: 'overwhelmed and hopeless' → {LABEL_NAMES[pred]}\")\n",
+    "assert LABEL_NAMES[pred] == \"distress\", f\"Expected distress, got {LABEL_NAMES[pred]}\"\n",
+    "print(\"✅ Model loaded and verified on GPU.\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9e8ef3c2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sqlite3\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "BATCH_SIZE = 512   # A100 can handle 512 at max_len=128 comfortably\n",
+    "\n",
+    "conn = sqlite3.connect(DB_LOCAL_PATH)\n",
+    "\n",
+    "# Only fetch rows not yet annotated — resumes safely if interrupted\n",
+    "rows = conn.execute(\n",
+    "    \"SELECT id, text FROM chunks WHERE emotion_label = -1\"\n",
+    ").fetchall()\n",
+    "\n",
+    "print(f\"Rows to annotate: {len(rows):,}\")\n",
+    "if len(rows) == 0:\n",
+    "    print(\"Nothing to do — all rows already annotated.\")\n",
+    "else:\n",
+    "    updates = []\n",
+    "\n",
+    "    for i in tqdm(range(0, len(rows), BATCH_SIZE), desc=\"Annotating\"):\n",
+    "        batch  = rows[i : i + BATCH_SIZE]\n",
+    "        ids    = [r[0] for r in batch]\n",
+    "        texts  = [r[1] for r in batch]\n",
+    "\n",
+    "        enc = tokenizer(\n",
+    "            texts,\n",
+    "            truncation=True,\n",
+    "            max_length=128,\n",
+    "            padding=True,\n",
+    "            return_tensors=\"pt\"\n",
+    "        ).to(DEVICE)\n",
+    "\n",
+    "        with torch.no_grad():\n",
+    "            logits = model(**enc).logits\n",
+    "\n",
+    "        labels = logits.argmax(-1).tolist()\n",
+    "\n",
+    "        for rid, lbl in zip(ids, labels):\n",
+    "            score = SAFETY_SCORE_MAP[lbl]\n",
+    "            updates.append((lbl, score, rid))\n",
+    "\n",
+    "        # Commit every 50 batches (~25,600 rows) to avoid losing work if disconnected\n",
+    "        if (i // BATCH_SIZE) % 50 == 0 and updates:\n",
+    "            conn.executemany(\n",
+    "                \"UPDATE chunks SET emotion_label=?, safety_score=? WHERE id=?\",\n",
+    "                updates\n",
+    "            )\n",
+    "            conn.commit()\n",
+    "            updates = []\n",
+    "\n",
+    "    # Final commit\n",
+    "    if updates:\n",
+    "        conn.executemany(\n",
+    "            \"UPDATE chunks SET emotion_label=?, safety_score=? WHERE id=?\",\n",
+    "            updates\n",
+    "        )\n",
+    "        conn.commit()\n",
+    "\n",
+    "print(\"\\n✅ Annotation complete.\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f434d2b2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from collections import Counter\n",
+    "\n",
+    "total      = conn.execute(\"SELECT COUNT(*) FROM chunks\").fetchone()[0]\n",
+    "annotated  = conn.execute(\"SELECT COUNT(*) FROM chunks WHERE emotion_label != -1\").fetchone()[0]\n",
+    "unannotated= conn.execute(\"SELECT COUNT(*) FROM chunks WHERE emotion_label = -1\").fetchone()[0]\n",
+    "\n",
+    "print(f\"Total rows   : {total:,}\")\n",
+    "print(f\"Annotated    : {annotated:,}\")\n",
+    "print(f\"Unannotated  : {unannotated:,}  ← must be 0\")\n",
+    "\n",
+    "rows = conn.execute(\"SELECT emotion_label FROM chunks\").fetchall()\n",
+    "dist = Counter(r[0] for r in rows)\n",
+    "print(\"\\nEmotion distribution:\")\n",
+    "for i, name in enumerate(LABEL_NAMES):\n",
+    "    pct = 100 * dist[i] / total\n",
+    "    print(f\"  {i} {name:<15} {dist[i]:>9,}  ({pct:.1f}%)\")\n",
+    "\n",
+    "conn.close()\n",
+    "assert unannotated == 0, f\"{unannotated:,} rows still unannotated!\"\n",
+    "print(\"\\n✅ All rows annotated.\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "945552e5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import shutil\n",
+    "\n",
+    "print(f\"Copying annotated DB to Drive: {DB_DRIVE_PATH}\")\n",
+    "os.makedirs(os.path.dirname(DB_DRIVE_PATH), exist_ok=True)\n",
+    "shutil.copy2(DB_LOCAL_PATH, DB_DRIVE_PATH)\n",
+    "\n",
+    "size_mb = os.path.getsize(DB_DRIVE_PATH) / 1e6\n",
+    "print(f\"Saved: {DB_DRIVE_PATH}  ({size_mb:.1f} MB)\")\n",
+    "print(\"\\n✅ Done. Download metadata.db from Drive and replace data/indexes/metadata.db locally.\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ae95d9cd",
+   "metadata": {},
+   "source": [
+    "## ✅ Annotation Complete — Download Steps\n",
+    "\n",
+    "1. In Google Drive, navigate to `My Drive/empathrag/metadata.db`\n",
+    "2. Right-click → Download\n",
+    "3. Replace your local `data/indexes/metadata.db` with the downloaded file\n",
+    "4. Verify locally:\n",
+    "```python\n",
+    "import sqlite3\n",
+    "from collections import Counter\n",
+    "conn = sqlite3.connect('data/indexes/metadata.db')\n",
+    "unannotated = conn.execute('SELECT COUNT(*) FROM chunks WHERE emotion_label = -1').fetchone()[0]\n",
+    "print(f'Unannotated: {unannotated}')  # must be 0\n",
+    "conn.close()\n",
+    "```\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "A100",
+   "name": "colab_annotate_corpus.ipynb",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

notebooks/colab_emotion_classifier.ipynb

@@ -0,0 +1,499 @@
+{
+ "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
+}