Add Unified LoRA MRPC Benchmark Example Notebook

This notebook demonstrates Unified LoRA on the GLUE MRPC (paraphrase detection) task, including setup, data loading, and evaluation of both baseline and unified LoRA models.

notebooks/mrpc_example.ipynb

@@ -0,0 +1,212 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Unified LoRA - MRPC Benchmark Example\n",
+    "\n",
+    "This notebook demonstrates Unified LoRA on the GLUE MRPC (paraphrase detection) task.\n",
+    "\n",
+    "**Expected results:**\n",
+    "- Baseline LoRA: F1 ~0.78-0.79\n",
+    "- Unified LoRA: F1 ~0.78-0.79 (performance parity with adaptive control)\n",
+    "- φ(t) convergence from 0.5 → ~0.35-0.40\n",
+    "- Mode primarily stays in Multi (1) for stable training"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install -q transformers datasets peft evaluate scikit-learn accelerate"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"WANDB_DISABLED\"] = \"true\"\n",
+    "\n",
+    "import torch\n",
+    "from datasets import load_dataset\n",
+    "from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments\n",
+    "from peft import LoraConfig, get_peft_model\n",
+    "from torch.utils.data import DataLoader\n",
+    "import evaluate\n",
+    "\n",
+    "# Import UnifiedController\n",
+    "import sys\n",
+    "sys.path.append('..')\n",
+    "from controller import UnifiedController\n",
+    "\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"Using device: {device}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load Data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset = load_dataset(\"glue\", \"mrpc\")[\"train\"].train_test_split(test_size=0.2, seed=42)\n",
+    "print(f\"Dataset: {len(dataset['train'])} train | {len(dataset['test'])} test\")\n",
+    "\n",
+    "model_name = \"distilbert-base-uncased\"\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+    "if tokenizer.pad_token is None:\n",
+    "    tokenizer.pad_token = tokenizer.eos_token\n",
+    "\n",
+    "def tokenize(examples):\n",
+    "    return tokenizer(examples['sentence1'], examples['sentence2'], truncation=True, max_length=128, padding=True)\n",
+    "\n",
+    "tokenized_train = dataset['train'].map(tokenize, batched=True).rename_column(\"label\", \"labels\")\n",
+    "tokenized_test = dataset['test'].map(tokenize, batched=True).rename_column(\"label\", \"labels\")\n",
+    "\n",
+    "metric = evaluate.combine([\"accuracy\", \"f1\"])\n",
+    "\n",
+    "def compute_metrics(eval_pred):\n",
+    "    logits, labels = eval_pred\n",
+    "    predictions = torch.argmax(torch.tensor(logits), axis=-1)\n",
+    "    return metric.compute(predictions=predictions, references=labels)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Baseline LoRA"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"🟡 BASELINE LoRA\")\n",
+    "model_base = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)\n",
+    "model_base = get_peft_model(model_base, LoraConfig(r=16, lora_alpha=32, target_modules=[\"q_lin\", \"v_lin\"]))\n",
+    "\n",
+    "trainer_base = Trainer(\n",
+    "    model=model_base,\n",
+    "    train_dataset=tokenized_train,\n",
+    "    eval_dataset=tokenized_test,\n",
+    "    args=TrainingArguments(\n",
+    "        output_dir=\"./baseline\",\n",
+    "        num_train_epochs=3,\n",
+    "        per_device_train_batch_size=16,\n",
+    "        fp16=True,\n",
+    "        eval_strategy=\"epoch\",\n",
+    "        save_strategy=\"no\",\n",
+    "        report_to=None\n",
+    "    ),\n",
+    "    tokenizer=tokenizer,\n",
+    "    compute_metrics=compute_metrics\n",
+    ")\n",
+    "\n",
+    "trainer_base.train()\n",
+    "results_base = trainer_base.evaluate()\n",
+    "print(f\"F1: {results_base['eval_f1']:.3f} | Acc: {results_base['eval_accuracy']:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Unified LoRA"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"🔵 UNIFIED LoRA\")\n",
+    "controller = UnifiedController()\n",
+    "\n",
+    "model_unified = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)\n",
+    "model_unified = get_peft_model(model_unified, LoraConfig(r=16, lora_alpha=32, target_modules=[\"q_lin\", \"v_lin\"]))\n",
+    "model_unified = model_unified.to(device)\n",
+    "\n",
+    "train_loader = DataLoader(tokenized_train.remove_columns(['sentence1', 'sentence2', 'idx']), batch_size=16, shuffle=True)\n",
+    "optimizer = torch.optim.AdamW(model_unified.parameters(), lr=3e-5)\n",
+    "model_unified.train()\n",
+    "\n",
+    "for epoch in range(3):\n",
+    "    print(f\"Epoch {epoch+1}/3\")\n",
+    "    for i, batch in enumerate(train_loader):\n",
+    "        inputs = {k: v.to(device) for k, v in batch.items() if k in ['input_ids', 'attention_mask', 'labels']}\n",
+    "        outputs = model_unified(**inputs)\n",
+    "        new_lr = controller.update(outputs.loss.item())\n",
+    "        for g in optimizer.param_groups: g['lr'] = new_lr\n",
+    "        outputs.loss.backward()\n",
+    "        optimizer.step()\n",
+    "        optimizer.zero_grad()\n",
+    "        if (i+1) % 50 == 0:\n",
+    "            print(f\"  [{controller.step}] φ={controller.phi:.3f} M={controller.mode}\")\n",
+    "\n",
+    "model_unified.eval()\n",
+    "trainer_unified = Trainer(\n",
+    "    model=model_unified, eval_dataset=tokenized_test,\n",
+    "    args=TrainingArguments(output_dir=\"./u\", per_device_eval_batch_size=16, fp16=True, report_to=None),\n",
+    "    tokenizer=tokenizer, compute_metrics=compute_metrics\n",
+    ")\n",
+    "results_unified = trainer_unified.evaluate()\n",
+    "print(f\"F1: {results_unified['eval_f1']:.3f} | Acc: {results_unified['eval_accuracy']:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"\\n📊 COMPARISON\")\n",
+    "print(\"| Method   | F1    | Acc   |\")\n",
+    "print(\"|----------|-------|-------|\")\n",
+    "print(f\"| Baseline | {results_base['eval_f1']:.3f} | {results_base['eval_accuracy']:.3f} |\")\n",
+    "print(f\"| Unified  | {results_unified['eval_f1']:.3f} | {results_unified['eval_accuracy']:.3f} |\")\n",
+    "print(f\"\\nΔF1: {results_unified['eval_f1'] - results_base['eval_f1']:+.3f}\")\n",
+    "print(f\"Final φ: {controller.phi:.3f} | Mode: {controller.mode}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}