bert_shap_faithfulness.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluating Faithfulness on our model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from transformers import AutoTokenizer, AutoModel, AutoConfig\n",
    "import torch\n",
    "from datasets import load_dataset\n",
    "import torch.nn.functional as F\n",
    "import shap\n",
    "import shap\n",
    "from transformers import Pipeline\n",
    "\n",
    "import os \n",
    "import numpy\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "# device = \"cpu\"\n",
    "data_dir = \"output/\"\n",
    "destination_dir = \"./\"\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset loaded?\n"
     ]
    }
   ],
   "source": [
    "\n",
    "test_short_path = \"data/test_10_top50_short.csv\"\n",
    "labels_10_top50 = pd.read_csv('data/icd10_codes_top50.csv')\n",
    "code_labels_10 = pd.read_csv(\"data/icd10_codes.csv\")\n",
    "print(\"dataset loaded?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model Parameters\n",
    "MAX_POSITION_EMBEDDINGS = 512\n",
    "MODEL = \"emilyalsentzer/Bio_ClinicalBERT\"\n",
    "CKPT = os.path.join(data_dir,\"best_model_state.bin\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "classes\n",
      "bert model and tokenizer initialized\n"
     ]
    }
   ],
   "source": [
    "# Create class dictionaries\n",
    "classes = [class_ for class_ in code_labels_10[\"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",
    "print(\"classes\")\n",
    "\n",
    "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",
    "tokenizer_bert = AutoTokenizer.from_pretrained(MODEL)\n",
    "model_bert = AutoModel.from_pretrained(MODEL, config=config, cache_dir='./model_ckpt/')\n",
    "print(\"bert model and tokenizer initialized\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset loaded\n"
     ]
    }
   ],
   "source": [
    "class TokenizerWrapper:\n",
    "    def __init__(self, tokenizer, length, classes):\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_length = length\n",
    "        self.classes = classes\n",
    "        self.class2id = {class_: id for id, class_ in enumerate(self.classes)}\n",
    "        self.id2class = {id: class_ for class_, id in self.class2id.items()}\n",
    "        \n",
    "    def multi_labels_to_ids(self, labels: list[str]) -> list[float]:\n",
    "        ids = [0.0] * len(self.class2id)  # BCELoss requires float as target type\n",
    "        for label in labels:\n",
    "            ids[self.class2id[label]] = 1.0\n",
    "        return ids\n",
    "    \n",
    "    def tokenize_function(self, example):\n",
    "        result = self.tokenizer(\n",
    "            example[\"text\"],\n",
    "            max_length = self.max_length,\n",
    "            padding = 'max_length',\n",
    "            truncation = True,\n",
    "            return_tensors='pt'\n",
    "        ).to(device)\n",
    "        result[\"label\"] = torch.tensor([self.multi_labels_to_ids(eval(label)) for label in example[\"label\"]])\n",
    "        return result\n",
    "        \n",
    "data_files = {\n",
    "        \"test\": test_short_path,\n",
    "    }\n",
    "\n",
    "tokenizer_wrapper = TokenizerWrapper(tokenizer_bert, MAX_POSITION_EMBEDDINGS, classes)\n",
    "dataset = load_dataset(\"csv\", data_files=data_files)\n",
    "dataset = dataset.map(tokenizer_wrapper.tokenize_function, batched=True, num_proc=1)\n",
    "dataset = dataset.with_format(\"torch\")\n",
    "print(\"dataset loaded\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class BERTClass(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BERTClass, self).__init__()\n",
    "        self.config = config\n",
    "        self.device = device\n",
    "        self.bert_model = model_bert\n",
    "        self.can_generate = model_bert.can_generate\n",
    "        self.base_model_prefix = model_bert.base_model_prefix\n",
    "        self.get_input_embeddings = model_bert.get_input_embeddings\n",
    "        self.dropout = torch.nn.Dropout(0.2)\n",
    "        self.linear = torch.nn.Linear(self.bert_model.config.hidden_size, 50)\n",
    "    \n",
    "    def forward(self, input_ids, attn_mask, token_type_ids):\n",
    "        output = self.bert_model(\n",
    "            input_ids, \n",
    "            attention_mask=attn_mask, \n",
    "            token_type_ids=token_type_ids\n",
    "        )\n",
    "        output_dropout = self.dropout(output.pooler_output)\n",
    "        output = self.linear(output_dropout)\n",
    "        return output\n",
    "    \n",
    "model_bert = BERTClass()\n",
    "model_bert.load_state_dict(torch.load(CKPT))\n",
    "model_bert = model_bert.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pipeline Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BERT_ICD10_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",
    "            max_length = MAX_POSITION_EMBEDDINGS,\n",
    "            padding = 'max_length',\n",
    "            truncation = True,\n",
    "            return_tensors='pt'\n",
    "        ).to(self.device)\n",
    "\n",
    "    def _forward(self, model_inputs):\n",
    "        ids = model_inputs['input_ids'].to(self.device, dtype = torch.long)\n",
    "        mask = model_inputs['attention_mask'].to(self.device, dtype = torch.long)\n",
    "        token_type_ids = model_inputs['token_type_ids'].to(self.device, dtype = torch.long)\n",
    "        outputs = self.model(ids, mask, token_type_ids).to(self.device)\n",
    "        return outputs\n",
    "\n",
    "    def postprocess(self, model_outputs):\n",
    "        probs = F.sigmoid(model_outputs).detach().cpu().numpy() # if there's more than one possible diagnosis\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",
    "        # print(output)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test code for faithfulness calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pipeline initialized\n"
     ]
    }
   ],
   "source": [
    "\n",
    "pipeline = BERT_ICD10_Pipeline(model=model_bert, tokenizer=tokenizer_bert, device = device)\n",
    "print(\"pipeline initialized\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# shap.sample(shap_input, 2)\n",
    "# shap_values = explainer(\n",
    "#         shap_input,\n",
    "#         batch_size=5,\n",
    "#         outputs=shap.Explanation.argsort.flip[:2]\n",
    "#         )\n",
    "# print(shap_values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_index_arrays_shap(inputs, pred_func, model, tokenizer, top_k = 5):\n",
    "    \"\"\" Function to create the arrays corresponding to the shap \n",
    "    \n",
    "    The output is in the format [[input_index_0, input_index_0, ... input_index_n, input_index_n], \n",
    "    [rationale_token_index_0 (for input 0), rationale_token_index_1 (for input 0), ... rationale_token_index_k-1 (for input n), rationale_token_index_k (for input n)]]. \n",
    "    This is used as an indexing array for masking.\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    # get the shap values over the inputs\n",
    "    shap_values = explainer(inputs, batch_size=5)\n",
    "    \n",
    "    # get the mode inferences for the inputs\n",
    "    inferences = pred_func(inputs, model, tokenizer)\n",
    "    indices_array = None\n",
    "    # get the longest \n",
    "    \n",
    "    for i, val in enumerate(shap_values):\n",
    "        # get the choosen labels\n",
    "        print(\"Inferences: \", inferences)\n",
    "        choosen_labels = np.where(inferences[i] > 0.5)\n",
    "        choosen_labels = np.unique(choosen_labels)\n",
    "        \n",
    "        # convert the indices to labels\n",
    "        choosen_labels = np.array(choosen_labels).astype(int)\n",
    "        print(\"Choosen label: \", choosen_labels)\n",
    "        choosen_labels = [id2class[label] for label in choosen_labels]\n",
    "        print(id2class)\n",
    "        print(choosen_labels)\n",
    "        \n",
    "        # for each shap value, index in via it's choosen labels\n",
    "        total_top_k_indices = np.array([])\n",
    "        top_val = min(top_k, len(choosen_labels))\n",
    "        print(top_val)\n",
    "        print(choosen_labels)\n",
    "        \n",
    "        \n",
    "        for label in choosen_labels:\n",
    "            # get the top k shap value indices\n",
    "            print(label)\n",
    "            top_k_indices = np.argsort(shap_values[i, :, label].values)[-top_val:]\n",
    "            print(top_k_indices)\n",
    "            total_top_k_indices = np.append(total_top_k_indices, top_k_indices)\n",
    "        \n",
    "        # sort the indices array to be in ascending order\n",
    "        total_top_k_indices = np.sort(total_top_k_indices)\n",
    "        # remove duplicates\n",
    "        total_top_k_indices = np.unique(total_top_k_indices)\n",
    "        # this might be wrong, it seems like shap returns indices outside of the token range\n",
    "        # so I'm not sure if shap is using the same tokenization function as ours.\n",
    "        total_top_k_indices = total_top_k_indices[total_top_k_indices < 2048]\n",
    "        \n",
    "        # create a array of the same shape of total_top_k_indices and fill with value i\n",
    "        index_array = np.full(total_top_k_indices.shape, i)\n",
    "        \n",
    "        if i == 0:\n",
    "            indices_array = [index_array.tolist(), total_top_k_indices.tolist()]\n",
    "        else:\n",
    "            # append index array to indices array[0]\n",
    "            indices_array[0] = indices_array[0] + index_array.tolist()\n",
    "            # append total_top_k_indices to indices array[1]\n",
    "            indices_array[1] = indices_array[1] + total_top_k_indices.tolist()\n",
    "    \n",
    "    return np.array(indices_array).astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predictor_model_token(texts, model, tokenizer_bert):\n",
    "    # print(len(texts))\n",
    "    # tk = tokenizer(texts, return_tensors=\"pt\",truncation=True, padding=True, max_length=MAX_POSITION_EMBEDDINGS).to(device)\n",
    "    # print(type(tokenizer(texts, return_tensors=\"pt\",truncation=True, padding=True, max_length=MAX_POSITION_EMBEDDINGS)))\n",
    "    # print(\"token_att: \", dir(tk))\n",
    "    tk = tokenizer_bert(\n",
    "            texts,\n",
    "            max_length = MAX_POSITION_EMBEDDINGS,\n",
    "            padding = 'max_length',\n",
    "            truncation = True,\n",
    "            return_tensors='pt'\n",
    "        ).to(device)\n",
    "    ids = tk['input_ids'].to(device, dtype = torch.long)\n",
    "    mask = tk['attention_mask'].to(device, dtype = torch.long)\n",
    "    token_type_ids = tk['token_type_ids'].to(device, dtype = torch.long)\n",
    "    outputs = model_bert(ids, mask, token_type_ids)\n",
    "    # tensor_logits = outputs[0]\n",
    "    # probas = tensor_logits.sigmoid().detach().cpu().numpy()\n",
    "    probas = F.sigmoid(outputs).detach().cpu().numpy()\n",
    "    return probas\n",
    "\n",
    "def predictor_model_no_token(texts, model, tokenizer_bert):\n",
    "    # print(len(texts))\n",
    "    # tk = tokenizer(texts, return_tensors=\"pt\",truncation=True, padding=True, max_length=MAX_POSITION_EMBEDDINGS).to(device)\n",
    "    # tokenization is removed but still need to set texts to device\n",
    "    # i'm not sure why this is a list and don't have time to debug\n",
    "    # print(\"Texts_type:\", type(texts))\n",
    "    # print(\"Texts_dir:\",  dir(texts))\n",
    "    # texts.to(device)\n",
    "    # outputs = model(**texts)\n",
    "    # tensor_logits = outputs[0]\n",
    "    # probas = tensor_logits.sigmoid().detach().cpu().numpy()\n",
    "    ids = texts['input_ids'].to(device, dtype = torch.long)\n",
    "    mask = texts['attention_mask'].to(device, dtype = torch.long)\n",
    "    token_type_ids = texts['token_type_ids'].to(device, dtype = torch.long)\n",
    "    outputs = model_bert(ids, mask, token_type_ids)\n",
    "    tensor_logits = outputs\n",
    "    probas = tensor_logits.sigmoid().detach().cpu().numpy()\n",
    "    return probas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import faithfulness\n",
    "# this reimports the library for easy testing in the notebook\n",
    "import importlib\n",
    "import numpy as np\n",
    "importlib.reload(faithfulness)\n",
    "\n",
    "MAX_LEN=512\n",
    "\n",
    "\n",
    "# tokenize the test dataset\n",
    "test_data =  dataset['test']['text'][:1]\n",
    "print(len(test_data))\n",
    "print(len(test_data[0]))\n",
    "\n",
    "masker = shap.maskers.Text(pipeline.tokenizer)\n",
    "explainer = shap.Explainer(pipeline, masker)\n",
    "print(\"shap computed\")\n",
    "\n",
    "inputs = tokenizer_bert(test_data, max_length=MAX_LEN, padding='max_length', truncation=True, return_tensors='pt')\n",
    "print(\"input type: \", type(inputs))\n",
    "\n",
    "indices_array = get_index_arrays_shap(test_data, predictor_model_token, model_bert, tokenizer_bert)\n",
    "print(\"indices_array:\", indices_array)\n",
    "\n",
    "inputs_rationale_removed = faithfulness.remove_rationale_words(inputs, indices_array, join=False, tokenized=True)\n",
    "inputs_other_removed = faithfulness.remove_other_words(inputs, indices_array, join=False, tokenized=True)\n",
    "\n",
    "# print(\"rational removed: \", inputs_rationale_removed)\n",
    "# print(\"other removed: \", inputs_other_removed)\n",
    "print(\"rational removed type: \", type(inputs_rationale_removed))\n",
    "print(\"other removed type: \", type(inputs_other_removed))\n",
    "\n",
    "ind, faith = faithfulness.calculate_faithfulness(inputs, [inputs_rationale_removed], [inputs_other_removed ], model_bert, tokenizer_bert, predictor_model_no_token, tokenized=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import math\n",
    "\n",
    "##################################################################################################################################\n",
    "# Returns the faithfulness results for choice of:\n",
    "# - input_data being the same dimension/characteristics as the test/val/train dataset used for our classifier of choice Bert or OPT.\n",
    "# - start_index being the starting point index of the input_data\n",
    "# - N being the size of input dataset (this is how many texts you want an explanation for)\n",
    "# - B being the size of explanation batch (this is how many texts your machine can explain at a given instance)\n",
    "# - k being the top k features defined as our rationales for explanation\n",
    "# Precondition: pipeline with tokenizer and model correctly initialized along with masker and explainer for SHAP\n",
    "######################################################################################################################################\n",
    "def get_faith_shap(input_data, start_index, N, B, k):\n",
    "    num_steps = math.ceil(N/B)\n",
    "    tail_n =  N % B    \n",
    "    overall_ind = []\n",
    "    overall_faith = []\n",
    "\n",
    "    for i in range(num_steps):\n",
    "        print(i)\n",
    "        if i == (num_steps - 1): # on the final step\n",
    "            end_index = start_index + tail_n\n",
    "            input_subset = input_data[start_index:end_index]\n",
    "            start_index += tail_n\n",
    "            \n",
    "        else:\n",
    "            end_index = start_index + B\n",
    "            input_subset = input_data[start_index:end_index]\n",
    "            start_index += B\n",
    "    \n",
    "        inputs = tokenizer_bert(test_data, max_length=MAX_LEN, padding='max_length', truncation=True, return_tensors='pt')\n",
    "        print(\"input type: \", type(inputs))\n",
    "    \n",
    "        indices_array = get_index_arrays_shap(test_data, predictor_model_token, model_bert, tokenizer_bert, k)\n",
    "        print(\"indices_array:\", indices_array)\n",
    "    \n",
    "        inputs_rationale_removed = faithfulness.remove_rationale_words(inputs, indices_array, join=False, tokenized=True)\n",
    "        inputs_other_removed = faithfulness.remove_other_words(inputs, indices_array, join=False, tokenized=True)\n",
    "    \n",
    "        # print(\"rational removed: \", inputs_rationale_removed)\n",
    "        # print(\"other removed: \", inputs_other_removed)\n",
    "        print(\"rational removed type: \", type(inputs_rationale_removed))\n",
    "        print(\"other removed type: \", type(inputs_other_removed))\n",
    "    \n",
    "        ind, faith = faithfulness.calculate_faithfulness(inputs, [inputs_rationale_removed], [inputs_other_removed ], model_bert, tokenizer_bert, predictor_model_no_token, tokenized=True)\n",
    "        overall_ind.append(ind)\n",
    "        overall_faith.extend(faith)\n",
    "        \n",
    "    return overall_ind, overall_faith, np.mean(overall_faith)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#  For our evaluation, we fix for 10 examples and choose k = 5, 10, 15.\n",
    "start_index = 0\n",
    "N = 10\n",
    "B = 1\n",
    "masker = shap.maskers.Text(pipeline.tokenizer)\n",
    "explainer = shap.Explainer(pipeline, masker)\n",
    "input_data =  dataset['test']['text']\n",
    "\n",
    "k = 5\n",
    "_, overall_faith_5, avg_faith_5 = get_faith_shap(input_data, start_index, N, B, k)\n",
    "k = 10\n",
    "_, overall_faith_10, avg_faith_10 = get_faith_shap(input_data, start_index, N, B, k)\n",
    "k = 15\n",
    "_, overall_faith_15, avg_faith_15 = get_faith_shap(input_data, start_index, N, B, k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "For k = 5:\n",
      "overall_faith: [0.12282419, 0.117213875, 0.064900264, 0.11008728, 0.10610219, 0.21381745, 0.103904635, 0.083802536, 0.09123538, 0.12878683]\n",
      "avg faith: 0.11426747\n",
      "For k = 10:\n",
      "overall_faith: [0.09830999, 0.10666947, 0.06897583, 0.14220554, 0.12074634, 0.122300334, 0.0930727, 0.47959316, 0.056425303, 0.118119985]\n",
      "avg faith: 0.14064187\n",
      "For k = 15:\n",
      "overall_faith: [0.48845753, 0.0691035, 0.18289092, 0.088048644, 0.11520019, 0.11343779, 0.057432324, 0.07895278, 0.093656234, 0.13393429]\n",
      "avg faith: 0.1421114\n"
     ]
    }
   ],
   "source": [
    "print('For k = 5:')\n",
    "print('overall_faith:', overall_faith_5)\n",
    "print('avg faith:', avg_faith_5)\n",
    "print('For k = 10:')\n",
    "print('overall_faith:', overall_faith_10)\n",
    "print('avg faith:', avg_faith_10)\n",
    "print('For k = 15:')\n",
    "print('overall_faith:', overall_faith_15)\n",
    "print('avg faith:', avg_faith_15)"
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}