inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/h/u9/c0/00/kuochihh/year_4/CSC413-Project/.venv/lib/python3.11/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"
     ]
    }
   ],
   "source": [
    "import argparse\n",
    "from pprint import pprint\n",
    "\n",
    "import evaluate\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from torch import nn\n",
    "from datasets import load_dataset\n",
    "from peft import LoraConfig, TaskType, get_peft_model\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoTokenizer,\n",
    "    AutoModelForSequenceClassification,\n",
    "    DataCollatorWithPadding,\n",
    "    EvalPrediction,\n",
    "    OPTForSequenceClassification,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    ")\n",
    "import wandb\n",
    "\n",
    "MODEL = \"facebook/opt-350m\"\n",
    "MAX_POSITION_EMBEDDINGS = 2048\n",
    "\n",
    "from dataclasses import dataclass\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "torch.backends.cuda.matmul.allow_tf32 = True\n",
    "torch.backends.cudnn.allow_tf32 = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "CHECKPOINT_DIR = \"OPT-350m-mimic-full\"\n",
    "VAL_DATASET_PATH = \"data/val_9.csv\"\n",
    "CODE_PATH = \"data/icd9_codes.csv\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load dataset\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL, use_fast=True, device=device)\n",
    "\n",
    "code_labels = pd.read_csv(\"data/icd9_codes.csv\")\n",
    "dataset = load_dataset(\"csv\", data_files=VAL_DATASET_PATH)\n",
    "\n",
    "# Create class dictionaries\n",
    "classes = [class_ for class_ in code_labels[\"icd_code\"] if class_]\n",
    "class2id = {class_: id for id, class_ in enumerate(classes)}\n",
    "id2class = {id: class_ for class_, id in class2id.items()}\n",
    "\n",
    "\n",
    "def multi_labels_to_ids(labels: list[str]) -> list[float]:\n",
    "    ids = [0.0] * len(class2id)  # BCELoss requires float as target type\n",
    "    for label in labels:\n",
    "        ids[class2id[label]] = 1.0\n",
    "    return ids\n",
    "\n",
    "\n",
    "def preprocess_function(example):\n",
    "    result = tokenizer(\n",
    "        example[\"text\"], truncation=True, max_length=MAX_POSITION_EMBEDDINGS\n",
    "    )\n",
    "    result[\"labels\"] = [multi_labels_to_ids(eval(label)) for label in example[\"labels\"]]\n",
    "    return result\n",
    "\n",
    "\n",
    "dataset = dataset.map(\n",
    "    preprocess_function, load_from_cache_file=True, batched=True, num_proc=8\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of OPTForSequenceClassification were not initialized from the model checkpoint at facebook/opt-350m and are newly initialized: ['score.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": [
    "config, unused_kwargs = AutoConfig.from_pretrained(\n",
    "    MODEL,\n",
    "    num_labels=len(classes),\n",
    "    id2label=id2class,\n",
    "    label2id=class2id,\n",
    "    problem_type=\"multi_label_classification\",\n",
    "    return_unused_kwargs=True,\n",
    ")\n",
    "\n",
    "if unused_kwargs:\n",
    "    print(f\"Unused kwargs: {unused_kwargs}\")\n",
    "\n",
    "model = OPTForSequenceClassification.from_pretrained(\n",
    "    MODEL,\n",
    "    config=config,\n",
    ").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_adapter(CHECKPOINT_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "untokenized_dataset = load_dataset(\"csv\", data_files=VAL_DATASET_PATH)\n",
    "\n",
    "print(untokenized_dataset['train'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tokenizer(untokenized_dataset[\"train\"][0]['text'], return_tensors=\"pt\", truncation=True, max_length=MAX_POSITION_EMBEDDINGS).to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    logits = model(**inputs).logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logits = logits.to('cpu')\n",
    "\n",
    "predicted_class_ids = torch.arange(0, logits.shape[-1])[\n",
    "    torch.sigmoid(logits).squeeze(dim=0) > 0.5\n",
    "]\n",
    "\n",
    "# Get the predicted class names\n",
    "for id in predicted_class_ids:\n",
    "    predicted_class = id2class[int(id)]\n",
    "    pprint(code_labels[code_labels.icd_code == predicted_class])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, logit in enumerate(logits[0]):\n",
    "    pprint(f'{classes[i]}, {logit}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lime\n",
    "from lime import lime_text\n",
    "from lime.lime_text import LimeTextExplainer\n",
    "from lime.lime_text import IndexedString\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "from time import time\n",
    "\n",
    "\n",
    "explainer = LimeTextExplainer(class_names=classes, bow=False)\n",
    "\n",
    "def predictor_opt(texts):\n",
    "    tk = tokenizer(texts, return_tensors=\"pt\",truncation=True, max_length=MAX_POSITION_EMBEDDINGS).to(device)\n",
    "    outputs = model(**tk)\n",
    "    tensor_logits = outputs[0]\n",
    "    probas = F.sigmoid(tensor_logits).detach().cpu().numpy()\n",
    "    return probas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "sentence = untokenized_dataset[\"train\"][2][\"text\"]\n",
    "n_samples = 10\n",
    "k = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    exp_bert = explainer.explain_instance(\n",
    "        sentence, predictor_opt, num_samples=n_samples, top_labels=k\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_bert.show_in_notebook(text=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "untokenized_dataset[\"train\"][2]['labels']"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}