opt_shap.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": [
    "from pprint import pprint\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import shap\n",
    "import torch\n",
    "from datasets import load_dataset\n",
    "from peft import LoraConfig, TaskType, get_peft_model\n",
    "from torch import nn\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorWithPadding,\n",
    "    EvalPrediction,\n",
    "    OPTForSequenceClassification,\n",
    "    Pipeline,\n",
    ")\n",
    "\n",
    "import wandb\n",
    "\n",
    "MODEL = \"facebook/opt-350m\"\n",
    "MAX_POSITION_EMBEDDINGS = 2048\n",
    "\n",
    "from dataclasses import dataclass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "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\n",
    "\n",
    "print(device)"
   ]
  },
  {
   "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": [
      "You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.\n",
      "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",
      "The BetterTransformer implementation does not support padding during training, as the fused kernels do not support attention masks. Beware that passing padded batched data during training may result in unexpected outputs. Please refer to https://huggingface.co/docs/optimum/bettertransformer/overview for more details.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "OPTForSequenceClassification(\n",
       "  (model): OPTModel(\n",
       "    (decoder): OPTDecoder(\n",
       "      (embed_tokens): Embedding(50272, 512, padding_idx=1)\n",
       "      (embed_positions): OPTLearnedPositionalEmbedding(2050, 1024)\n",
       "      (project_out): Linear(in_features=1024, out_features=512, bias=False)\n",
       "      (project_in): Linear(in_features=512, out_features=1024, bias=False)\n",
       "      (layers): ModuleList(\n",
       "        (0-23): 24 x OPTDecoderLayer(\n",
       "          (self_attn): OptFlashAttention2(\n",
       "            (k_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (v_proj): lora.Linear(\n",
       "              (base_layer): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (lora_dropout): ModuleDict(\n",
       "                (default): Dropout(p=0.05, inplace=False)\n",
       "              )\n",
       "              (lora_A): ModuleDict(\n",
       "                (default): Linear(in_features=1024, out_features=16, bias=False)\n",
       "              )\n",
       "              (lora_B): ModuleDict(\n",
       "                (default): Linear(in_features=16, out_features=1024, bias=False)\n",
       "              )\n",
       "              (lora_embedding_A): ParameterDict()\n",
       "              (lora_embedding_B): ParameterDict()\n",
       "            )\n",
       "            (q_proj): lora.Linear(\n",
       "              (base_layer): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (lora_dropout): ModuleDict(\n",
       "                (default): Dropout(p=0.05, inplace=False)\n",
       "              )\n",
       "              (lora_A): ModuleDict(\n",
       "                (default): Linear(in_features=1024, out_features=16, bias=False)\n",
       "              )\n",
       "              (lora_B): ModuleDict(\n",
       "                (default): Linear(in_features=16, out_features=1024, bias=False)\n",
       "              )\n",
       "              (lora_embedding_A): ParameterDict()\n",
       "              (lora_embedding_B): ParameterDict()\n",
       "            )\n",
       "            (out_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          )\n",
       "          (activation_fn): ReLU()\n",
       "          (self_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (score): Linear(in_features=512, out_features=51, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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",
    "    torch_dtype=torch.bfloat16,\n",
    "    attn_implementation=\"flash_attention_2\",\n",
    ").to(device)\n",
    "\n",
    "model.load_adapter(CHECKPOINT_DIR)\n",
    "model.to_bettertransformer()"
   ]
  },
  {
   "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": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tokenizer(\n",
    "    untokenized_dataset[\"train\"][0][\"text\"],\n",
    "    return_tensors=\"pt\",\n",
    "    truncation=True,\n",
    "    max_length=MAX_POSITION_EMBEDDINGS,\n",
    ").to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    logits = model(**inputs).logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class OPT_ICD9_Pipeline(Pipeline):\n",
    "    def _sanitize_parameters(self, **kwargs):\n",
    "        preprocess_kwargs = {}\n",
    "        if \"maybe_arg\" in kwargs:\n",
    "            preprocess_kwargs[\"maybe_arg\"] = kwargs[\"maybe_arg\"]\n",
    "        return preprocess_kwargs, {}, {}\n",
    "\n",
    "    def preprocess(self, text):\n",
    "        return self.tokenizer(\n",
    "            text,\n",
    "            truncation=True,\n",
    "            max_length=MAX_POSITION_EMBEDDINGS,\n",
    "            return_tensors=\"pt\",\n",
    "        )\n",
    "\n",
    "    def _forward(self, model_inputs):\n",
    "        outputs = self.model(**model_inputs)\n",
    "        return outputs\n",
    "\n",
    "    def postprocess(self, model_outputs):\n",
    "        # logits = model_outputs.logits[0].numpy()\n",
    "        #print(logits)\n",
    "        probs = model_outputs[\"logits\"].sigmoid()\n",
    "\n",
    "        output = []\n",
    "        for i, prob in enumerate(probs[0]):\n",
    "            label = self.model.config.id2label[i]\n",
    "            score = prob\n",
    "            output.append({\"label\": label, \"score\": score})\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = OPT_ICD9_Pipeline(model=model, tokenizer=tokenizer, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline(untokenized_dataset[\"train\"][2][\"text\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "masker = shap.maskers.Text(pipeline.tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = shap.sample(untokenized_dataset[\"train\"][\"text\"], 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "explainer = shap.Explainer(pipeline, masker)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "untokenized_dataset[\"train\"][:2][\"text\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "PartitionExplainer explainer:  20%|██        | 1/5 [00:00<?, ?it/s]"
     ]
    }
   ],
   "source": [
    "shap_values = explainer(untokenized_dataset[\"train\"][:5][\"text\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shap.plots.text(shap_values[0, :, \"d-2749\"])"
   ]
  }
 ],
 "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
}