LayerInteractions_Gumini_Demo_Commented.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": [],
   "gpuType": "T4"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cell 1 — Install `LayerInteractions` directly from Hugging Face\n",
    "\n",
    "**What the next cell does**\n",
    "\n",
    "This cell installs the public `LayerInteractions` Python package directly from the Hugging Face repository `Angshul/LayerInteractions`. It uses `pip` with the repository's Git URL, so the code being tested is the same code that an external user would obtain from Hugging Face rather than a local copy.\n",
    "\n",
    "**Why this cell is here**\n",
    "\n",
    "The purpose is to validate the distribution path itself: if this installation succeeds, the Hugging Face repository is installable as a normal Python package. This is the only package-installation step for `LayerInteractions` in the notebook.\n",
    "\n",
    "**What it does not do**\n",
    "\n",
    "It does **not** load a language model, allocate the Gumini model on the GPU, prune any layers, or run an experiment. It only installs the package and its declared dependencies. No runtime restart is required by this notebook unless Colab itself reports an unrelated dependency conflict."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b2BcOsh3EcwF"
   },
   "outputs": [],
   "source": [
    "!pip install \"git+https://huggingface.co/Angshul/LayerInteractions\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cell 2 — Verify that the public pruning API imports correctly\n",
    "\n",
    "**What the next cell does**\n",
    "\n",
    "This cell imports `Order2Pruner`, the main user-facing class of `LayerInteractions`, and prints a short confirmation message.\n",
    "\n",
    "**Why this matters**\n",
    "\n",
    "Successful installation alone does not prove that the package exports its public API correctly. This import checks that Python can find the installed package and that `Order2Pruner` is exposed from `layer_interactions` exactly as intended for users.\n",
    "\n",
    "**Computational cost**\n",
    "\n",
    "Negligible. The cell does **not** instantiate a model and therefore does not download Gumini or consume meaningful GPU memory."
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "from layer_interactions import Order2Pruner\n",
    "print(\"LayerInteractions import OK\")"
   ],
   "metadata": {
    "id": "cf3WKk3BFHQz"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cell 3 — Verify that the result object is also part of the public API\n",
    "\n",
    "**What the next cell does**\n",
    "\n",
    "This cell imports `Order2Result`, the object used to store the measurements and selections produced by the second-order pruning procedure. It then prints the class so that the notebook confirms both major public objects are available from the installed package.\n",
    "\n",
    "**Role of `Order2Result`**\n",
    "\n",
    "During the real experiment, this object records the baseline loss, single-layer deletion losses, pairwise deletion losses, derived first- and second-order interaction terms, and the greedy deletion path. Its `select(...)` method is later used to obtain the retained and deleted layer indices for a requested target depth.\n",
    "\n",
    "**Computational cost**\n",
    "\n",
    "Again, negligible. No model is downloaded or loaded in this cell. The actual compression experiment starts in the following cell."
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "from layer_interactions import Order2Result\n",
    "\n",
    "print(\"Order2Result import OK\")\n",
    "print(Order2Result)"
   ],
   "metadata": {
    "id": "GpjE59UOFZo1"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cell 4 — Run the actual Order-2 layer-interaction experiment: Gumini 10 → 5\n",
    "\n",
    "**What the next cell does**\n",
    "\n",
    "This is the main compression cell. It installs the Hugging Face `datasets` package if needed, loads a small fixed calibration subset from WikiText-2, loads `GuminiResearch/Gumini-1B-Base`, and runs the `LayerInteractions` second-order pruning procedure. The original model has **10 transformer layers**, and the experiment asks for enough measurements to support deletion of 5 layers, leaving a **5-layer model**.\n",
    "\n",
    "**Calibration data**\n",
    "\n",
    "The cell uses the WikiText-2 **training split** only for calibration. It keeps the first 64 nontrivial text entries whose stripped length is at least 100 characters. The held-out test split is deliberately reserved for the next cell so that the final evaluation is not performed on the calibration data.\n",
    "\n",
    "**What `fit(...)` measures**\n",
    "\n",
    "For a 10-layer model, the second-order model needs the unpruned baseline, the effect of deleting each individual layer, and every pairwise layer deletion. There are `C(10,2) = 45` layer pairs. From these measurements, the package constructs the second-order interaction surrogate and a greedy deletion path. With `max_delete=5`, it computes the path far enough to obtain a 10 → 5 selection.\n",
    "\n",
    "**Checkpointing**\n",
    "\n",
    "Measurements are written to `/content/gumini_order2.json`. This allows the expensive measurements to be reused if the fitting process is interrupted while the same Colab runtime/storage remains available.\n",
    "\n",
    "**Important distinction**\n",
    "\n",
    "`fit(...)` computes the measurements and the layer selection, but it does **not yet physically remove the selected layers from the in-memory model**. The final line prints the five layers that will be retained. Physical pruning is performed only in the evaluation cell that follows.\n",
    "\n",
    "**Expected scale**\n",
    "\n",
    "Gumini's weights are about 2.16 GB on Hugging Face and the model has only 10 transformer layers, so this is deliberately chosen as a quick end-to-end demonstration of the package on a T4-class GPU without the hundreds of pair evaluations required by 28–32 layer models."
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "# Cell 3 — LayerInteractions: Gumini-1B-Base, 10 -> 5\n",
    "\n",
    "!pip install -q datasets\n",
    "\n",
    "from datasets import load_dataset\n",
    "from layer_interactions import Order2Pruner\n",
    "\n",
    "MODEL_ID = \"GuminiResearch/Gumini-1B-Base\"\n",
    "\n",
    "# Fixed, small calibration corpus\n",
    "ds = load_dataset(\n",
    "    \"Salesforce/wikitext\",\n",
    "    \"wikitext-2-raw-v1\",\n",
    "    split=\"train\"\n",
    ")\n",
    "\n",
    "texts = [\n",
    "    x[\"text\"].strip()\n",
    "    for x in ds\n",
    "    if len(x[\"text\"].strip()) >= 100\n",
    "][:64]\n",
    "\n",
    "print(f\"Calibration texts: {len(texts)}\")\n",
    "print(f\"Model: {MODEL_ID}\")\n",
    "print(\"Compression: 10 -> 5 layers\")\n",
    "print(\"Order-2 pairs: 45\")\n",
    "\n",
    "pruner = Order2Pruner(MODEL_ID)\n",
    "\n",
    "result = pruner.fit(\n",
    "    texts=texts,\n",
    "    checkpoint_path=\"/content/gumini_order2.json\",\n",
    "    max_delete=5,\n",
    ")\n",
    "\n",
    "print(\"\\nDONE\")\n",
    "print(\"Retained layers:\", result.select(5))"
   ],
   "metadata": {
    "id": "EIzouN4bGcqq"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cell 5 — Evaluate the original model, apply the computed pruning decision, and evaluate the 5-layer model\n",
    "\n",
    "**What the next cell does**\n",
    "\n",
    "This cell performs the final held-out comparison. It uses WikiText-2's **test split**, not the training/calibration split. The same fixed set of 64 nontrivial test texts is used for both the original and compressed models. The metric is token-level negative log-likelihood (NLL), with perplexity (PPL) reported as `exp(NLL)`.\n",
    "\n",
    "**Step 1: evaluate the untouched 10-layer model**\n",
    "\n",
    "Before modifying the model, the cell measures the baseline NLL and perplexity. This establishes the reference performance on exactly the text that will later be used for the compressed model.\n",
    "\n",
    "**Step 2: recover the already-computed selection**\n",
    "\n",
    "`result.select(5)` reads the result from Cell 4 and reports the retained and deleted layer indices. No interaction measurements are recomputed here.\n",
    "\n",
    "**Step 3: physically prune the model**\n",
    "\n",
    "`pruner.apply(5, result=result)` applies the already-computed 10 → 5 decision to the in-memory Hugging Face model. After this call, `pruner.model` contains only five transformer layers. This is a real structural reduction, not merely masking or reporting indices.\n",
    "\n",
    "**Step 4: evaluate the compressed model**\n",
    "\n",
    "The now-5-layer model is evaluated on the exact same held-out text. The cell reports original and compressed NLL, the NLL increase, original and compressed perplexity, and the perplexity ratio.\n",
    "\n",
    "**Important: this cell mutates the model in place**\n",
    "\n",
    "Run this cell **once after Cell 4**. After `pruner.apply(...)`, the object in `pruner.model` is already the 5-layer model. Re-running only this evaluation cell would therefore no longer start from the original 10-layer baseline. To repeat the complete comparison, rerun the experiment cell first so that a fresh 10-layer model is loaded.\n",
    "\n",
    "**Result from the demonstrated run**\n",
    "\n",
    "The previously observed run retained layers `[0, 1, 2, 7, 9]`. On 9,061 held-out tokens, the original model obtained NLL 4.6452 / PPL 104.08, while the 5-layer model obtained NLL 7.4391 / PPL 1701.28. That result demonstrates that the API works end to end; it also shows that this aggressive 50% pruning point substantially degrades this particular model."
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "# Cell 4 — FINAL evaluation\n",
    "# Uses the already-computed Order-2 result from Cell 3.\n",
    "# NO refitting. NO pair calculations.\n",
    "\n",
    "import math\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from datasets import load_dataset\n",
    "\n",
    "# ------------------------------------------------------------\n",
    "# Held-out WikiText-2 test set\n",
    "# ------------------------------------------------------------\n",
    "\n",
    "test_ds = load_dataset(\n",
    "    \"Salesforce/wikitext\",\n",
    "    \"wikitext-2-raw-v1\",\n",
    "    split=\"test\"\n",
    ")\n",
    "\n",
    "test_texts = [\n",
    "    x[\"text\"].strip()\n",
    "    for x in test_ds\n",
    "    if len(x[\"text\"].strip()) >= 100\n",
    "][:64]\n",
    "\n",
    "tokenizer = pruner.tokenizer\n",
    "\n",
    "\n",
    "def evaluate_nll(model, texts, max_length=512):\n",
    "    model.eval()\n",
    "\n",
    "    total_loss = 0.0\n",
    "    total_tokens = 0\n",
    "\n",
    "    device = next(model.parameters()).device\n",
    "\n",
    "    for text in texts:\n",
    "        enc = tokenizer(\n",
    "            text,\n",
    "            return_tensors=\"pt\",\n",
    "            truncation=True,\n",
    "            max_length=max_length\n",
    "        )\n",
    "\n",
    "        input_ids = enc[\"input_ids\"].to(device)\n",
    "\n",
    "        if input_ids.shape[1] < 2:\n",
    "            continue\n",
    "\n",
    "        with torch.no_grad():\n",
    "            logits = model(input_ids=input_ids).logits\n",
    "\n",
    "        shift_logits = logits[:, :-1, :].contiguous()\n",
    "        shift_labels = input_ids[:, 1:].contiguous()\n",
    "\n",
    "        loss = F.cross_entropy(\n",
    "            shift_logits.view(-1, shift_logits.size(-1)),\n",
    "            shift_labels.view(-1),\n",
    "            reduction=\"sum\"\n",
    "        )\n",
    "\n",
    "        total_loss += loss.item()\n",
    "        total_tokens += shift_labels.numel()\n",
    "\n",
    "    nll = total_loss / total_tokens\n",
    "    ppl = math.exp(nll)\n",
    "\n",
    "    return nll, ppl, total_tokens\n",
    "\n",
    "\n",
    "# ============================================================\n",
    "# ORIGINAL MODEL\n",
    "# ============================================================\n",
    "\n",
    "print(\"Evaluating original 10-layer model...\")\n",
    "\n",
    "base_nll, base_ppl, n_tokens = evaluate_nll(\n",
    "    pruner.model,\n",
    "    test_texts\n",
    ")\n",
    "\n",
    "selection = result.select(5)\n",
    "\n",
    "print(f\"Baseline NLL : {base_nll:.4f}\")\n",
    "print(f\"Baseline PPL : {base_ppl:.2f}\")\n",
    "print(f\"Tokens       : {n_tokens:,}\")\n",
    "print(\"Retained     :\", selection[\"retained_layers\"])\n",
    "print(\"Deleted      :\", selection[\"deleted_layers\"])\n",
    "\n",
    "\n",
    "# ============================================================\n",
    "# APPLY THE ALREADY-COMPUTED 10 -> 5 SELECTION\n",
    "# Correct LayerInteractions API signature:\n",
    "#     apply(target_layers, *, result=...)\n",
    "# ============================================================\n",
    "\n",
    "pruner.apply(5, result=result)\n",
    "\n",
    "print(\"\\nApplied retained-layer selection.\")\n",
    "print(\"Physical depth:\", len(pruner.model.model.layers))\n",
    "\n",
    "\n",
    "# ============================================================\n",
    "# COMPRESSED MODEL\n",
    "# ============================================================\n",
    "\n",
    "print(\"\\nEvaluating compressed 5-layer model...\")\n",
    "\n",
    "comp_nll, comp_ppl, _ = evaluate_nll(\n",
    "    pruner.model,\n",
    "    test_texts\n",
    ")\n",
    "\n",
    "\n",
    "# ============================================================\n",
    "# RESULTS\n",
    "# ============================================================\n",
    "\n",
    "print(\"\\n\" + \"=\" * 60)\n",
    "print(\"Gumini-1B-Base — LayerInteractions\")\n",
    "print(\"=\" * 60)\n",
    "\n",
    "print(\"Compression      : 10 -> 5\")\n",
    "print(\"Retained layers  :\", selection[\"retained_layers\"])\n",
    "print(\"Deleted layers   :\", selection[\"deleted_layers\"])\n",
    "print(f\"Evaluation tokens: {n_tokens:,}\")\n",
    "\n",
    "print(\"-\" * 60)\n",
    "\n",
    "print(f\"Original NLL     : {base_nll:.4f}\")\n",
    "print(f\"Compressed NLL   : {comp_nll:.4f}\")\n",
    "print(f\"NLL increase     : {comp_nll - base_nll:+.4f}\")\n",
    "\n",
    "print()\n",
    "\n",
    "print(f\"Original PPL     : {base_ppl:.2f}\")\n",
    "print(f\"Compressed PPL   : {comp_ppl:.2f}\")\n",
    "print(f\"PPL ratio        : {comp_ppl / base_ppl:.3f}x\")\n",
    "\n",
    "print(\"=\" * 60)"
   ],
   "metadata": {
    "id": "06RcwysfH-Lj"
   },
   "execution_count": null,
   "outputs": []
  }
 ]
}