TaskB - Attribution - Notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "02c2a8c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import plotly\n",
    "import plotly.express as px\n",
    "import numpy as np\n",
    "from transformers import AutoConfig, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForSequenceClassification\n",
    "from pathlib import Path\n",
    "from datasets import Dataset,DatasetDict,load_dataset,load_metric\n",
    "import evaluate\n",
    "import re\n",
    "from sklearn.model_selection import KFold, StratifiedKFold, train_test_split\n",
    "from iterstrat.ml_stratifiers import MultilabelStratifiedKFold\n",
    "from sklearn.metrics import roc_auc_score, average_precision_score\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.optim import AdamW\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset, RandomSampler, SequentialSampler\n",
    "from transformers import get_linear_schedule_with_warmup\n",
    "import huggingface_hub as hf_hub\n",
    "import numpy as np\n",
    "import random\n",
    "import time\n",
    "import GPUtil\n",
    "import wandb\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import config as code_config\n",
    "import captum\n",
    "from captum.attr import LayerIntegratedGradients\n",
    "import json\n",
    "import deepdiff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "722d98c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "# device = \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28b3b5c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ML_TOKENIZER = \"emilyalsentzer/Bio_ClinicalBERT\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad746979",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dict = dict()\n",
    "config = AutoConfig.from_pretrained(ML_TOKENIZER, \\\n",
    "                                    force_download=True)\n",
    "config.num_labels = 20\n",
    "tokenizer = AutoTokenizer.from_pretrained(ML_TOKENIZER, \\\n",
    "                                          do_lower_case=True, \\\n",
    "                                          force_download=True)\n",
    "for split in [0,1,2,3,4]:\n",
    "    ML_CHECKPOINT = \\\n",
    "    f\"suryakiran786/bio-clinicalbert-multilabel-focal-loss-seed-42-complete-data-{split}-roc-pr\"\n",
    "    model = AutoModelForSequenceClassification.from_pretrained(ML_CHECKPOINT, \\\n",
    "                                                               num_labels=num_labels, \\\n",
    "                                                               problem_type=problem_type, \\\n",
    "                                                               force_download=True)\n",
    "    model_dict[split] = model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69c90502",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e54b4159",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_path = Path.cwd().joinpath(\"mediqa-chat-data\",\"TaskB\",\"TaskB-TrainingSet.csv\")\n",
    "validation_path = Path.cwd().joinpath(\"mediqa-chat-data\",\"TaskB\",\"TaskB-ValidationSet.csv\")\n",
    "\n",
    "train_df = pd.read_csv(train_path)\n",
    "valid_df = pd.read_csv(validation_path)\n",
    "merge_df = pd.concat([train_df,valid_df],axis=0,ignore_index=True)\n",
    "# merge_df = merge_df.sample(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82096aa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "merge_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38e3230d",
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold_dict = \\\n",
    "{\n",
    "  \"ALLERGY\": 2.0464603900909424,\n",
    "  \"ASSESSMENT\": 0.522149384021759,\n",
    "  \"CC\": -1.7761383056640625,\n",
    "  \"DIAGNOSIS\": -0.8881509304046631,\n",
    "  \"DISPOSITION\": -0.5025293827056885,\n",
    "  \"EDCOURSE\": -3.002713441848755,\n",
    "  \"EXAM\": 0.526688814163208,\n",
    "  \"FAM/SOCHX\": -0.6969256401062012,\n",
    "  \"GENHX\": -1.4150243997573853,\n",
    "  \"GYNHX\": -4.687745094299316,\n",
    "  \"IMAGING\": -3.692301034927368,\n",
    "  \"IMMUNIZATIONS\": -2.4603474140167236,\n",
    "  \"LABS\": -3.9070963859558105,\n",
    "  \"MEDICATIONS\": -0.04675881192088127,\n",
    "  \"OTHER_HISTORY\": -6.0084547996521,\n",
    "  \"PASTMEDICALHX\": -1.7836127281188965,\n",
    "  \"PASTSURGICAL\": 1.3906686305999756,\n",
    "  \"PLAN\": -1.1809145212173462,\n",
    "  \"PROCEDURES\": -3.0918960571289062,\n",
    "  \"ROS\": 0.8316808938980103\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65b563ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_forward_func(input_ids, token_type_ids=None, \n",
    "                         position_ids=None, attention_mask=None):\n",
    "    \"\"\"Function passed to ig constructors\"\"\"\n",
    "    return model(input_ids=input_ids, \n",
    "                 token_type_ids=token_type_ids, \n",
    "                 position_ids=position_ids, \n",
    "                 attention_mask=attention_mask)[0]  \n",
    "\n",
    "\n",
    "def prepare_input(text):\n",
    "    \"\"\"Prepare ig attribution input: tokenize sample and baseline text.\"\"\"\n",
    "    tokenized_text = tokenizer(text, \\\n",
    "                               return_tensors=\"pt\", \\\n",
    "                               padding=\"max_length\", \\\n",
    "                               max_length = code_config.MULTI_LABEL_MAX_LENGTH, \\\n",
    "                               truncation=True, \\\n",
    "                               return_attention_mask=True)\n",
    "    seq_len = tokenized_text[\"input_ids\"].shape[1]\n",
    "    position_ids = torch.arange(seq_len, dtype=torch.long).unsqueeze(0)\n",
    "\n",
    "    # Construct the baseline (a reference sample).\n",
    "    # A sequence of [PAD] tokens of length equal to that of the processed sample\n",
    "    ref_text = tokenizer.pad_token * (seq_len - 2) # special tokens\n",
    "    tokenized_ref_text = tokenizer(ref_text, return_tensors=\"pt\") \n",
    "    ref_position_ids = torch.arange(seq_len, dtype=torch.long).unsqueeze(0)\n",
    "\n",
    "    return (tokenized_text[\"input_ids\"],\n",
    "            tokenized_text[\"token_type_ids\"], \n",
    "            position_ids,\n",
    "            tokenized_ref_text[\"input_ids\"],\n",
    "            tokenized_ref_text[\"token_type_ids\"], \n",
    "            ref_position_ids,\n",
    "            tokenized_text[\"attention_mask\"])\n",
    "\n",
    "def place_on_device(*tensors):\n",
    "    tensors_device = []\n",
    "    for t in tensors:\n",
    "        tensors_device.append(t.to(device))\n",
    "    return tuple(tensors_device)  \n",
    "\n",
    "def lig_attribute(lig, class_index, input_data):\n",
    "    return lig.attribute(\n",
    "        inputs=input_data[0], \\\n",
    "        baselines=input_data[3],\n",
    "        additional_forward_args=(input_data[1], input_data[2], input_data[6]), \\\n",
    "        return_convergence_delta=True, \\\n",
    "        target=class_index, \\\n",
    "        n_steps=200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3345130c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_preds(sentence,model):\n",
    "\n",
    "    model = model.to(device)\n",
    "    model.eval()\n",
    "    \n",
    "    # Input for lig attributions (model with no special layers configured)\n",
    "    input_data = place_on_device(*prepare_input(sentence))\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        input_ids = input_data[0]\n",
    "        token_type_ids = input_data[1]\n",
    "        attention_mask = input_data[-1]\n",
    "        \n",
    "        output = model(input_ids=input_ids, \\\n",
    "                       token_type_ids=token_type_ids, \\\n",
    "                       attention_mask=attention_mask, \\\n",
    "                       return_dict=True)\n",
    "    logits = output.logits.detach().cpu().squeeze()\n",
    "    \n",
    "    predicted_idx = []\n",
    "    for idx,(section,threshold) in enumerate(threshold_dict.items()):\n",
    "        if logits[idx] > threshold:\n",
    "            predicted_idx.append(idx)\n",
    "    \n",
    "    return predicted_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10f655cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_word_embedding_attribution(sentence,idx,model):\n",
    "    \"\"\"\n",
    "    Getting layer level attributions for given sentence\n",
    "    Shape -> (batch, max_length, embedding_dimensions)\n",
    "    \"\"\"\n",
    "    model = model.to(device)\n",
    "    model.eval()\n",
    "    \n",
    "    # Input for lig attributions (model with no special layers configured)\n",
    "    input_data = place_on_device(*prepare_input(sentence))\n",
    "    \n",
    "    def predict_forward_func(input_ids, token_type_ids=None, \n",
    "                         position_ids=None, attention_mask=None):\n",
    "        \"\"\"Function passed to ig constructors\"\"\"\n",
    "        return model(input_ids=input_ids, \n",
    "                     token_type_ids=token_type_ids, \n",
    "                     position_ids=position_ids, \n",
    "                     attention_mask=attention_mask)[0]  \n",
    "\n",
    "    # 1. Layer: model.bert.embeddings.word_embeddings\n",
    "    lig_we = LayerIntegratedGradients(\n",
    "        predict_forward_func, \n",
    "        model.bert.embeddings.word_embeddings)\n",
    "    layer_attributions_we, _ = lig_attribute(lig_we, idx, input_data)\n",
    "    \n",
    "    return layer_attributions_we"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "129d4b03",
   "metadata": {},
   "outputs": [],
   "source": [
    "def reorder_encounter_data(encounter_data_dict):\n",
    "    \"\"\"\n",
    "    Reorder encounter_data to encounter_id, section_id, utterance_id,utterance,attribution\n",
    "    \"\"\"\n",
    "    encounter_data_ranked = dict()\n",
    "    for idx,(encounter_id,encounter_data) in enumerate(encounter_data_dict.items()):\n",
    "        encounter_data_ranked[encounter_id] = dict()\n",
    "        for utterance_id,utterance_data in encounter_data.items():\n",
    "            utterance = utterance_data[0]\n",
    "            section_dict = utterance_data[1]\n",
    "            for section_id, attribution in section_dict.items():\n",
    "                if section_id not in encounter_data_ranked[encounter_id]:\n",
    "                    encounter_data_ranked[encounter_id][section_id] = []\n",
    "                encounter_data_ranked[encounter_id][section_id].append([utterance_id,utterance,attribution])\n",
    "    return encounter_data_ranked"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20901b50",
   "metadata": {},
   "outputs": [],
   "source": [
    "merge_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b064171",
   "metadata": {},
   "outputs": [],
   "source": [
    "section_dict = dict()\n",
    "for idx,row in merge_df.iterrows():\n",
    "    section_dict[row[\"encounter_id\"]] = dict()\n",
    "    dialogue_list = row[\"dialogue\"].split(\"\\n\")\n",
    "    for utterance_idx,utterance in enumerate(dialogue_list):\n",
    "        section_dict[row[\"encounter_id\"]][utterance_idx] = dict()\n",
    "        utterance_dict = dict()\n",
    "        for split,model in model_dict.items():\n",
    "            preds = get_preds(utterance,model)\n",
    "            if isinstance(preds,int):\n",
    "                preds = [preds]\n",
    "            for pred in preds:\n",
    "                if pred not in utterance_dict:\n",
    "                    utterance_dict[pred] = [split]\n",
    "                else:\n",
    "                    utterance_dict[pred].append(split)\n",
    "            section_dict[row[\"encounter_id\"]][utterance_idx] = (utterance,utterance_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ec71d26",
   "metadata": {},
   "outputs": [],
   "source": [
    "section_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d03a65f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(section_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bc2f706",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get attribution for every prediction\n",
    "if Path(\"encounter_data_with_attribution.json\").exists():\n",
    "    with open(\"encounter_data_with_attribution.json\",\"r\") as f:\n",
    "        encounter_data_ranked = json.load(f)\n",
    "else:\n",
    "    encounter_data_ranked = dict()\n",
    "encounter_data_dict = dict()\n",
    "for encounter_id,encounter_body in tqdm(section_dict.items()):\n",
    "    if encounter_id in encounter_data_ranked:\n",
    "        print(f\"Skipping {encounter_id}\")\n",
    "        continue\n",
    "    encounter_data_dict[encounter_id] = dict()\n",
    "    for utterance_id,utterance_body in encounter_body.items():\n",
    "        utterance = utterance_body[0]\n",
    "        section_attribution_dict = dict()\n",
    "        for section_id,split_list in utterance_body[1].items():\n",
    "            attribute_by_split = []\n",
    "            for split in split_list:\n",
    "                attribution = get_word_embedding_attribution(utterance,section_id,model_dict[split])\n",
    "                attribution = attribution.mean(dim=-1)\n",
    "                attribution = attribution.abs().mean(dim=-1).squeeze(0)\n",
    "                attribute_by_split.append(attribution.item())\n",
    "            section_attribution_dict[str(section_id)] = np.mean(attribute_by_split).item()\n",
    "        encounter_data_dict[str(encounter_id)][str(utterance_id)] = (utterance,section_attribution_dict)\n",
    "    reordered_encounter_data = reorder_encounter_data(encounter_data_dict)\n",
    "    encounter_data_ranked.update(reordered_encounter_data)\n",
    "    with open(\"encounter_data_with_attribution_tmp.json\",\"w\") as f:\n",
    "        json.dump(encounter_data_ranked,f,indent=2)\n",
    "    with open(\"encounter_data_with_attribution_tmp.json\",\"r\") as f:\n",
    "        encounter_data_ranked_2 = json.load(f)\n",
    "    if deepdiff.diff.DeepDiff(encounter_data_ranked,encounter_data_ranked_2) != {}:\n",
    "        raise Exception(\"encounter_data_ranked must be equal to encounter_data_ranked_2\")\n",
    "    else:\n",
    "        with open(\"encounter_data_with_attribution.json\",\"w\") as f:\n",
    "            json.dump(encounter_data_ranked,f,indent=2)\n",
    "    print(f\"{encounter_id} done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcab01a1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}