notebooks/distilbert_imdb.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# DistilBERT calibration experiment on IMDB Dataset\n",
    "\n",
    "We compare multiple calibration algorithms with their reduced, \n",
    "class-wise and class-wise reduced counterparts. We use for this experiment 4 cross-validation splits\n",
    "and 20 bins for ECE and cwECE instead of 5 and 25, respectively, like for the other experiments\n",
    "because of the low number of samples in the dataset.\n",
    "\n",
    "We work with a pre-trained DistilBERT classifier trained by distilling the\n",
    "BERT base model and then fine-tuned on the [IMDB Dataset](http://ai.stanford.edu/~amaas/data/sentiment/),\n",
    "a dataset for binary sentiment classification consisting of 50000 movie reviews.\n",
    "\n",
    "The model achieves an accuracy of roughly 92% of the test set.\n",
    "\n",
    "Since the model's accuracy is pretty high it is, as expected, well calibrated\n",
    "(pre-calibration ECE ≈ 0.043, post-calibration ECE <= 0.02).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "from typing import List\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torchtext\n",
    "from kyle.evaluation import EvalStats\n",
    "from sklearn.metrics import accuracy_score\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm.auto import tqdm\n",
    "from transformers import DistilBertForSequenceClassification, DistilBertTokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is needed for notebooks in case jupyter is started directly in the notebooks directory\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_working_directory = Path(\".\").resolve()\n",
    "if current_working_directory.name == \"notebooks\":\n",
    "    sys.path.insert(0, os.fspath(current_working_directory.parent))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.constants import DATA_DIR, OUTPUT_DIR, RANDOM_SEED\n",
    "from src.utils import (\n",
    "    configure_plots,\n",
    "    perform_default_evaluation,\n",
    "    plot_evaluation_results_from_dataframe,\n",
    "    set_random_seed,\n",
    ")\n",
    "from src.utils.evaluation import combined_results_into_dataframe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constants\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = OUTPUT_DIR / \"distilbert_imdb\"\n",
    "output_dir.mkdir(exist_ok=True)\n",
    "output_file = output_dir / \"results.csv\"\n",
    "\n",
    "imdb_dir = DATA_DIR / \"distilbert_imdb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"textattack/distilbert-base-uncased-imdb\"\n",
    "batch_size = 64\n",
    "classes = [\"neg\", \"pos\"]\n",
    "n_classes = len(classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_random_seed(RANDOM_SEED)\n",
    "configure_plots()\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "logging.basicConfig(level=logging.INFO, format=\"%(message)s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = DistilBertTokenizer.from_pretrained(model_name)\n",
    "datapipe = torchtext.datasets.IMDB(root=os.fspath(imdb_dir), split=\"test\")\n",
    "datapipe = datapipe.batch(batch_size).rows2columnar([\"label\", \"text\"])\n",
    "datapipe = datapipe.map(\n",
    "    lambda x: [\n",
    "        tokenizer(x[\"text\"], return_tensors=\"pt\", truncation=True, padding=True),\n",
    "        torch.LongTensor(list(map(lambda k: classes.index(k), x[\"label\"]))),\n",
    "    ]\n",
    ")\n",
    "dataloader = DataLoader(datapipe, batch_size=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = DistilBertForSequenceClassification.from_pretrained(model_name)\n",
    "model = model.to(device)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating Calibration\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.info(\"Generating model predictions on test set\")\n",
    "\n",
    "all_logits = []\n",
    "y_true = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for inputs, labels in tqdm(dataloader, total=25000 // batch_size):\n",
    "        y_true.append(labels)\n",
    "        inputs = inputs.to(device)\n",
    "        output = model(**inputs)\n",
    "        logits = output.logits\n",
    "        logits = logits.to(\"cpu\")\n",
    "        all_logits.append(logits)\n",
    "\n",
    "all_logits = torch.cat(all_logits)\n",
    "uncalibrated_confidences = F.softmax(all_logits, dim=1).numpy()\n",
    "y_true = torch.cat(y_true).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = np.argmax(uncalibrated_confidences, axis=1)\n",
    "model_accuracy = accuracy_score(y_true, y_pred)\n",
    "logger.info(f\"Model accuracy: {model_accuracy*100}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_stats = EvalStats(y_true, uncalibrated_confidences, bins=25)\n",
    "logger.info(f\"ECE before calibration: {eval_stats.expected_calibration_error()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recalibration\n",
    "We evaluate reduction wrappers on multiple metrics with different calibration algorithms\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Balanced\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.info(\"Performing evaluation\")\n",
    "\n",
    "eval_results = perform_default_evaluation(\n",
    "    confidences=uncalibrated_confidences,\n",
    "    gt_labels=y_true,\n",
    "    cv=4,\n",
    "    bins=20,\n",
    ")\n",
    "\n",
    "results_df = combined_results_into_dataframe(\n",
    "    eval_results,\n",
    "    model_name=\"DistilBERT\",\n",
    "    dataset_name=\"IMDB\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reduction_methods_order: List[str] = results_df[\"Reduction Method\"].unique().tolist()\n",
    "reduction_methods_order = [reduction_methods_order[0]] + sorted(\n",
    "    reduction_methods_order[1:], key=len\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save Results\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.info(\"Saving results\")\n",
    "results_df.to_csv(output_file, sep=\";\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plots\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.info(\"Plotting results\")\n",
    "\n",
    "results_df = results_df.query(\"(Metric != 'condition') & (Metric != 'weak_condition')\")\n",
    "\n",
    "plot_evaluation_results_from_dataframe(\n",
    "    results_df,\n",
    "    hue_order=reduction_methods_order,\n",
    "    output_file=(output_dir / \"evaluation_ECE_distilbert_imdb.eps\"),\n",
    "    show=False,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}