new.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ed3c9aa0-655a-49b8-a3e1-104c7267d05f",
   "metadata": {},
   "source": [
    "# Nayana Barai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "126eb6d5-5bb4-4493-8788-74b29a8a5c7a",
   "metadata": {},
   "source": [
    "# Part 1: Baseline Model Fine-Tuning (FP32/FP16)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d56957e4-e901-4283-b0fa-4177437f970e",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## 1.1. Install and Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "aa750fb7-b8cc-4364-adde-dc886e6dcc2e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already up-to-date: transformers in /home/saisab/.local/lib/python3.8/site-packages (4.46.3)\n",
      "Requirement already up-to-date: datasets in /home/saisab/.local/lib/python3.8/site-packages (3.1.0)\n",
      "Requirement already up-to-date: scikit-learn in /home/saisab/.local/lib/python3.8/site-packages (1.3.2)\n",
      "Requirement already up-to-date: accelerate in /home/saisab/.local/lib/python3.8/site-packages (1.0.1)\n",
      "Requirement already satisfied, skipping upgrade: tokenizers<0.21,>=0.20 in /home/saisab/.local/lib/python3.8/site-packages (from transformers) (0.20.3)\n",
      "Requirement already satisfied, skipping upgrade: safetensors>=0.4.1 in /home/saisab/.local/lib/python3.8/site-packages (from transformers) (0.5.3)\n",
      "Requirement already satisfied, skipping upgrade: numpy>=1.17 in /usr/local/lib/python3.8/dist-packages (from transformers) (1.24.4)\n",
      "Requirement already satisfied, skipping upgrade: huggingface-hub<1.0,>=0.23.2 in /usr/local/lib/python3.8/dist-packages (from transformers) (0.26.2)\n",
      "Requirement already satisfied, skipping upgrade: filelock in /usr/local/lib/python3.8/dist-packages (from transformers) (3.15.4)\n",
      "Requirement already satisfied, skipping upgrade: pyyaml>=5.1 in /usr/lib/python3/dist-packages (from transformers) (5.3.1)\n",
      "Requirement already satisfied, skipping upgrade: tqdm>=4.27 in /usr/local/lib/python3.8/dist-packages (from transformers) (4.66.2)\n",
      "Requirement already satisfied, skipping upgrade: regex!=2019.12.17 in /usr/local/lib/python3.8/dist-packages (from transformers) (2024.5.15)\n",
      "Requirement already satisfied, skipping upgrade: requests in /home/saisab/.local/lib/python3.8/site-packages (from transformers) (2.32.4)\n",
      "Requirement already satisfied, skipping upgrade: packaging>=20.0 in /home/saisab/.local/lib/python3.8/site-packages (from transformers) (24.2)\n",
      "Requirement already satisfied, skipping upgrade: fsspec[http]<=2024.9.0,>=2023.1.0 in /usr/local/lib/python3.8/dist-packages (from datasets) (2024.6.1)\n",
      "Requirement already satisfied, skipping upgrade: aiohttp in /home/saisab/.local/lib/python3.8/site-packages (from datasets) (3.10.11)\n",
      "Requirement already satisfied, skipping upgrade: xxhash in /usr/local/lib/python3.8/dist-packages (from datasets) (3.4.1)\n",
      "Requirement already satisfied, skipping upgrade: multiprocess<0.70.17 in /home/saisab/.local/lib/python3.8/site-packages (from datasets) (0.70.16)\n",
      "Requirement already satisfied, skipping upgrade: dill<0.3.9,>=0.3.0 in /home/saisab/.local/lib/python3.8/site-packages (from datasets) (0.3.8)\n",
      "Requirement already satisfied, skipping upgrade: pandas in /usr/local/lib/python3.8/dist-packages (from datasets) (2.0.3)\n",
      "Requirement already satisfied, skipping upgrade: pyarrow>=15.0.0 in /home/saisab/.local/lib/python3.8/site-packages (from datasets) (17.0.0)\n",
      "Requirement already satisfied, skipping upgrade: joblib>=1.1.1 in /usr/local/lib/python3.8/dist-packages (from scikit-learn) (1.4.2)\n",
      "Requirement already satisfied, skipping upgrade: scipy>=1.5.0 in /home/saisab/.local/lib/python3.8/site-packages (from scikit-learn) (1.10.1)\n",
      "Requirement already satisfied, skipping upgrade: threadpoolctl>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from scikit-learn) (3.5.0)\n",
      "Requirement already satisfied, skipping upgrade: torch>=1.10.0 in /usr/local/lib/python3.8/dist-packages (from accelerate) (2.4.1+cpu)\n",
      "Requirement already satisfied, skipping upgrade: psutil in /usr/lib/python3/dist-packages (from accelerate) (5.5.1)\n",
      "Requirement already satisfied, skipping upgrade: typing-extensions>=3.7.4.3 in /home/saisab/.local/lib/python3.8/site-packages (from huggingface-hub<1.0,>=0.23.2->transformers) (4.13.2)\n",
      "Requirement already satisfied, skipping upgrade: charset_normalizer<4,>=2 in /home/saisab/.local/lib/python3.8/site-packages (from requests->transformers) (3.4.4)\n",
      "Requirement already satisfied, skipping upgrade: urllib3<3,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->transformers) (2.2.3)\n",
      "Requirement already satisfied, skipping upgrade: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests->transformers) (2.8)\n",
      "Requirement already satisfied, skipping upgrade: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests->transformers) (2019.11.28)\n",
      "Requirement already satisfied, skipping upgrade: async-timeout<6.0,>=4.0; python_version < \"3.11\" in /home/saisab/.local/lib/python3.8/site-packages (from aiohttp->datasets) (4.0.3)\n",
      "Requirement already satisfied, skipping upgrade: aiohappyeyeballs>=2.3.0 in /home/saisab/.local/lib/python3.8/site-packages (from aiohttp->datasets) (2.4.4)\n",
      "Requirement already satisfied, skipping upgrade: aiosignal>=1.1.2 in /home/saisab/.local/lib/python3.8/site-packages (from aiohttp->datasets) (1.3.1)\n",
      "Requirement already satisfied, skipping upgrade: yarl<2.0,>=1.12.0 in /home/saisab/.local/lib/python3.8/site-packages (from aiohttp->datasets) (1.15.2)\n",
      "Requirement already satisfied, skipping upgrade: attrs>=17.3.0 in /home/saisab/.local/lib/python3.8/site-packages (from aiohttp->datasets) (25.3.0)\n",
      "Requirement already satisfied, skipping upgrade: multidict<7.0,>=4.5 in /home/saisab/.local/lib/python3.8/site-packages (from aiohttp->datasets) (6.1.0)\n",
      "Requirement already satisfied, skipping upgrade: frozenlist>=1.1.1 in /home/saisab/.local/lib/python3.8/site-packages (from aiohttp->datasets) (1.5.0)\n",
      "Requirement already satisfied, skipping upgrade: python-dateutil>=2.8.2 in /usr/local/lib/python3.8/dist-packages (from pandas->datasets) (2.9.0.post0)\n",
      "Requirement already satisfied, skipping upgrade: tzdata>=2022.1 in /usr/local/lib/python3.8/dist-packages (from pandas->datasets) (2024.1)\n",
      "Requirement already satisfied, skipping upgrade: pytz>=2020.1 in /usr/local/lib/python3.8/dist-packages (from pandas->datasets) (2024.1)\n",
      "Requirement already satisfied, skipping upgrade: networkx in /usr/local/lib/python3.8/dist-packages (from torch>=1.10.0->accelerate) (3.1)\n",
      "Requirement already satisfied, skipping upgrade: sympy in /usr/local/lib/python3.8/dist-packages (from torch>=1.10.0->accelerate) (1.12.1)\n",
      "Requirement already satisfied, skipping upgrade: jinja2 in /home/saisab/.local/lib/python3.8/site-packages (from torch>=1.10.0->accelerate) (3.1.6)\n",
      "Requirement already satisfied, skipping upgrade: propcache>=0.2.0 in /home/saisab/.local/lib/python3.8/site-packages (from yarl<2.0,>=1.12.0->aiohttp->datasets) (0.2.0)\n",
      "Requirement already satisfied, skipping upgrade: six>=1.5 in /usr/local/lib/python3.8/dist-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.17.0)\n",
      "Requirement already satisfied, skipping upgrade: mpmath<1.4.0,>=1.1.0 in /usr/local/lib/python3.8/dist-packages (from sympy->torch>=1.10.0->accelerate) (1.3.0)\n",
      "Requirement already satisfied, skipping upgrade: MarkupSafe>=2.0 in /home/saisab/.local/lib/python3.8/site-packages (from jinja2->torch>=1.10.0->accelerate) (2.1.5)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/saisab/py10/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "/home/saisab/py10/lib/python3.10/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: '/home/saisab/py10/lib/python3.10/site-packages/torchvision/image.so: undefined symbol: _ZN3c1017RegisterOperatorsD1Ev'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?\n",
      "  warn(\n",
      "2025-11-19 17:40:23.856880: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2025-11-19 17:40:23.901681: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2025-11-19 17:40:23.901728: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2025-11-19 17:40:23.903288: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2025-11-19 17:40:23.911990: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2025-11-19 17:40:25.468369: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------------------------------------------------------------\n",
      "✅ Transformers version: 4.55.4\n",
      "✅ Datasets version: 4.3.0\n",
      "✅ PyTorch version: 2.7.0+cu126\n",
      "------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# 1. INSTALL LIBRARIES\n",
    "!pip install transformers datasets scikit-learn accelerate -U\n",
    "\n",
    "# 2. IMPORT EVERYTHING\n",
    "import transformers\n",
    "import datasets\n",
    "import torch\n",
    "import numpy as np\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "from transformers import TrainingArguments, Trainer\n",
    "from sklearn.metrics import accuracy_score, f1_score\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "# 3. VERIFY THE VERSION (This is the critical check)\n",
    "print(\"---\" * 20)\n",
    "print(f\"✅ Transformers version: {transformers.__version__}\")\n",
    "print(f\"✅ Datasets version: {datasets.__version__}\")\n",
    "print(f\"✅ PyTorch version: {torch.__version__}\")\n",
    "print(\"---\" * 20)\n",
    "\n",
    "# If the transformers version is 4.x.x or higher, the rest of the code will work.\n",
    "# If it is still 3.x.x or lower, the Colab environment is broken."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6f182e1-1331-4248-a9ee-34e9d2cac30f",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## 1.2. Load the Dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f7ab2507-4b6f-4999-ad0c-11f1718f904f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "# Load dataset\n",
    "emotion_dataset = load_dataset(\"dair-ai/emotion\")\n",
    "\n",
    "# Load tokenizer\n",
    "model_checkpoint = \"bert-base-uncased\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[\"text\"], padding=\"max_length\", truncation=True)\n",
    "\n",
    "tokenized_datasets = emotion_dataset.map(tokenize_function, batched=True)\n",
    "\n",
    "# Load model\n",
    "labels = emotion_dataset[\"train\"].features[\"label\"].names\n",
    "num_labels = len(labels)\n",
    "id2label = {i: label for i, label in enumerate(labels)}\n",
    "label2id = {label: i for i, label in enumerate(labels)}\n",
    "\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    model_checkpoint,\n",
    "    num_labels=num_labels,\n",
    "    id2label=id2label,\n",
    "    label2id=label2id\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c396ac7-db63-40eb-9b40-3dd0b11cd255",
   "metadata": {},
   "source": [
    "## 1.3. Entire Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ae2f0908-c375-4f72-9940-1e2f5ea59107",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                    | 0/3000 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting manual training loop...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 33%|████████████▋                         | 1001/3000 [27:00<10:54:13, 19.64s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "--- Epoch 1 / 3 ---\n",
      "Validation Macro F1: 0.9084\n",
      "Validation Accuracy: 0.9325\n",
      "Best model saved at epoch 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 67%|██████████████████████████             | 2001/3000 [54:32<5:38:14, 20.32s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "--- Epoch 2 / 3 ---\n",
      "Validation Macro F1: 0.9168\n",
      "Validation Accuracy: 0.9390\n",
      "Best model saved at epoch 2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████| 3000/3000 [1:19:34<00:00,  1.48s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "--- Epoch 3 / 3 ---\n",
      "Validation Macro F1: 0.9159\n",
      "Validation Accuracy: 0.9390\n",
      "\n",
      "🎉 Training complete. Best model has been loaded.\n",
      "Baseline model saved to ./finetuned-bert-emotion-baseline\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from torch.utils.data import DataLoader\n",
    "from transformers import DataCollatorWithPadding, get_scheduler\n",
    "from sklearn.metrics import accuracy_score, f1_score\n",
    "from tqdm.auto import tqdm\n",
    "import copy\n",
    "\n",
    "# --- 1. Prepare DataLoaders ---\n",
    "# We must re-run this preparation in case the notebook state was lost.\n",
    "tokenized_datasets.set_format(\"torch\")\n",
    "tokenized_datasets = tokenized_datasets.remove_columns([\"text\"])\n",
    "tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n",
    "\n",
    "data_collator = DataCollatorWithPadding(tokenizer=tokenizer)\n",
    "\n",
    "train_dataloader = DataLoader(\n",
    "    tokenized_datasets[\"train\"], shuffle=True, batch_size=16, collate_fn=data_collator\n",
    ")\n",
    "eval_dataloader = DataLoader(\n",
    "    tokenized_datasets[\"validation\"], batch_size=64, collate_fn=data_collator\n",
    ")\n",
    "\n",
    "# --- 2. Set up Optimizer and Device ---\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "model.to(device)\n",
    "\n",
    "# --- 3. Set up Scheduler and Training Parameters ---\n",
    "num_epochs = 3\n",
    "num_training_steps = num_epochs * len(train_dataloader)\n",
    "lr_scheduler = get_scheduler(\n",
    "    name=\"linear\", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps\n",
    ")\n",
    "\n",
    "# --- 4. The Training and Evaluation Loop ---\n",
    "progress_bar = tqdm(range(num_training_steps))\n",
    "best_metric = -1.0\n",
    "best_model_state = None\n",
    "\n",
    "print(\"Starting manual training loop...\")\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    for batch in train_dataloader:\n",
    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
    "        outputs = model(**batch)\n",
    "        loss = outputs.loss\n",
    "        loss.backward()\n",
    "\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "        optimizer.zero_grad()\n",
    "        progress_bar.update(1)\n",
    "\n",
    "    # --- Evaluation after each epoch ---\n",
    "    model.eval()\n",
    "    all_predictions = []\n",
    "    all_labels = []\n",
    "    for batch in eval_dataloader:\n",
    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
    "        with torch.no_grad():\n",
    "            outputs = model(**batch)\n",
    "        \n",
    "        logits = outputs.logits\n",
    "        predictions = torch.argmax(logits, dim=-1).cpu().numpy()\n",
    "        labels = batch[\"labels\"].cpu().numpy()\n",
    "        all_predictions.extend(predictions)\n",
    "        all_labels.extend(labels)\n",
    "\n",
    "    macro_f1 = f1_score(all_labels, all_predictions, average=\"macro\")\n",
    "    accuracy = accuracy_score(all_labels, all_predictions)\n",
    "    \n",
    "    print(f\"\\n--- Epoch {epoch + 1} / {num_epochs} ---\")\n",
    "    print(f\"Validation Macro F1: {macro_f1:.4f}\")\n",
    "    print(f\"Validation Accuracy: {accuracy:.4f}\")\n",
    "\n",
    "    # Save the best model\n",
    "    if macro_f1 > best_metric:\n",
    "        best_metric = macro_f1\n",
    "        best_model_state = copy.deepcopy(model.state_dict())\n",
    "        print(f\"Best model saved at epoch {epoch + 1}\")\n",
    "\n",
    "# --- 5. Load the Best Model ---\n",
    "model.load_state_dict(best_model_state)\n",
    "print(\"\\n🎉 Training complete. Best model has been loaded.\")\n",
    "\n",
    "# --- 6. Save the Final Baseline Model ---\n",
    "baseline_model_dir = \"./finetuned-bert-emotion-baseline\"\n",
    "model.save_pretrained(baseline_model_dir)\n",
    "tokenizer.save_pretrained(baseline_model_dir)\n",
    "print(f\"Baseline model saved to {baseline_model_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a48c9f2-6503-4d3d-8541-14963654e5dc",
   "metadata": {},
   "source": [
    "## 1.8. Evaluate the Baseline Model on the Test Set\n",
    "Now that we have our best model, we must evaluate its performance on the test split of the dataset. This provides the final, unbiased measure of our model's performance. We will create a test_dataloader and run a similar evaluation loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b27ff9b5-1dfd-4b11-ad9c-1da187224648",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training is commented out. Loading the saved baseline model directly.\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import os\n",
    "import pandas as pd\n",
    "from sklearn.metrics import classification_report, confusion_matrix\n",
    "\n",
    "# --- 1. Create the Test DataLoader ---\n",
    "test_dataloader = DataLoader(\n",
    "    tokenized_datasets[\"test\"], batch_size=64, collate_fn=data_collator\n",
    ")\n",
    "\n",
    "# --- 2. Run Evaluation on the Test Set ---\n",
    "model.eval()\n",
    "all_predictions = []\n",
    "all_labels = []\n",
    "inference_times = []\n",
    "\n",
    "print(\"Evaluating on the test set...\")\n",
    "for batch in tqdm(test_dataloader):\n",
    "    batch = {k: v.to(device) for k, v in batch.items()}\n",
    "    start_time = time.time()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        outputs = model(**batch)\n",
    "        \n",
    "    end_time = time.time()\n",
    "    inference_times.append(end_time - start_time)\n",
    "\n",
    "    logits = outputs.logits\n",
    "    predictions = torch.argmax(logits, dim=-1).cpu().numpy()\n",
    "    labels = batch[\"labels\"].cpu().numpy()\n",
    "    all_predictions.extend(predictions)\n",
    "    all_labels.extend(labels)\n",
    "\n",
    "# --- 3. Calculate Final Metrics ---\n",
    "# Performance Metrics\n",
    "test_macro_f1 = f1_score(all_labels, all_predictions, average=\"macro\")\n",
    "test_accuracy = accuracy_score(all_labels, all_predictions)\n",
    "\n",
    "# Model Size (CORRECTED FILENAME)\n",
    "model_size_mb = os.path.getsize(f\"{baseline_model_dir}/model.safetensors\") / (1024 * 1024)\n",
    "\n",
    "# Latency Metrics\n",
    "avg_batch_latency_ms = (sum(inference_times) / len(inference_times)) * 1000\n",
    "avg_example_latency_ms = avg_batch_latency_ms / test_dataloader.batch_size\n",
    "\n",
    "# Store results in a dictionary for later\n",
    "baseline_results = {\n",
    "    \"model\": \"Baseline (FP32)\",\n",
    "    \"accuracy\": test_accuracy,\n",
    "    \"macro_f1\": test_macro_f1,\n",
    "    \"model_size_mb\": model_size_mb,\n",
    "    \"latency_ms_per_batch\": avg_batch_latency_ms\n",
    "}\n",
    "\n",
    "print(\"\\n--- Baseline Model Test Results ---\")\n",
    "print(f\"Accuracy: {test_accuracy:.4f}\")\n",
    "print(f\"Macro F1-Score: {test_macro_f1:.4f}\")\n",
    "print(f\"Model Size: {model_size_mb:.2f} MB\")\n",
    "print(f\"Avg. Latency per Batch: {avg_batch_latency_ms:.2f} ms\")\n",
    "\n",
    "\n",
    "# \"\"\"\n",
    "# import torch\n",
    "# import numpy as np\n",
    "# from torch.utils.data import DataLoader\n",
    "# from transformers import DataCollatorWithPadding, get_scheduler\n",
    "# from sklearn.metrics import accuracy_score, f1_score\n",
    "# from tqdm.auto import tqdm\n",
    "# import copy\n",
    "\n",
    "# # --- 1. Prepare DataLoaders ---\n",
    "# # We must re-run this preparation in case the notebook state was lost.\n",
    "# tokenized_datasets.set_format(\"torch\")\n",
    "# tokenized_datasets = tokenized_datasets.remove_columns([\"text\"])\n",
    "# tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n",
    "\n",
    "# data_collator = DataCollatorWithPadding(tokenizer=tokenizer)\n",
    "\n",
    "# train_dataloader = DataLoader(\n",
    "#     tokenized_datasets[\"train\"], shuffle=True, batch_size=16, collate_fn=data_collator\n",
    "# )\n",
    "# eval_dataloader = DataLoader(\n",
    "#     tokenized_datasets[\"validation\"], batch_size=64, collate_fn=data_collator\n",
    "# )\n",
    "\n",
    "# # --- 2. Set up Optimizer and Device ---\n",
    "# optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)\n",
    "# device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "# model.to(device)\n",
    "\n",
    "# # --- 3. Set up Scheduler and Training Parameters ---\n",
    "# num_epochs = 3\n",
    "# num_training_steps = num_epochs * len(train_dataloader)\n",
    "# lr_scheduler = get_scheduler(\n",
    "#     name=\"linear\", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps\n",
    "# )\n",
    "\n",
    "# # --- 4. The Training and Evaluation Loop ---\n",
    "# progress_bar = tqdm(range(num_training_steps))\n",
    "# best_metric = -1.0\n",
    "# best_model_state = None\n",
    "\n",
    "# print(\"Starting manual training loop...\")\n",
    "# for epoch in range(num_epochs):\n",
    "#     model.train()\n",
    "#     for batch in train_dataloader:\n",
    "#         batch = {k: v.to(device) for k, v in batch.items()}\n",
    "#         outputs = model(**batch)\n",
    "#         loss = outputs.loss\n",
    "#         loss.backward()\n",
    "\n",
    "#         optimizer.step()\n",
    "#         lr_scheduler.step()\n",
    "#         optimizer.zero_grad()\n",
    "#         progress_bar.update(1)\n",
    "\n",
    "#     # --- Evaluation after each epoch ---\n",
    "#     model.eval()\n",
    "#     all_predictions = []\n",
    "#     all_labels = []\n",
    "#     for batch in eval_dataloader:\n",
    "#         batch = {k: v.to(device) for k, v in batch.items()}\n",
    "#         with torch.no_grad():\n",
    "#             outputs = model(**batch)\n",
    "        \n",
    "#         logits = outputs.logits\n",
    "#         predictions = torch.argmax(logits, dim=-1).cpu().numpy()\n",
    "#         labels = batch[\"labels\"].cpu().numpy()\n",
    "#         all_predictions.extend(predictions)\n",
    "#         all_labels.extend(labels)\n",
    "\n",
    "#     macro_f1 = f1_score(all_labels, all_predictions, average=\"macro\")\n",
    "#     accuracy = accuracy_score(all_labels, all_predictions)\n",
    "    \n",
    "#     print(f\"\\n--- Epoch {epoch + 1} / {num_epochs} ---\")\n",
    "#     print(f\"Validation Macro F1: {macro_f1:.4f}\")\n",
    "#     print(f\"Validation Accuracy: {accuracy:.4f}\")\n",
    "\n",
    "#     # Save the best model\n",
    "#     if macro_f1 > best_metric:\n",
    "#         best_metric = macro_f1\n",
    "#         best_model_state = copy.deepcopy(model.state_dict())\n",
    "#         print(f\"Best model saved at epoch {epoch + 1}\")\n",
    "\n",
    "# # --- 5. Load the Best Model ---\n",
    "# model.load_state_dict(best_model_state)\n",
    "# print(\"\\n🎉 Training complete. Best model has been loaded.\")\n",
    "\n",
    "# # --- 6. Save the Final Baseline Model ---\n",
    "# baseline_model_dir = \"./finetuned-bert-emotion-baseline\"\n",
    "# model.save_pretrained(baseline_model_dir)\n",
    "# tokenizer.save_pretrained(baseline_model_dir)\n",
    "# print(f\"Baseline model saved to {baseline_model_dir}\")\n",
    "# \"\"\"\n",
    "\n",
    "# # --- TRAINING SKIPPED: Load the already fine-tuned model from disk ---\n",
    "# print(\"Training is commented out. Loading the saved baseline model directly.\")\n",
    "# baseline_model_dir = \"./finetuned-bert-emotion-baseline\"\n",
    "# # The 'model' variable is what the next cells expect. We load our saved model into it.\n",
    "# model = AutoModelForSequenceClassification.from_pretrained(baseline_model_dir)\n",
    "# # Also need to define the device for the evaluation cell\n",
    "# device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "# model.to(device)\n",
    "\n",
    "# # We still need the DataCollator for the evaluation cell\n",
    "# from transformers import DataCollatorWithPadding\n",
    "# data_collator = DataCollatorWithPadding(tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f01d7baf-7080-488e-bf62-fc4df9c6e738",
   "metadata": {},
   "source": [
    "## 1.9. Analysis: Confusion Matrix and Classification Report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2f17fc36-936c-42b6-a2ab-a3310e217a5e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Classification Report ---\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "     sadness       0.95      0.97      0.96       550\n",
      "         joy       0.96      0.95      0.95       704\n",
      "        love       0.86      0.88      0.87       178\n",
      "       anger       0.98      0.92      0.95       275\n",
      "        fear       0.89      0.92      0.90       212\n",
      "    surprise       0.83      0.89      0.86        81\n",
      "\n",
      "    accuracy                           0.94      2000\n",
      "   macro avg       0.91      0.92      0.92      2000\n",
      "weighted avg       0.94      0.94      0.94      2000\n",
      "\n",
      "\n",
      "--- Confusion Matrix ---\n",
      "          sadness  joy  love  anger  fear  surprise\n",
      "sadness       536    1     1      4     8         0\n",
      "joy             4  668    23      0     2         7\n",
      "love            0   22   156      0     0         0\n",
      "anger          14    2     1    252     6         0\n",
      "fear            7    1     0      2   194         8\n",
      "surprise        1    1     0      0     7        72\n"
     ]
    }
   ],
   "source": [
    "from sklearn.metrics import classification_report, confusion_matrix\n",
    "import pandas as pd\n",
    "\n",
    "# Get the label names from our label mapping\n",
    "label_names = list(label2id.keys())\n",
    "\n",
    "# Generate and print the classification report\n",
    "print(\"--- Classification Report ---\")\n",
    "report = classification_report(all_labels, all_predictions, target_names=label_names)\n",
    "print(report)\n",
    "\n",
    "# Generate and print the confusion matrix\n",
    "print(\"\\n--- Confusion Matrix ---\")\n",
    "conf_matrix = confusion_matrix(all_labels, all_predictions)\n",
    "\n",
    "# Use pandas for a more readable confusion matrix\n",
    "conf_matrix_df = pd.DataFrame(conf_matrix, index=label_names, columns=label_names)\n",
    "print(conf_matrix_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "144c6d74-ea3c-44da-8ab7-82a5430c73b8",
   "metadata": {},
   "source": [
    "# Part 2: Post-Training Quantization (PTQ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a22fbad-64be-41b1-a455-d570573933db",
   "metadata": {},
   "source": [
    "## 2.1. Load the Baseline Model and Apply Dynamic Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b8139f89-73a5-43f9-9f71-4c2112499f48",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Original FP32 Model ---\n",
      "BertForSequenceClassification(\n",
      "  (bert): BertModel(\n",
      "    (embeddings): BertEmbeddings(\n",
      "      (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
      "      (position_embeddings): Embedding(512, 768)\n",
      "      (token_type_embeddings): Embedding(2, 768)\n",
      "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
      "      (dropout): Dropout(p=0.1, inplace=False)\n",
      "    )\n",
      "    (encoder): BertEncoder(\n",
      "      (layer): ModuleList(\n",
      "        (0-11): 12 x BertLayer(\n",
      "          (attention): BertAttention(\n",
      "            (self): BertSdpaSelfAttention(\n",
      "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
      "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
      "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
      "              (dropout): Dropout(p=0.1, inplace=False)\n",
      "            )\n",
      "            (output): BertSelfOutput(\n",
      "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
      "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
      "              (dropout): Dropout(p=0.1, inplace=False)\n",
      "            )\n",
      "          )\n",
      "          (intermediate): BertIntermediate(\n",
      "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
      "            (intermediate_act_fn): GELUActivation()\n",
      "          )\n",
      "          (output): BertOutput(\n",
      "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
      "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (pooler): BertPooler(\n",
      "      (dense): Linear(in_features=768, out_features=768, bias=True)\n",
      "      (activation): Tanh()\n",
      "    )\n",
      "  )\n",
      "  (dropout): Dropout(p=0.1, inplace=False)\n",
      "  (classifier): Linear(in_features=768, out_features=6, bias=True)\n",
      ")\n",
      "\n",
      "--- Quantized INT8 Model ---\n",
      "BertForSequenceClassification(\n",
      "  (bert): BertModel(\n",
      "    (embeddings): BertEmbeddings(\n",
      "      (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
      "      (position_embeddings): Embedding(512, 768)\n",
      "      (token_type_embeddings): Embedding(2, 768)\n",
      "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
      "      (dropout): Dropout(p=0.1, inplace=False)\n",
      "    )\n",
      "    (encoder): BertEncoder(\n",
      "      (layer): ModuleList(\n",
      "        (0-11): 12 x BertLayer(\n",
      "          (attention): BertAttention(\n",
      "            (self): BertSdpaSelfAttention(\n",
      "              (query): DynamicQuantizedLinear(in_features=768, out_features=768, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
      "              (key): DynamicQuantizedLinear(in_features=768, out_features=768, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
      "              (value): DynamicQuantizedLinear(in_features=768, out_features=768, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
      "              (dropout): Dropout(p=0.1, inplace=False)\n",
      "            )\n",
      "            (output): BertSelfOutput(\n",
      "              (dense): DynamicQuantizedLinear(in_features=768, out_features=768, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
      "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
      "              (dropout): Dropout(p=0.1, inplace=False)\n",
      "            )\n",
      "          )\n",
      "          (intermediate): BertIntermediate(\n",
      "            (dense): DynamicQuantizedLinear(in_features=768, out_features=3072, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
      "            (intermediate_act_fn): GELUActivation()\n",
      "          )\n",
      "          (output): BertOutput(\n",
      "            (dense): DynamicQuantizedLinear(in_features=3072, out_features=768, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
      "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (pooler): BertPooler(\n",
      "      (dense): DynamicQuantizedLinear(in_features=768, out_features=768, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
      "      (activation): Tanh()\n",
      "    )\n",
      "  )\n",
      "  (dropout): Dropout(p=0.1, inplace=False)\n",
      "  (classifier): DynamicQuantizedLinear(in_features=768, out_features=6, dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch.quantization import quantize_dynamic\n",
    "from transformers import AutoModelForSequenceClassification\n",
    "\n",
    "# --- 1. Load the fine-tuned FP32 model ---\n",
    "baseline_model_dir = \"./finetuned-bert-emotion-baseline\"\n",
    "model_fp32 = AutoModelForSequenceClassification.from_pretrained(baseline_model_dir)\n",
    "\n",
    "# --- 2. Apply dynamic quantization ---\n",
    "# We specify that we want to quantize the Linear layers of the model.\n",
    "# The model must be in evaluation mode and on the CPU.\n",
    "model_quantized = quantize_dynamic(\n",
    "    model=model_fp32.to(\"cpu\").eval(),\n",
    "    qconfig_spec={torch.nn.Linear},\n",
    "    dtype=torch.qint8\n",
    ")\n",
    "\n",
    "print(\"--- Original FP32 Model ---\")\n",
    "print(model_fp32)\n",
    "print(\"\\n--- Quantized INT8 Model ---\")\n",
    "print(model_quantized)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1030dff5-668f-4178-a360-b0cd067aefb3",
   "metadata": {},
   "source": [
    "## 2.2. Evaluate the Quantized Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b235484a-04e6-4bc9-b509-6929d4f2b12f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluating quantized model on the CPU...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  0%|                                                      | 0/32 [00:00<?, ?it/s]\u001b[A\n",
      "  3%|█▍                                            | 1/32 [00:06<03:23,  6.57s/it]\u001b[A\n",
      "  6%|██▉                                           | 2/32 [00:13<03:18,  6.62s/it]\u001b[A\n",
      "  9%|████▎                                         | 3/32 [00:20<03:16,  6.78s/it]\u001b[A\n",
      " 12%|█████▊                                        | 4/32 [00:27<03:14,  6.95s/it]\u001b[A\n",
      " 16%|███████▏                                      | 5/32 [00:34<03:12,  7.14s/it]\u001b[A\n",
      " 19%|████████▋                                     | 6/32 [00:42<03:06,  7.19s/it]\u001b[A\n",
      " 22%|██████████                                    | 7/32 [00:49<02:58,  7.12s/it]\u001b[A\n",
      " 25%|███████████▌                                  | 8/32 [00:55<02:48,  7.00s/it]\u001b[A\n",
      " 28%|████████████▉                                 | 9/32 [01:03<02:42,  7.08s/it]\u001b[A\n",
      " 31%|██████████████                               | 10/32 [01:10<02:36,  7.12s/it]\u001b[A\n",
      " 34%|███████████████▍                             | 11/32 [01:17<02:28,  7.09s/it]\u001b[A\n",
      " 38%|████████████████▉                            | 12/32 [01:24<02:19,  6.96s/it]\u001b[A\n",
      " 41%|██████████████████▎                          | 13/32 [01:30<02:10,  6.87s/it]\u001b[A\n",
      " 44%|███████████████████▋                         | 14/32 [01:37<02:02,  6.82s/it]\u001b[A\n",
      " 47%|█████████████████████                        | 15/32 [01:44<01:55,  6.79s/it]\u001b[A\n",
      " 50%|██████████████████████▌                      | 16/32 [01:51<01:50,  6.90s/it]\u001b[A\n",
      " 53%|███████████████████████▉                     | 17/32 [01:58<01:45,  7.00s/it]\u001b[A\n",
      " 56%|█████████████████████████▎                   | 18/32 [02:06<01:40,  7.19s/it]\u001b[A\n",
      " 59%|██████████████████████████▋                  | 19/32 [02:13<01:34,  7.27s/it]\u001b[A\n",
      " 62%|████████████████████████████▏                | 20/32 [02:21<01:27,  7.30s/it]\u001b[A\n",
      " 66%|█████████████████████████████▌               | 21/32 [02:28<01:20,  7.36s/it]\u001b[A\n",
      " 69%|██████████████████████████████▉              | 22/32 [02:35<01:13,  7.37s/it]\u001b[A\n",
      " 72%|████████████████████████████████▎            | 23/32 [02:43<01:06,  7.40s/it]\u001b[A\n",
      " 75%|█████████████████████████████████▊           | 24/32 [02:50<00:58,  7.37s/it]\u001b[A\n",
      " 78%|███████████████████████████████████▏         | 25/32 [02:58<00:51,  7.39s/it]\u001b[A\n",
      " 81%|████████████████████████████████████▌        | 26/32 [03:05<00:44,  7.40s/it]\u001b[A\n",
      " 84%|█████████████████████████████████████▉       | 27/32 [03:12<00:36,  7.40s/it]\u001b[A\n",
      " 88%|███████████████████████████████████████▍     | 28/32 [03:20<00:29,  7.40s/it]\u001b[A\n",
      " 91%|████████████████████████████████████████▊    | 29/32 [03:27<00:21,  7.31s/it]\u001b[A\n",
      " 94%|██████████████████████████████████████████▏  | 30/32 [03:34<00:14,  7.30s/it]\u001b[A\n",
      " 97%|███████████████████████████████████████████▌ | 31/32 [03:42<00:07,  7.35s/it]\u001b[A\n",
      "100%|█████████████████████████████████████████████| 32/32 [03:43<00:00,  7.00s/it]\u001b[A"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "--- PTQ Model Test Results ---\n",
      "Accuracy: 0.9265\n",
      "Macro F1-Score: 0.8869\n",
      "Model Size: 173.08 MB\n",
      "Avg. Latency per Batch (CPU): 6970.90 ms\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import time\n",
    "\n",
    "# --- 1. Save the quantized model to disk to measure its size ---\n",
    "quantized_model_dir = \"./quantized-bert-emotion-ptq\"\n",
    "os.makedirs(quantized_model_dir, exist_ok=True)\n",
    "quantized_model_path = f\"{quantized_model_dir}/model.pth\"\n",
    "torch.save(model_quantized.state_dict(), quantized_model_path)\n",
    "\n",
    "# --- 2. Set up the evaluation device and dataloader ---\n",
    "# IMPORTANT: Evaluation must be on the CPU\n",
    "device_cpu = torch.device(\"cpu\")\n",
    "model_quantized.to(device_cpu)\n",
    "model_quantized.eval()\n",
    "\n",
    "# Re-create the test dataloader just to be safe\n",
    "test_dataloader = DataLoader(\n",
    "    tokenized_datasets[\"test\"], batch_size=64, collate_fn=data_collator\n",
    ")\n",
    "\n",
    "# --- 3. Run Evaluation on the Test Set ---\n",
    "all_predictions = []\n",
    "all_labels = []\n",
    "inference_times = []\n",
    "\n",
    "print(\"Evaluating quantized model on the CPU...\")\n",
    "for batch in tqdm(test_dataloader):\n",
    "    # No need to move to device, as it's already CPU data\n",
    "    batch = {k: v.to(device_cpu) for k, v in batch.items()}\n",
    "    start_time = time.time()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        outputs = model_quantized(**batch)\n",
    "        \n",
    "    end_time = time.time()\n",
    "    inference_times.append(end_time - start_time)\n",
    "\n",
    "    logits = outputs.logits\n",
    "    predictions = torch.argmax(logits, dim=-1).numpy()\n",
    "    labels = batch[\"labels\"].numpy()\n",
    "    all_predictions.extend(predictions)\n",
    "    all_labels.extend(labels)\n",
    "\n",
    "# --- 4. Calculate Final Metrics ---\n",
    "# Performance Metrics\n",
    "ptq_macro_f1 = f1_score(all_labels, all_predictions, average=\"macro\")\n",
    "ptq_accuracy = accuracy_score(all_labels, all_predictions)\n",
    "\n",
    "# Model Size\n",
    "ptq_model_size_mb = os.path.getsize(quantized_model_path) / (1024 * 1024)\n",
    "\n",
    "# Latency Metrics\n",
    "ptq_avg_batch_latency_ms = (sum(inference_times) / len(inference_times)) * 1000\n",
    "\n",
    "# Store results for the final comparison table\n",
    "ptq_results = {\n",
    "    \"model\": \"PTQ (INT8)\",\n",
    "    \"accuracy\": ptq_accuracy,\n",
    "    \"macro_f1\": ptq_macro_f1,\n",
    "    \"model_size_mb\": ptq_model_size_mb,\n",
    "    \"latency_ms_per_batch\": ptq_avg_batch_latency_ms\n",
    "}\n",
    "\n",
    "print(\"\\n--- PTQ Model Test Results ---\")\n",
    "print(f\"Accuracy: {ptq_accuracy:.4f}\")\n",
    "print(f\"Macro F1-Score: {ptq_macro_f1:.4f}\")\n",
    "print(f\"Model Size: {ptq_model_size_mb:.2f} MB\")\n",
    "print(f\"Avg. Latency per Batch (CPU): {ptq_avg_batch_latency_ms:.2f} ms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae54a97c-99f3-4a50-bcfd-d3763f6f03f2",
   "metadata": {},
   "source": [
    "## 2.3. Comparison: Baseline vs. PTQ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e40a6927-cbec-45e2-808b-478dc5a2f3d5",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'baseline_results' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[8], line 4\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpandas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpd\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m# Create a DataFrame from our stored results\u001b[39;00m\n\u001b[0;32m----> 4\u001b[0m comparison_df \u001b[38;5;241m=\u001b[39m pd\u001b[38;5;241m.\u001b[39mDataFrame([\u001b[43mbaseline_results\u001b[49m, ptq_results])\n\u001b[1;32m      5\u001b[0m comparison_df\u001b[38;5;241m.\u001b[39mset_index(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mmodel\u001b[39m\u001b[38;5;124m'\u001b[39m, inplace\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m      7\u001b[0m \u001b[38;5;66;03m# Add a note about the latency measurement context\u001b[39;00m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'baseline_results' is not defined"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# Create a DataFrame from our stored results\n",
    "comparison_df = pd.DataFrame([baseline_results, ptq_results])\n",
    "comparison_df.set_index('model', inplace=True)\n",
    "\n",
    "# Add a note about the latency measurement context\n",
    "comparison_df.rename(columns={'latency_ms_per_batch': 'Latency (ms/batch)'}, inplace=True)\n",
    "comparison_df['Latency (ms/batch)'] = comparison_df['Latency (ms/batch)'].round(2).astype(str)\n",
    "comparison_df.loc['Baseline (FP32)', 'Latency (ms/batch)'] += ' (GPU)'\n",
    "comparison_df.loc['PTQ (INT8)', 'Latency (ms/batch)'] += ' (CPU)'\n",
    "\n",
    "\n",
    "print(\"--- Performance and Resource Comparison ---\")\n",
    "print(comparison_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25229ca4-4f12-4422-845a-24a8cddfe65d",
   "metadata": {},
   "source": [
    "# **Part 3: Quantization-Aware Training (QAT)**\n",
    "Quantization-Aware Training (QAT) is a more advanced technique. Instead of quantizing a fully trained model, we introduce the \"simulation\" of quantization *during* the fine-tuning process. This allows the model to learn and adapt to the effects of reduced precision, which can lead to significantly better accuracy compared to PTQ.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a550134b-a9ef-416d-a87d-93ca72629d63",
   "metadata": {},
   "source": [
    "## 3.1. Prepare the Model for QAT\n",
    "We start by loading a fresh copy of our fine-tuned FP32 model. Then, we attach a quantization configuration and use a PyTorch helper function, `prepare_qat`, to insert special modules (`FakeQuantize`) into the model. These modules will simulate the effect of INT8 quantization during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69926273-f2d8-493e-ac00-6652f6eb9680",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.quantization import get_default_qat_qconfig, prepare_qat\n",
    "\n",
    "# --- 1. Load a fresh copy of the fine-tuned FP32 model ---\n",
    "model_for_qat = AutoModelForSequenceClassification.from_pretrained(baseline_model_dir)\n",
    "model_for_qat.train()\n",
    "\n",
    "# --- 2. Get the default QAT configuration ---\n",
    "qconfig = get_default_qat_qconfig('fbgemm')\n",
    "model_for_qat.qconfig = qconfig\n",
    "\n",
    "# --- 3. CRITICAL FIX: Exclude embeddings from quantization ---\n",
    "# By setting the qconfig of the embeddings module to None, we tell PyTorch to skip it.\n",
    "model_for_qat.bert.embeddings.qconfig = None\n",
    "print(\"Disabled quantization for the following module:\")\n",
    "print(model_for_qat.bert.embeddings)\n",
    "\n",
    "\n",
    "# --- 4. Prepare the model for QAT ---\n",
    "model_qat_prepared = prepare_qat(model_for_qat)\n",
    "\n",
    "# You can inspect the model below and see that the Embedding layers\n",
    "# no longer have FakeQuantize modules attached.\n",
    "print(\"\\n--- Model Prepared for QAT (Embeddings Skipped) ---\")\n",
    "print(model_qat_prepared.bert.embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4f1c146-708a-4560-a446-b22d24590974",
   "metadata": {},
   "source": [
    "## 3.2. Fine-Tune the QAT-Prepared Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bab9b605-0fca-4ef9-834c-e1fed6c05e83",
   "metadata": {},
   "outputs": [],
   "source": [
    "import optimum\n",
    "from optimum.intel import INCQuantizer\n",
    "from transformers import AutoModelForSequenceClassification\n",
    "\n",
    "# --- 1. Load the fine-tuned FP32 model ---\n",
    "# This is the same model we started with in Part 2\n",
    "baseline_model_dir = \"./finetuned-bert-emotion-baseline\"\n",
    "model_fp32 = AutoModelForSequenceClassification.from_pretrained(baseline_model_dir)\n",
    "\n",
    "# --- 2. Initialize the Quantizer ---\n",
    "quantizer = INCQuantizer.from_pretrained(model_fp32)\n",
    "\n",
    "# --- 3. Create a calibration function ---\n",
    "# The quantizer needs a small sample of data to observe the model's behavior.\n",
    "# We'll use 100 examples from the training set.\n",
    "calibration_dataset = tokenized_datasets[\"train\"].select(range(100))\n",
    "# Remove columns the model doesn't expect\n",
    "calibration_dataset = calibration_dataset.remove_columns([\"labels\", \"text\"])\n",
    "\n",
    "def calibration_func(model):\n",
    "    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)\n",
    "    dataloader = DataLoader(calibration_dataset, batch_size=8, collate_fn=data_collator)\n",
    "    for batch in dataloader:\n",
    "        try:\n",
    "            model(**batch)\n",
    "        except Exception:\n",
    "            continue\n",
    "\n",
    "# --- 4. Define the quantization configuration and run quantization ---\n",
    "from optimum.intel.neural_compressor import INCConfig\n",
    "# This is a default static quantization configuration\n",
    "quantization_config = INCConfig(quantization={\"approach\": \"static\"})\n",
    "\n",
    "static_quantized_model_dir = \"./quantized-bert-emotion-static-optimum\"\n",
    "\n",
    "quantizer.quantize(\n",
    "    quantization_config=quantization_config,\n",
    "    calibration_function=calibration_func,\n",
    "    save_directory=static_quantized_model_dir,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a26880e-6b30-4b81-8d93-a5fad4d89c3f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (py10)",
   "language": "python",
   "name": "py10"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}