assignment4/Assignment_4_complete.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GKTWhKtNo8OP"
   },
   "source": [
    "## Part A: Environment Setup and Dataset Loading\n",
    "\n",
    "In accordance with Assignments 2 and 3, we begin by preparing the computational environment using the same libraries introduced during the course. The dataset is loaded directly from Hugging Face using the `load_dataset()` interface, ensuring reproducibility and alignment with the class methodology. This approach maintains consistency across assignments and facilitates later integration with Hugging Face model repositories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0R2G4gY9pBnl",
    "outputId": "78a760f2-26ee-4e83-c106-5e50e9355dc8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[?25l   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/520.7 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m520.7/520.7 kB\u001b[0m \u001b[31m18.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m10.4/10.4 MB\u001b[0m \u001b[31m63.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m60.7/60.7 MB\u001b[0m \u001b[31m18.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m528.8/528.8 kB\u001b[0m \u001b[31m38.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m84.1/84.1 kB\u001b[0m \u001b[31m6.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.9/8.9 MB\u001b[0m \u001b[31m60.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m47.6/47.6 MB\u001b[0m \u001b[31m12.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h"
     ]
    }
   ],
   "source": [
    "# Install required libraries (as used in class)\n",
    "!pip -q install -U datasets transformers accelerate peft bitsandbytes evaluate scikit-learn\n",
    "\n",
    "# Imports\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from collections import Counter\n",
    "from datasets import DatasetDict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZFzLNxR_pEFl"
   },
   "source": [
    "### Loading the Patent Dataset\n",
    "\n",
    "As required, we load the patent claims dataset using the Hugging Face `load_dataset()` interface. This ensures that the dataset is accessed in a standardized and reproducible manner consistent with prior assignments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 219
    },
    "id": "vwqpezZSpGbR",
    "outputId": "0898ad99-a4a9-4d36-fea3-b8f639b58a79"
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'load_dataset' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_276/819393921.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mDATASET_ID\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"AI-Growth-Lab/patents_claims_1.5m_traim_test\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m \u001b[0mdataset\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mload_dataset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mDATASET_ID\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'load_dataset' is not defined"
     ]
    }
   ],
   "source": [
    "DATASET_ID = \"AI-Growth-Lab/patents_claims_1.5m_traim_test\"\n",
    "\n",
    "dataset = load_dataset(DATASET_ID)\n",
    "\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rO2Ve77WuwBB"
   },
   "source": [
    "# Part B: Benchmark Dataset Loading and 35–10–5 Split\n",
    "\n",
    "In this step, we load the fixed 50,000-sample benchmark dataset used in prior iterations to ensure a controlled and comparable experimental setup. Maintaining the same benchmark dataset across Assignments 2, 3, and the final project is necessary for fair model comparison and reliable reporting of performance differences.\n",
    "\n",
    "We then apply a deterministic 35–10–5 split (train/test/eval) to support training, benchmarking, and evaluation under the same protocol required by the assignment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HoroqWLjuw0m"
   },
   "source": [
    "## B.0 Loading the Balanced 50k Dataset from Hugging Face\n",
    "\n",
    "As required, we load the pre-prepared balanced 50,000-sample dataset (`patents_50k_green.parquet`) from a public Hugging Face dataset repository. This ensures the dataset can be accessed reproducibly by external evaluators and enables seamless integration with subsequent Hugging Face model uploads."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MkckB5Byuywe"
   },
   "outputs": [],
   "source": [
    "\n",
    "# HF dataset repo provided by the project\n",
    "HF_DATASET_REPO = \"Asalun/ADLML_course\"\n",
    "\n",
    "# Load the parquet file directly from the repo (reproducible, shareable)\n",
    "ds = load_dataset(\n",
    "    \"parquet\",\n",
    "    data_files=\"https://huggingface.co/datasets/Asalun/ADLML_course/resolve/main/patents_50k_green.parquet\"\n",
    ")\n",
    "\n",
    "data = ds[\"train\"]\n",
    "print(\"Rows:\", len(data))\n",
    "print(\"Columns (first 20):\", data.column_names[:20])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Nke3hBynu9pq"
   },
   "source": [
    "## B.1 Schema Validation: Text and Silver Label\n",
    "\n",
    "Before splitting, we validate the dataset schema by explicitly identifying the claim text field and the binary silver label used for training. This prevents silent errors caused by schema mismatches and ensures that subsequent operations (balancing, uncertainty sampling, and training) are performed on the intended target variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ptHU84UYu-Td"
   },
   "outputs": [],
   "source": [
    "# Prefer class-consistent column names\n",
    "TEXT_CANDIDATES  = [\"text\", \"claim\", \"claim_text\", \"claims\"]\n",
    "LABEL_CANDIDATES = [\"is_green_silver\", \"label\", \"labels\", \"is_green\"]\n",
    "\n",
    "def pick_first(cols, candidates):\n",
    "    for c in candidates:\n",
    "        if c in cols:\n",
    "            return c\n",
    "    return None\n",
    "\n",
    "cols = data.column_names\n",
    "TEXT_COL  = pick_first(cols, TEXT_CANDIDATES)\n",
    "LABEL_COL = pick_first(cols, LABEL_CANDIDATES)\n",
    "\n",
    "print(\"Detected TEXT_COL:\", TEXT_COL)\n",
    "print(\"Detected LABEL_COL:\", LABEL_COL)\n",
    "\n",
    "# If the dataset is Y02-indicator style, derive a binary green label from Y02* columns.\n",
    "if LABEL_COL is None:\n",
    "    Y02_COLS = [c for c in cols if c.startswith(\"Y02\")]\n",
    "    assert len(Y02_COLS) > 0, \"No label column found and no Y02* columns available to derive labels.\"\n",
    "\n",
    "    def add_is_green_silver(ex):\n",
    "        ex[\"is_green_silver\"] = int(any(int(ex[c]) == 1 for c in Y02_COLS))\n",
    "        return ex\n",
    "\n",
    "    data = data.map(add_is_green_silver)\n",
    "    LABEL_COL = \"is_green_silver\"\n",
    "    print(\"Derived LABEL_COL:\", LABEL_COL)\n",
    "\n",
    "assert TEXT_COL is not None, \"No claim text column detected. Inspect columns and set TEXT_COL manually.\"\n",
    "assert LABEL_COL is not None, \"No label column detected/derived. Inspect schema.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6FIOszXxvBco"
   },
   "source": [
    "## B.2 Balanced Dataset Validation\n",
    "\n",
    "The benchmark dataset is expected to contain an equal number of green and non-green claims. We validate class counts to confirm balance. If the dataset is not perfectly balanced, we re-balance it to preserve comparability and prevent class imbalance from confounding model performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gm5wNZ3cvF6O"
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "labels = data[LABEL_COL]\n",
    "count = Counter([int(x) for x in labels])\n",
    "print(\"Class counts:\", dict(count))\n",
    "\n",
    "# If already balanced and exactly 50k, keep as-is\n",
    "if len(data) == 50000 and count.get(0, 0) == count.get(1, 0):\n",
    "    balanced_50k = data\n",
    "else:\n",
    "    # Rebalance to 25k/25k (class-consistent, deterministic)\n",
    "    SEED = 42\n",
    "    pos = data.filter(lambda x: int(x[LABEL_COL]) == 1).shuffle(seed=SEED).select(range(25000))\n",
    "    neg = data.filter(lambda x: int(x[LABEL_COL]) == 0).shuffle(seed=SEED).select(range(25000))\n",
    "    balanced_50k = (pos.add_items(neg.to_dict())).shuffle(seed=SEED)\n",
    "\n",
    "print(\"Final balanced rows:\", len(balanced_50k))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UCZjOrylvfEH"
   },
   "source": [
    "## B.3 35–10–5 Split (Train/Test/Eval)\n",
    "\n",
    "Following the assignment requirement, we split the balanced benchmark dataset into 35,000 training examples, 10,000 test examples, and 5,000 evaluation examples. A fixed random seed is used to ensure reproducibility and to maintain identical splits across runs and model versions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kL7XJb9nvfvq"
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "SEED = 42\n",
    "balanced_50k = balanced_50k.shuffle(seed=SEED)\n",
    "\n",
    "train_silver = balanced_50k.select(range(35000))\n",
    "test_silver  = balanced_50k.select(range(35000, 45000))\n",
    "eval_silver  = balanced_50k.select(range(45000, 50000))\n",
    "\n",
    "splits = DatasetDict({\n",
    "    \"train_silver\": train_silver,\n",
    "    \"test_silver\": test_silver,\n",
    "    \"eval_silver\": eval_silver\n",
    "})\n",
    "\n",
    "print({k: len(v) for k, v in splits.items()})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "g6Z3aBonyKdq"
   },
   "source": [
    "## Part B.3.5 Training the PatentSBERTa Baseline Classifier\n",
    "\n",
    "Before performing uncertainty sampling, we train a baseline PatentSBERTa classifier on the silver-labeled training set. This classifier provides probability estimates \\(p(\\text{green})\\) required to compute uncertainty scores for claim selection.\n",
    "\n",
    "The model is fine-tuned on the 35,000-sample training split and validated on the 5,000-sample evaluation set. Training follows the standard Hugging Face `Trainer` workflow introduced during the course."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "daREgS9syLmf"
   },
   "outputs": [],
   "source": [
    "!pip install transformers datasets evaluate accelerate sentencepiece -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FhNWhFwEyN3C"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "\n",
    "MODEL_ID = \"AI-Growth-Lab/PatentSBERTa\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\n",
    "\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    MODEL_ID,\n",
    "    num_labels=2\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jo4V2o7nyQVk"
   },
   "outputs": [],
   "source": [
    "MAX_LENGTH = 256\n",
    "\n",
    "def tokenize(example):\n",
    "    return tokenizer(\n",
    "        example[TEXT_COL],\n",
    "        truncation=True,\n",
    "        padding=\"max_length\",\n",
    "        max_length=MAX_LENGTH\n",
    "    )\n",
    "\n",
    "train_tok = train_silver.map(tokenize, batched=True)\n",
    "eval_tok  = eval_silver.map(tokenize, batched=True)\n",
    "test_tok  = test_silver.map(tokenize, batched=True)\n",
    "\n",
    "train_tok = train_tok.rename_column(LABEL_COL, \"labels\")\n",
    "eval_tok  = eval_tok.rename_column(LABEL_COL, \"labels\")\n",
    "test_tok  = test_tok.rename_column(LABEL_COL, \"labels\")\n",
    "\n",
    "train_tok.set_format(\"torch\")\n",
    "eval_tok.set_format(\"torch\")\n",
    "test_tok.set_format(\"torch\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "weqb25ntySRu"
   },
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./patentsberta_classifier\",\n",
    "\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=8,\n",
    "    per_device_eval_batch_size=8,\n",
    "\n",
    "    num_train_epochs=1,   # fast training\n",
    "\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "\n",
    "    logging_steps=100,\n",
    "\n",
    "    load_best_model_at_end=True,\n",
    "\n",
    "    fp16=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7hFd7gcKyUjr"
   },
   "outputs": [],
   "source": [
    "import evaluate\n",
    "import numpy as np\n",
    "\n",
    "f1_metric = evaluate.load(\"f1\")\n",
    "accuracy_metric = evaluate.load(\"accuracy\")\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    logits, labels = eval_pred\n",
    "    preds = np.argmax(logits, axis=1)\n",
    "\n",
    "    return {\n",
    "        \"accuracy\": accuracy_metric.compute(predictions=preds, references=labels)[\"accuracy\"],\n",
    "        \"f1\": f1_metric.compute(predictions=preds, references=labels)[\"f1\"]\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Rhc8ZEZS0JdU"
   },
   "outputs": [],
   "source": [
    "from transformers import Trainer\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_tok,\n",
    "    eval_dataset=eval_tok,\n",
    "    compute_metrics=compute_metrics\n",
    ")\n",
    "trainer.train()\n",
    "trainer.evaluate(test_tok)\n",
    "trainer.save_model(\"patentsberta_green_classifier\")\n",
    "tokenizer.save_pretrained(\"patentsberta_green_classifier\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CMRt6xOM7Y5E"
   },
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir=\"patentsberta-green-classifier\",\n",
    "\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=16,\n",
    "    per_device_eval_batch_size=16,\n",
    "\n",
    "    num_train_epochs=1,\n",
    "\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "\n",
    "    load_best_model_at_end=True,\n",
    "\n",
    "    fp16=True,\n",
    "\n",
    "    push_to_hub=True,\n",
    "    hub_model_id=\"your_username/patentsberta-green-classifier\" # Change 'Asalun' to your actual Hugging Face username\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sgulCqDr8kvk"
   },
   "outputs": [],
   "source": [
    "trainer.save_model(\"patentsberta-green-classifier\")\n",
    "tokenizer.save_pretrained(\"patentsberta-green-classifier\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MywsmpuN_8Er"
   },
   "source": [
    "## Part B.4: Uncertainty Sampling for High-Risk Claim Selection\n",
    "\n",
    "In this step, we compute baseline probabilities using a fine-tuned PatentSBERTa sequence classifier and convert them into an uncertainty score \\(u = 1 - 2|p - 0.5|\\). Claims with probabilities closest to 0.5 are considered most ambiguous and are therefore prioritized for advanced labeling.\n",
    "\n",
    "We then select and export the top 100 claims with the highest uncertainty. This fixed high-risk subset serves as the standardized input for the QLoRA-powered multi-agent labeling workflow in Part C."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UkzcwZEUAdHR"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "import numpy as np, pandas as pd, torch, os\n",
    "\n",
    "# The model was saved locally to './patentsberta-green-classifier' in the previous step.\n",
    "# We will load it from there directly.\n",
    "LOCAL_MODEL_PATH = \"./patentsberta-green-classifier\"\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "tok = AutoTokenizer.from_pretrained(LOCAL_MODEL_PATH)\n",
    "mdl = AutoModelForSequenceClassification.from_pretrained(LOCAL_MODEL_PATH).to(device).eval()\n",
    "\n",
    "# 2) Predict p(green) on pool, compute uncertainty u, select top-100\n",
    "pool = splits[\"test_silver\"]\n",
    "bs, max_len = 16, 256\n",
    "probs = []\n",
    "\n",
    "for i in range(0, len(pool), bs):\n",
    "    texts = pool[i:i+bs][TEXT_COL]\n",
    "    enc = tok(texts, padding=True, truncation=True, max_length=max_len, return_tensors=\"pt\").to(device)\n",
    "    with torch.no_grad():\n",
    "        logits = mdl(**enc).logits\n",
    "        p1 = torch.softmax(logits, dim=-1)[:, 1] if logits.shape[-1] == 2 else torch.sigmoid(logits.squeeze(-1))\n",
    "    probs.extend(p1.cpu().numpy().tolist())\n",
    "\n",
    "probs = np.array(probs)\n",
    "u = 1.0 - 2.0 * np.abs(probs - 0.5)\n",
    "\n",
    "top_idx = np.argsort(-u)[:100].tolist()\n",
    "high_risk_100 = pool.select(top_idx).add_column(\"baseline_p_green\", probs[top_idx].tolist()).add_column(\"uncertainty_u\", u[top_idx].tolist())\n",
    "\n",
    "pd.DataFrame(high_risk_100).to_csv(\"high_risk_100.csv\", index=False)\n",
    "print(\"Saved high_risk_100.csv | u-range:\", float(u[top_idx].min()), \"to\", float(u[top_idx].max()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IPIP6NiQCPUf"
   },
   "source": [
    "# Part C.1: Domain Adaptation via QLoRA Fine-Tuning (Mistral-7B)\n",
    "\n",
    "Following the assignment specification, we perform domain adaptation by fine-tuning a generative LLM using QLoRA on the `train_silver` split. The objective is to adapt the model to the linguistic style of patent claims and the reasoning patterns relevant to Y02 classifications.\n",
    "\n",
    "We implement QLoRA using the same components demonstrated in class: 4-bit quantization (BitsAndBytes), LoRA adapters (PEFT), and the Hugging Face Trainer training loop. A minimal configuration (1 epoch, max_length=256) is used to ensure computational feasibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Xd1aY8iBCVoA"
   },
   "outputs": [],
   "source": [
    "\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "\n",
    "BASE_LLM_ID = \"mistralai/Mistral-7B-v0.1\"\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.float16,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    ")\n",
    "\n",
    "tok_llm = AutoTokenizer.from_pretrained(BASE_LLM_ID, use_fast=True)\n",
    "if tok_llm.pad_token is None:\n",
    "    tok_llm.pad_token = tok_llm.eos_token\n",
    "\n",
    "llm = AutoModelForCausalLM.from_pretrained(\n",
    "    BASE_LLM_ID,\n",
    "    quantization_config=bnb_config,\n",
    "    device_map=\"auto\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uO3vfwmRCcQP"
   },
   "source": [
    "## C.1.2 Applying LoRA Adapters\n",
    "\n",
    "We apply LoRA adapters to fine-tune only a small subset of parameters while keeping the base model frozen. This significantly reduces training cost and is consistent with the parameter-efficient fine-tuning approach demonstrated in class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DsBLr2g9CdK9"
   },
   "outputs": [],
   "source": [
    "from peft import LoraConfig, get_peft_model\n",
    "\n",
    "lora_cfg = LoraConfig(\n",
    "    r=8,\n",
    "    lora_alpha=16,\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\"]\n",
    ")\n",
    "\n",
    "llm = get_peft_model(llm, lora_cfg)\n",
    "llm.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R9jqKIr5Cf3b"
   },
   "source": [
    "## C.1.3 Prompt Construction for Domain Adaptation\n",
    "\n",
    "To adapt the model to patent claim language and Y02-oriented decisions, we convert each training example into a concise instruction-style prompt. The target label is included as the expected completion, enabling the generative model to learn the classification decision format in a causal language modeling setting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xqMUuW1OCiAE"
   },
   "outputs": [],
   "source": [
    "def to_prompt(ex):\n",
    "    claim = ex[TEXT_COL]\n",
    "    label = int(ex[LABEL_COL])\n",
    "    ex[\"prompt\"] = (\n",
    "        \"You are classifying patent claims for Y02 (green) relevance.\\n\"\n",
    "        f\"Claim: {claim}\\n\"\n",
    "        \"Output a single digit: 1 if Y02/green-related, otherwise 0.\\n\"\n",
    "        f\"Answer: {label}\"\n",
    "    )\n",
    "    return ex\n",
    "\n",
    "train_sft = splits[\"train_silver\"].map(to_prompt)\n",
    "train_sft = train_sft.remove_columns([c for c in train_sft.column_names if c != \"prompt\"])\n",
    "train_sft"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9K9BoWOLEcva"
   },
   "source": [
    "## C.1.4 Tokenization\n",
    "\n",
    "We tokenize the prompt text using the Mistral tokenizer with truncation to a maximum sequence length of 256. This constraint is selected to ensure training feasibility while still capturing the core content of patent claims for domain adaptation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uq0CJUbnEhD4"
   },
   "outputs": [],
   "source": [
    "MAX_LEN = 256\n",
    "\n",
    "def tok_fn(batch):\n",
    "    return tok_llm(\n",
    "        batch[\"prompt\"],\n",
    "        truncation=True,\n",
    "        padding=\"max_length\",\n",
    "        max_length=MAX_LEN\n",
    "    )\n",
    "\n",
    "train_tok = train_sft.map(tok_fn, batched=True, remove_columns=[\"prompt\"])\n",
    "train_tok.set_format(type=\"torch\")\n",
    "train_tok"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mIld8PLrEoZ-"
   },
   "source": [
    "## C.1.6 Saving the Domain-Adapted QLoRA Model\n",
    "\n",
    "After training, we save the LoRA adapter weights. These adapters represent the domain-adapted model and will be used as the reasoning “brain” for at least one agent in the multi-agent labeling workflow in Part C.2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "O5zDpJgmErTQ"
   },
   "outputs": [],
   "source": [
    "llm.save_pretrained(\"mistral-qlora-patent-adapter\")\n",
    "tok_llm.save_pretrained(\"mistral-qlora-patent-adapter\")\n",
    "print(\"Saved adapter to: mistral-qlora-patent-adapter/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nYldP2jMFMX4"
   },
   "source": [
    "## Part C.2: Multi-Agent Debate\n",
    "\n",
    "In this step we implement a simplified Multi-Agent System using the Mistral language model.\n",
    "Three agents are simulated through role-specific prompts: an Advocate arguing that the claim is green, a Skeptic arguing against it, and a Judge that evaluates both arguments.\n",
    "\n",
    "This debate-style reasoning allows the system to analyze ambiguous patent claims more carefully before producing a final classification decision."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hIFrB4GZFy4o"
   },
   "source": [
    "### C.2.1 Loading the LLM for Agent Reasoning\n",
    "\n",
    "We load the Mistral-7B model in 4-bit mode to reduce memory requirements.\n",
    "This model will serve as the shared reasoning engine for all agents in the debate workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dZFWpa8UI_6d"
   },
   "outputs": [],
   "source": [
    "import torch, gc\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "print(\"GPU available:\", torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kkE-wWP1OMSo"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.float16,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    ")\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    \"mistralai/Mistral-7B-v0.1\",\n",
    "    quantization_config=bnb_config,\n",
    "    device_map=\"auto\",\n",
    "    max_memory={0: \"13GB\", \"cpu\": \"30GB\"},\n",
    "    torch_dtype=torch.float16,\n",
    "    low_cpu_mem_usage=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8e65076",
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_MODEL = \"mistralai/Mistral-7B-v0.1\"\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.float16,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    ")\n",
    "\n",
    "tok = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)\n",
    "if tok.pad_token is None:\n",
    "    tok.pad_token = tok.eos_token\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    BASE_MODEL,\n",
    "    quantization_config=bnb_config,\n",
    "    device_map=\"auto\",\n",
    "    max_memory={0: \"13GB\", \"cpu\": \"30GB\"},\n",
    "    torch_dtype=torch.float16,\n",
    "    low_cpu_mem_usage=True,\n",
    ")\n",
    "model.eval()\n",
    "print(\"Loaded:\", BASE_MODEL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZG_8cDJUPs3k"
   },
   "source": [
    "## Part C.2: Multi-Agent Debate and Labeling of High-Risk Claims\n",
    "\n",
    "We implement a three-agent debate workflow to label the 100 high-risk claims selected in Part B.4. The Advocate argues that a claim is green (Y02-relevant), the Skeptic argues that it is not, and the Judge produces the final decision in a structured JSON format including a confidence score and rationale.\n",
    "\n",
    "This satisfies the required agent roles while keeping the implementation lightweight and reproducible."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6hNLXBO6P9IR"
   },
   "source": [
    "### C.2.1 Generation Utility\n",
    "\n",
    "A small deterministic generation function is used to ensure stable outputs across runs. We keep the token budget small to ensure the multi-agent loop completes within practical runtime limits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BAZr8Tu2P41o"
   },
   "outputs": [],
   "source": [
    "import json, re\n",
    "import pandas as pd\n",
    "import torch\n",
    "\n",
    "# If tokenizer variable isn't defined in your runtime, define it now:\n",
    "# tok = AutoTokenizer.from_pretrained(\"mistralai/Mistral-7B-v0.1\")\n",
    "\n",
    "def generate_text(prompt, max_new_tokens=160):\n",
    "    inputs = tok(prompt, return_tensors=\"pt\").to(model.device)\n",
    "    with torch.no_grad():\n",
    "        out = model.generate(\n",
    "            **inputs,\n",
    "            max_new_tokens=max_new_tokens,\n",
    "            do_sample=False,\n",
    "            pad_token_id=tok.eos_token_id\n",
    "        )\n",
    "    return tok.decode(out[0], skip_special_tokens=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BX6yd0GzQBLu"
   },
   "source": [
    "### C.2.2 Agent Role Definitions\n",
    "\n",
    "Each agent is implemented via a role-specific prompt. The Judge is required to return a JSON object containing the final label, a confidence score in [0,1], and a brief rationale, enabling downstream HITL selection in Part D."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ugGNyglDQDHN"
   },
   "outputs": [],
   "source": [
    "def advocate_prompt(claim):\n",
    "    return f\"\"\"You are Agent 1: The Advocate.\n",
    "Argue that the claim is green (Y02-related). Be concise and cite specific cues in the claim.\n",
    "\n",
    "Claim:\n",
    "{claim}\n",
    "\n",
    "Return 2-4 bullet points.\n",
    "\"\"\"\n",
    "\n",
    "def skeptic_prompt(claim):\n",
    "    return f\"\"\"You are Agent 2: The Skeptic.\n",
    "Argue that the claim is NOT green or may be greenwashing. Be concise and cite specific cues in the claim.\n",
    "\n",
    "Claim:\n",
    "{claim}\n",
    "\n",
    "Return 2-4 bullet points.\n",
    "\"\"\"\n",
    "\n",
    "def judge_prompt(claim, adv, skp):\n",
    "    return f\"\"\"You are Agent 3: The Judge.\n",
    "You must decide if the claim is green (Y02) based on both arguments.\n",
    "\n",
    "Claim:\n",
    "{claim}\n",
    "\n",
    "Advocate:\n",
    "{adv}\n",
    "\n",
    "Skeptic:\n",
    "{skp}\n",
    "\n",
    "Return ONLY valid JSON with this schema:\n",
    "{{\"is_green\": 0 or 1, \"confidence\": 0.0-1.0, \"rationale\": \"one short paragraph\"}}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6zu1NMgYQF8U"
   },
   "source": [
    "### C.2.3 Multi-Agent Labeling Run\n",
    "\n",
    "We run the three-agent workflow over the 100 high-risk claims. For each claim, we store the Advocate and Skeptic arguments and the Judge’s final JSON decision. The outputs are saved to a CSV file to support Part D (HITL) and final dataset creation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mV20s2AURhS9"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "high_risk_100 = pd.read_csv(\"high_risk_100.csv\")   # created earlier in B4\n",
    "TEXT_COL = \"text\"                                  # your dataset uses 'text'\n",
    "claims_high_risk = high_risk_100[TEXT_COL].tolist()\n",
    "print(\"High-risk claims loaded:\", len(claims_high_risk))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6a0ca90",
   "metadata": {},
   "source": [
    "# Part D: Targeted Human Review (HITL) and Gold Label Creation\n",
    "\n",
    "In accordance with the assignment specification, we apply exception-based Human-in-the-Loop (HITL) review only for cases where the automated decision is unreliable. We flag claims for review when the Judge output cannot be parsed as valid JSON or when the Judge reports low confidence.\n",
    "\n",
    "For all non-flagged claims, we accept the Judge decision directly. The result is a gold-labeled dataset of 100 claims (`is_green_gold`) for subsequent integration and final model training."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1346616a",
   "metadata": {},
   "source": [
    "## D.1 Parse Judge Decisions and Flag Cases for Review\n",
    "\n",
    "We parse the Judge JSON output into structured fields (`judge_is_green`, `judge_confidence`, `judge_rationale`). Claims are flagged for human review if parsing fails or if confidence falls below a chosen threshold. This implements exception-based HITL as required, minimizing manual effort while preserving label quality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cf0fef2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json, re\n",
    "import pandas as pd\n",
    "\n",
    "CONF_THRESHOLD = 0.6  # class-consistent, can be adjusted in report with justification\n",
    "\n",
    "agent_df = pd.read_csv(\"agent_labels.csv\")\n",
    "\n",
    "def extract_json(text):\n",
    "    if not isinstance(text, str):\n",
    "        return None\n",
    "    m = re.search(r\"\\{.*\\}\", text, flags=re.DOTALL)\n",
    "    if not m:\n",
    "        return None\n",
    "    try:\n",
    "        return json.loads(m.group(0))\n",
    "    except Exception:\n",
    "        return None\n",
    "\n",
    "parsed = agent_df[\"judge_raw\"].apply(extract_json)\n",
    "\n",
    "agent_df[\"judge_is_green\"] = parsed.apply(lambda x: x.get(\"is_green\") if isinstance(x, dict) else None)\n",
    "agent_df[\"judge_confidence\"] = parsed.apply(lambda x: x.get(\"confidence\") if isinstance(x, dict) else None)\n",
    "agent_df[\"judge_rationale\"] = parsed.apply(lambda x: x.get(\"rationale\") if isinstance(x, dict) else None)\n",
    "\n",
    "agent_df[\"needs_human_review\"] = agent_df[\"judge_is_green\"].isna() | agent_df[\"judge_confidence\"].isna() | (agent_df[\"judge_confidence\"] < CONF_THRESHOLD)\n",
    "\n",
    "print(\"Total claims:\", len(agent_df))\n",
    "print(\"Flagged for HITL:\", int(agent_df[\"needs_human_review\"].sum()))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bdc710e",
   "metadata": {},
   "source": [
    "## D.2 Create HITL Review Template and Final Gold Labels\n",
    "\n",
    "We export only the flagged cases to a compact review file (`hitl_review.csv`). The human reviewer fills `human_is_green` (0/1) for these cases. We then create the final gold label `is_green_gold` by using the human label when available, and otherwise accepting the Judge decision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07e39799",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export flagged cases for manual review\n",
    "review_cols = [\"id\", \"claim\", \"judge_is_green\", \"judge_confidence\", \"judge_rationale\", \"needs_human_review\"]\n",
    "if \"id\" not in agent_df.columns:\n",
    "    agent_df[\"id\"] = range(1, len(agent_df)+1)\n",
    "\n",
    "hitl_df = agent_df.loc[agent_df[\"needs_human_review\"], review_cols].copy()\n",
    "hitl_df[\"human_is_green\"] = \"\"  # to be filled with 0/1\n",
    "hitl_df.to_csv(\"hitl_review.csv\", index=False)\n",
    "print(\"Saved HITL template:\", \"hitl_review.csv\")\n",
    "\n",
    "# Create gold labels (human overrides judge when provided)\n",
    "# After filling hitl_review.csv, re-run the block below to generate gold_100.csv.\n",
    "\n",
    "try:\n",
    "    filled = pd.read_csv(\"hitl_review.csv\")\n",
    "    filled[\"human_is_green\"] = pd.to_numeric(filled[\"human_is_green\"], errors=\"coerce\")\n",
    "    human_map = dict(zip(filled[\"id\"], filled[\"human_is_green\"]))\n",
    "except Exception:\n",
    "    human_map = {}\n",
    "\n",
    "def final_label(row):\n",
    "    hid = row[\"id\"]\n",
    "    if hid in human_map and pd.notna(human_map[hid]):\n",
    "        return int(human_map[hid])\n",
    "    if pd.isna(row[\"judge_is_green\"]):\n",
    "        return None\n",
    "    return int(row[\"judge_is_green\"])\n",
    "\n",
    "agent_df[\"is_green_gold\"] = agent_df.apply(final_label, axis=1)\n",
    "\n",
    "gold_out = agent_df[[\"id\", \"claim\", \"is_green_gold\", \"judge_is_green\", \"judge_confidence\"]].copy()\n",
    "gold_out.to_csv(\"gold_100.csv\", index=False)\n",
    "print(\"Saved gold labels file:\", \"gold_100.csv\")\n",
    "print(\"Gold missing labels:\", int(gold_out['is_green_gold'].isna().sum()))\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}