Update mrpc_example.ipynb

notebooks/mrpc_example.ipynb

@@ -1,148 +1,135 @@
 {
-"cells": [
-{
-"cell_type": "markdown",
-"metadata": {},
-"source": [
-"# Orbital LoRA - MRPC Benchmark Example\n",
-"\n",
-"Expected: performance parity with baseline + adaptive behavior (no degradation)\n"
-]
-},
-{
-"cell_type": "code",
-"source": [
-"!pip install -q transformers datasets evaluate scikit-learn accelerate"
-]
-},
-{
-"cell_type": "code",
-"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\n",
-"from torch.utils.data import DataLoader\n",
-"import evaluate\n",
-"\n",
-"import sys\n",
-"sys.path.append('..')\n",
-"\n",
-"from nested_lora import inject_nested_lora\n",
-"from orbital_controller import OrbitalController\n",
-"from controller import set_rank\n",
-"\n",
-"device = torch.device("cuda" if torch.cuda.is_available() else "cpu")\n",
-"print(device)"
-]
-},
-{
-"cell_type": "code",
-"source": [
-"dataset = load_dataset("glue", "mrpc")\n",
-"tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")\n",
-"\n",
-"def tok(x):\n",
-"    return tokenizer(x['sentence1'], x['sentence2'], truncation=True, padding="max_length", max_length=128)\n",
-"\n",
-"train = dataset['train'].map(tok, batched=True)\n",
-"val   = dataset['validation'].map(tok, batched=True)\n",
-"\n",
-"train.set_format(type="torch", columns=["input_ids","attention_mask","label"])\n",
-"val.set_format(type="torch", columns=["input_ids","attention_mask","label"])\n",
-"\n",
-"train_loader = DataLoader(train, batch_size=16, shuffle=True)\n",
-"val_loader   = DataLoader(val, batch_size=16)\n",
-"\n",
-"metric = evaluate.load("glue","mrpc")"
-]
-},
-{
-"cell_type": "code",
-"source": [
-"def eval_model(model):\n",
-"    model.eval()\n",
-"    preds, labels = [], []\n",
-"    with torch.no_grad():\n",
-"        for b in val_loader:\n",
-"            x=b['input_ids'].to(device)\n",
-"            m=b['attention_mask'].to(device)\n",
-"            y=b['label'].to(device)\n",
-"            p=model(input_ids=x,attention_mask=m).logits.argmax(-1)\n",
-"            preds.extend(p.cpu().numpy()); labels.extend(y.cpu().numpy())\n",
-"    return metric.compute(predictions=preds,references=labels)['f1']"
-]
-},
-{
-"cell_type": "markdown",
-"source": ["## Baseline"]
-},
-{
-"cell_type": "code",
-"source": [
-"model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)\n",
-"model = inject_nested_lora(model,16).to(device)\n",
-"set_rank(model,16)\n",
-"\n",
-"opt = torch.optim.AdamW(model.parameters(), lr=5e-5)\n",
-"\n",
-"for step,b in enumerate(train_loader):\n",
-"    if step>200: break\n",
-"    x=b['input_ids'].to(device); m=b['attention_mask'].to(device); y=b['label'].to(device)\n",
-"    loss=model(input_ids=x,attention_mask=m,labels=y).loss\n",
-"    loss.backward(); opt.step(); opt.zero_grad()\n",
-"\n",
-"f1_base = eval_model(model)\n",
-"print("Baseline F1:", round(f1_base,3))"
-]
-},
-{
-"cell_type": "markdown",
-"source": ["## Orbital"]
-},
-{
-"cell_type": "code",
-"source": [
-"model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)\n",
-"model = inject_nested_lora(model,16).to(device)\n",
-"\n",
-"ctrl = OrbitalController(warmup=10, stable_window=6)\n",
-"set_rank(model,4)\n",
-"\n",
-"opt = torch.optim.AdamW(model.parameters(), lr=5e-5)\n",
-"\n",
-"for step,b in enumerate(train_loader):\n",
-"    if step>200: break\n",
-"    x=b['input_ids'].to(device); m=b['attention_mask'].to(device); y=b['label'].to(device)\n",
-"    loss=model(input_ids=x,attention_mask=m,labels=y).loss\n",
-"    loss.backward()\n",
-"\n",
-"    r = ctrl.step(loss.item())\n",
-"    r = max(4,min(16,r))\n",
-"    set_rank(model,r)\n",
-"\n",
-"    opt.step(); opt.zero_grad()\n",
-"\n",
-"f1_orb = eval_model(model)\n",
-"print("Orbital F1:", round(f1_orb,3))"
-]
-},
-{
-"cell_type": "markdown",
-"source": ["## Result"]
-},
-{
-"cell_type": "code",
-"source": [
-"print("\nBaseline:", round(f1_base,3))\n",
-"print("Orbital:", round(f1_orb,3))\n",
-"print("Delta:", round(f1_orb-f1_base,3))"
-]
-}
-],
-"metadata": {},
-"nbformat": 4,
-"nbformat_minor": 4
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Orbital LoRA - MRPC Benchmark Example\n",
+    "\n",
+    "**Expected:** performance parity with baseline + adaptive behavior\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "!pip install -q transformers datasets evaluate scikit-learn accelerate"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "import torch\n",
+    "from datasets import load_dataset\n",
+    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
+    "from torch.utils.data import DataLoader\n",
+    "import evaluate\n",
+    "\n",
+    "import sys\n",
+    "sys.path.append('..')\n",
+    "\n",
+    "from nested_lora import inject_nested_lora\n",
+    "from orbital_controller import OrbitalController\n",
+    "from controller import set_rank\n",
+    "\n",
+    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+    "print(device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "dataset = load_dataset('glue','mrpc')\n",
+    "tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')\n",
+    "\n",
+    "def tok(x):\n",
+    "    return tokenizer(x['sentence1'], x['sentence2'], truncation=True, padding='max_length', max_length=128)\n",
+    "\n",
+    "train = dataset['train'].map(tok, batched=True)\n",
+    "val   = dataset['validation'].map(tok, batched=True)\n",
+    "\n",
+    "train.set_format(type='torch', columns=['input_ids','attention_mask','label'])\n",
+    "val.set_format(type='torch', columns=['input_ids','attention_mask','label'])\n",
+    "\n",
+    "train_loader = DataLoader(train, batch_size=16, shuffle=True)\n",
+    "val_loader   = DataLoader(val, batch_size=16)\n",
+    "\n",
+    "metric = evaluate.load('glue','mrpc')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "def eval_model(model):\n",
+    "    model.eval()\n",
+    "    preds, labels = [], []\n",
+    "    with torch.no_grad():\n",
+    "        for b in val_loader:\n",
+    "            x=b['input_ids'].to(device)\n",
+    "            m=b['attention_mask'].to(device)\n",
+    "            y=b['label'].to(device)\n",
+    "            p=model(input_ids=x,attention_mask=m).logits.argmax(-1)\n",
+    "            preds.extend(p.cpu().numpy()); labels.extend(y.cpu().numpy())\n",
+    "    return metric.compute(predictions=preds,references=labels)['f1']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "# BASELINE\n",
+    "model = AutoModelForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=2)\n",
+    "model = inject_nested_lora(model,16).to(device)\n",
+    "set_rank(model,16)\n",
+    "\n",
+    "opt = torch.optim.AdamW(model.parameters(), lr=5e-5)\n",
+    "\n",
+    "for step,b in enumerate(train_loader):\n",
+    "    if step>200: break\n",
+    "    x=b['input_ids'].to(device); m=b['attention_mask'].to(device); y=b['label'].to(device)\n",
+    "    loss=model(input_ids=x,attention_mask=m,labels=y).loss\n",
+    "    loss.backward(); opt.step(); opt.zero_grad()\n",
+    "\n",
+    "f1_base = eval_model(model)\n",
+    "print('Baseline F1:', round(f1_base,3))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "# ORBITAL\n",
+    "model = AutoModelForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=2)\n",
+    "model = inject_nested_lora(model,16).to(device)\n",
+    "\n",
+    "ctrl = OrbitalController(warmup=10, stable_window=6)\n",
+    "set_rank(model,4)\n",
+    "\n",
+    "opt = torch.optim.AdamW(model.parameters(), lr=5e-5)\n",
+    "\n",
+    "for step,b in enumerate(train_loader):\n",
+    "    if step>200: break\n",
+    "    x=b['input_ids'].to(device); m=b['attention_mask'].to(device); y=b['label'].to(device)\n",
+    "    loss=model(input_ids=x,attention_mask=m,labels=y).loss\n",
+    "    loss.backward()\n",
+    "\n",
+    "    r = ctrl.step(loss.item())\n",
+    "    r = max(4,min(16,r))\n",
+    "    set_rank(model,r)\n",
+    "\n",
+    "    opt.step(); opt.zero_grad()\n",
+    "\n",
+    "f1_orb = eval_model(model)\n",
+    "print('Orbital F1:', round(f1_orb,3))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "print('\\nBaseline:', round(f1_base,3))\n",
+    "print('Orbital:', round(f1_orb,3))\n",
+    "print('Delta:', round(f1_orb-f1_base,3))"
+   ]
+  }
+ ],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 4
 }