Amazon.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d90a11d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import argparse\n",
    "import glob\n",
    "import math\n",
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tqdm as tqdm\n",
    "from scipy.special import softmax\n",
    "import scipy.stats as stats\n",
    "\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import display, HTML\n",
    "from captum.attr import visualization\n",
    "\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "import datasets\n",
    "from datasets import load_dataset, load_metric \n",
    "from datasets import list_datasets, list_metrics\n",
    "\n",
    "from BERT_explainability.modules.BERT.BertForSequenceClassification import BertForSequenceClassification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e0661e55",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load model   \n",
    "model = BertForSequenceClassification.from_pretrained(\"fabriceyhc/bert-base-uncased-amazon_polarity\").to(\"cuda\")\n",
    "model.eval()\n",
    "\n",
    "# load tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"fabriceyhc/bert-base-uncased-amazon_polarity\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ded2e972",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = \"amazon_polarity\" # select a dataset\n",
    "dataset = load_dataset(dataset) # load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "636a3f07",
   "metadata": {},
   "outputs": [],
   "source": [
    "special_tokens = {\"[CLS]\", \"[SEP]\"}\n",
    "special_idxs = {101,102}    \n",
    "mask = \"[PAD]\"\n",
    "mask_id = 0   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e85317ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_sample(text):\n",
    "    tokenized_input  = tokenizer(text, add_special_tokens=True, truncation=True)\n",
    "    input_ids = tokenized_input['input_ids']\n",
    "    text_ids = (torch.tensor([input_ids])).to(\"cuda\")\n",
    "    text_words = tokenizer.convert_ids_to_tokens(text_ids[0])\n",
    "    \n",
    "    # mask special tokens\n",
    "    att_mask = tokenized_input['attention_mask']\n",
    "    spe_idxs = [x for x, y in list(enumerate(input_ids)) if y in special_idxs]\n",
    "    att_mask = [0 if index in spe_idxs else 1 for index, item in enumerate(att_mask)]\n",
    "    att_mask = [0 if index in spe_idxs else 1 for index, item in enumerate(att_mask)]\n",
    "    att_mask = (torch.tensor([att_mask])).to(\"cuda\")\n",
    "    \n",
    "    return text_ids, att_mask, text_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "33db358e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(model, text_ids, target, att_mask=None, seg_ids=None):\n",
    "    out = model(text_ids, attention_mask=att_mask, token_type_ids=seg_ids)\n",
    "    prob = out[0]\n",
    "    pred_class = torch.argmax(prob, axis=1).cpu().detach().numpy()\n",
    "    pred_class_prob = softmax(prob.cpu().detach().numpy(), axis=1)\n",
    "    return pred_class[0], pred_class_prob[:, target][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0d7f37ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def truncate_words(sorted_idx, text_words, text_ids, replaced_num, seg_ids=None):\n",
    "    to_be_replaced_idx = []\n",
    "    i= 0\n",
    "    while len(to_be_replaced_idx) < replaced_num and i!=len(text_words)-1:\n",
    "        current_idx = sorted_idx[i]\n",
    "        if text_words[current_idx] not in special_tokens:\n",
    "            to_be_replaced_idx.append(current_idx)\n",
    "        i += 1\n",
    "    remaining_idx = sorted(list(set(sorted_idx) - set(to_be_replaced_idx)))\n",
    "    truncated_text_ids = text_ids[0, np.array(remaining_idx)]\n",
    "    if seg_ids is not None:\n",
    "        seg_ids = seg_ids[0, np.array(remaining_idx)]\n",
    "    truncated_text_words = np.array(text_words)[remaining_idx]\n",
    "    return truncated_text_ids.unsqueeze(0), truncated_text_words, seg_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0e875dd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def replace_words(sorted_idx, text_words, text_ids, replaced_num, mask, mask_id):\n",
    "    to_be_replaced_idx = []\n",
    "    i= 0\n",
    "    while len(to_be_replaced_idx) < replaced_num and i!=len(text_words)-1:\n",
    "        current_idx = sorted_idx[i]\n",
    "        if text_words[current_idx] not in special_tokens:\n",
    "            to_be_replaced_idx.append(current_idx)\n",
    "        i += 1\n",
    "    replaced_text_ids = text_ids.clone()\n",
    "    replaced_text_ids[0, to_be_replaced_idx] = mask_id\n",
    "    replaced_text_words = np.copy(text_words)\n",
    "    replaced_text_words[to_be_replaced_idx] = mask\n",
    "    return replaced_text_ids, replaced_text_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2eb4c203",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute rollout between attention layers\n",
    "def compute_rollout_attention(all_layer_matrices, start_layer=0):\n",
    "    \n",
    "    # adding residual consideration- code adapted from https://github.com/samiraabnar/attention_flow\n",
    "    num_tokens = all_layer_matrices[0].shape[1]\n",
    "    batch_size = all_layer_matrices[0].shape[0]\n",
    "    eye = torch.eye(num_tokens).expand(batch_size, num_tokens, num_tokens).to(all_layer_matrices[0].device)\n",
    "    all_layer_matrices = [all_layer_matrices[i] + eye for i in range(len(all_layer_matrices))]\n",
    "    matrices_aug = [all_layer_matrices[i] / all_layer_matrices[i].sum(dim=-1, keepdim=True)\n",
    "                          for i in range(len(all_layer_matrices))]\n",
    "    joint_attention = matrices_aug[start_layer]\n",
    "    for i in range(start_layer+1, len(matrices_aug)):\n",
    "        joint_attention = matrices_aug[i].bmm(joint_attention)\n",
    "        \n",
    "    return joint_attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "66061f0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator:\n",
    "    def __init__(self, model):\n",
    "        self.model = model\n",
    "        self.model.eval()\n",
    "\n",
    "    def forward(self, input_ids, attention_mask):\n",
    "        return self.model(input_ids, attention_mask)\n",
    "\n",
    "    \n",
    "    def generate_TransCAM(self, input_ids, attention_mask,\n",
    "                          index=None, start_layer=0):\n",
    "\n",
    "        result = self.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)\n",
    "        \n",
    "        output = result[0]\n",
    "        hs = result[1]\n",
    "\n",
    "        kwargs = {\"alpha\": 1}\n",
    "\n",
    "        blocks = self.model.bert.encoder.layer\n",
    "\n",
    "        for blk_id in range(len(blocks)):\n",
    "            hs[blk_id].retain_grad()\n",
    "\n",
    "        if index == None:\n",
    "            index = np.argmax(output.cpu().data.numpy(), axis=-1)\n",
    "\n",
    "        one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32)\n",
    "        one_hot[0, index] = 1\n",
    "        one_hot_vector = one_hot\n",
    "        one_hot = torch.from_numpy(one_hot).requires_grad_(True)\n",
    "        one_hot = torch.sum(one_hot.cuda() * output)\n",
    "\n",
    "        self.model.zero_grad()\n",
    "        one_hot.backward(retain_graph=True)\n",
    "\n",
    "        self.model.relprop(torch.tensor(one_hot_vector).to(input_ids.device), **kwargs)\n",
    "\n",
    "        cams = {}\n",
    "        \n",
    "        for blk_id in range(len(blocks)):\n",
    "            hs_grads = hs[blk_id].grad\n",
    "            \n",
    "            att = blocks[blk_id].attention.self.get_attn().squeeze(0)\n",
    "            att = att.mean(dim=0)\n",
    "            att = att.mean(dim=0)\n",
    "            \n",
    "            cat = (hs_grads * hs[blk_id]).sum(dim=-1).squeeze(0)\n",
    "            cat = cat * att\n",
    "            \n",
    "            cams[blk_id] = cat\n",
    "            \n",
    "        trans_expln = sum(cams.values())\n",
    "\n",
    "        return trans_expln\n",
    "\n",
    "    def generate_LRP(self, input_ids, attention_mask,\n",
    "                     index=None, start_layer=11):\n",
    "        \n",
    "        output = self.model(input_ids=input_ids, attention_mask=attention_mask)[0]\n",
    "        kwargs = {\"alpha\": 1}\n",
    "\n",
    "        if index == None:\n",
    "            index = np.argmax(output.cpu().data.numpy(), axis=-1)\n",
    "\n",
    "        one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32)\n",
    "        one_hot[0, index] = 1\n",
    "        one_hot_vector = one_hot\n",
    "        one_hot = torch.from_numpy(one_hot).requires_grad_(True)\n",
    "        one_hot = torch.sum(one_hot.cuda() * output)\n",
    "\n",
    "        self.model.zero_grad()\n",
    "        one_hot.backward(retain_graph=True)\n",
    "\n",
    "        self.model.relprop(torch.tensor(one_hot_vector).to(input_ids.device), **kwargs)\n",
    "\n",
    "        cams = []\n",
    "        blocks = self.model.bert.encoder.layer\n",
    "        for blk in blocks:\n",
    "            grad = blk.attention.self.get_attn_gradients()\n",
    "            cam = blk.attention.self.get_attn_cam()\n",
    "            cam = cam[0].reshape(-1, cam.shape[-1], cam.shape[-1])\n",
    "            grad = grad[0].reshape(-1, grad.shape[-1], grad.shape[-1])\n",
    "            cam = grad * cam\n",
    "            cam = cam.clamp(min=0).mean(dim=0)\n",
    "            cams.append(cam.unsqueeze(0))\n",
    "        rollout = compute_rollout_attention(cams, start_layer=start_layer)\n",
    "        rollout[:, 0, 0] = 0\n",
    "        return rollout[:, 0]\n",
    "\n",
    "    def generate_LRP_last_layer(self, input_ids, attention_mask,\n",
    "                     index=None):\n",
    "        output = self.model(input_ids=input_ids, attention_mask=attention_mask)[0]\n",
    "        kwargs = {\"alpha\": 1}\n",
    "        if index == None:\n",
    "            index = np.argmax(output.cpu().data.numpy(), axis=-1)\n",
    "\n",
    "        one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32)\n",
    "        one_hot[0, index] = 1\n",
    "        one_hot_vector = one_hot\n",
    "        one_hot = torch.from_numpy(one_hot).requires_grad_(True)\n",
    "        one_hot = torch.sum(one_hot.cuda() * output)\n",
    "\n",
    "        self.model.zero_grad()\n",
    "        one_hot.backward(retain_graph=True)\n",
    "\n",
    "        self.model.relprop(torch.tensor(one_hot_vector).to(input_ids.device), **kwargs)\n",
    "\n",
    "        cam = self.model.bert.encoder.layer[-1].attention.self.get_attn_cam()[0]\n",
    "        cam = cam.clamp(min=0).mean(dim=0).unsqueeze(0)\n",
    "        cam[:, 0, 0] = 0\n",
    "        return cam[:, 0]\n",
    "\n",
    "    def generate_full_lrp(self, input_ids, attention_mask,\n",
    "                     index=None):\n",
    "        output = self.model(input_ids=input_ids, attention_mask=attention_mask)[0]\n",
    "        kwargs = {\"alpha\": 1}\n",
    "\n",
    "        if index == None:\n",
    "            index = np.argmax(output.cpu().data.numpy(), axis=-1)\n",
    "\n",
    "        one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32)\n",
    "        one_hot[0, index] = 1\n",
    "        one_hot_vector = one_hot\n",
    "        one_hot = torch.from_numpy(one_hot).requires_grad_(True)\n",
    "        one_hot = torch.sum(one_hot.cuda() * output)\n",
    "\n",
    "        self.model.zero_grad()\n",
    "        one_hot.backward(retain_graph=True)\n",
    "\n",
    "        cam = self.model.relprop(torch.tensor(one_hot_vector).to(input_ids.device), **kwargs)\n",
    "        cam = cam.sum(dim=2)\n",
    "        cam[:, 0] = 0\n",
    "        return cam\n",
    "\n",
    "    def generate_attn_last_layer(self, input_ids, attention_mask,\n",
    "                     index=None):\n",
    "        output = self.model(input_ids=input_ids, attention_mask=attention_mask)[0]\n",
    "        cam = self.model.bert.encoder.layer[-1].attention.self.get_attn()[0]\n",
    "        cam = cam.mean(dim=0).unsqueeze(0)\n",
    "        cam[:, 0, 0] = 0\n",
    "        return cam[:, 0]\n",
    "\n",
    "    def generate_rollout(self, input_ids, attention_mask, start_layer=0, index=None):\n",
    "        self.model.zero_grad()\n",
    "        output = self.model(input_ids=input_ids, attention_mask=attention_mask)[0]\n",
    "        blocks = self.model.bert.encoder.layer\n",
    "        all_layer_attentions = []\n",
    "        for blk in blocks:\n",
    "            attn_heads = blk.attention.self.get_attn()\n",
    "            avg_heads = (attn_heads.sum(dim=1) / attn_heads.shape[1]).detach()\n",
    "            all_layer_attentions.append(avg_heads)\n",
    "        rollout = compute_rollout_attention(all_layer_attentions, start_layer=start_layer)\n",
    "        rollout[:, 0, 0] = 0\n",
    "        return rollout[:, 0]\n",
    "\n",
    "    def generate_attn_gradcam(self, input_ids, attention_mask, index=None):\n",
    "        output = self.model(input_ids=input_ids, attention_mask=attention_mask)[0]\n",
    "        kwargs = {\"alpha\": 1}\n",
    "\n",
    "        if index == None:\n",
    "            index = np.argmax(output.cpu().data.numpy(), axis=-1)\n",
    "\n",
    "        one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32)\n",
    "        one_hot[0, index] = 1\n",
    "        one_hot_vector = one_hot\n",
    "        one_hot = torch.from_numpy(one_hot).requires_grad_(True)\n",
    "        one_hot = torch.sum(one_hot.cuda() * output)\n",
    "\n",
    "        self.model.zero_grad()\n",
    "        one_hot.backward(retain_graph=True)\n",
    "\n",
    "        self.model.relprop(torch.tensor(one_hot_vector).to(input_ids.device), **kwargs)\n",
    "\n",
    "        cam = self.model.bert.encoder.layer[-1].attention.self.get_attn()\n",
    "        grad = self.model.bert.encoder.layer[-1].attention.self.get_attn_gradients()\n",
    "\n",
    "        cam = cam[0].reshape(-1, cam.shape[-1], cam.shape[-1])\n",
    "        grad = grad[0].reshape(-1, grad.shape[-1], grad.shape[-1])\n",
    "        grad = grad.mean(dim=[1, 2], keepdim=True)\n",
    "        cam = (cam * grad).mean(0).clamp(min=0).unsqueeze(0)\n",
    "        cam = (cam - cam.min()) / (cam.max() - cam.min())\n",
    "        cam[:, 0, 0] = 0\n",
    "        \n",
    "        return cam[:, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "584744a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_explns(explanations, input_ids, attention_mask, start_layer=0, true_class = 1, is_true = True):\n",
    "    \n",
    "    if is_true:\n",
    "        # TransCAM\n",
    "        TransCAM_expln = explanations.generate_TransCAM(input_ids=input_ids, attention_mask=attention_mask, \n",
    "                                                    index=true_class, start_layer=start_layer)\n",
    "        # LRP\n",
    "        LRP_expln = explanations.generate_LRP(input_ids=input_ids, attention_mask=attention_mask, \n",
    "                                                    index=true_class, start_layer=start_layer)[0]\n",
    "        # PartialLRP\n",
    "        PartialLRP_expln = explanations.generate_LRP_last_layer(input_ids=input_ids, \n",
    "                                                     index=true_class, attention_mask=attention_mask)[0]\n",
    "        # FullLRP\n",
    "        # FullLRP_expln = explanations.generate_full_lrp(input_ids=input_ids, attention_mask=attention_mask, \n",
    "        #                                           index=true_class)[0]\n",
    "        # Att\n",
    "        Att_expln = explanations.generate_attn_last_layer(input_ids=input_ids, attention_mask=attention_mask, \n",
    "                                                    index=true_class)[0]\n",
    "        # Rollout\n",
    "        Rollout_expln = explanations.generate_rollout(input_ids=input_ids, attention_mask=attention_mask,      \n",
    "                                                    index=true_class, start_layer=0)[0]\n",
    "        # Att_Gradcam\n",
    "        # Att_Gradcam_expln = explanations.generate_attn_gradcam(input_ids=input_ids, attention_mask=attention_mask, \n",
    "        #                                            index=true_class)[0]    \n",
    "        \n",
    "    else:\n",
    "        true_class = 1-true_class\n",
    "        # TransCAM\n",
    "        TransCAM_expln = explanations.generate_TransCAM(input_ids=input_ids, attention_mask=attention_mask, \n",
    "                                                    index=true_class, start_layer=start_layer)\n",
    "        # LRP\n",
    "        LRP_expln = explanations.generate_LRP(input_ids=input_ids, attention_mask=attention_mask, \n",
    "                                                    index=true_class, start_layer=start_layer)[0]\n",
    "        # PartialLRP\n",
    "        PartialLRP_expln = explanations.generate_LRP_last_layer(input_ids=input_ids, \n",
    "                                                    index=true_class, attention_mask=attention_mask)[0]\n",
    "        # FullLRP\n",
    "        # FullLRP_expln = explanations.generate_full_lrp(input_ids=input_ids, attention_mask=attention_mask, \n",
    "        #                                            index=true_class)[0]\n",
    "        # Att\n",
    "        Att_expln = explanations.generate_attn_last_layer(input_ids=input_ids, attention_mask=attention_mask, \n",
    "                                                    index=true_class)[0]\n",
    "        # Rollout\n",
    "        Rollout_expln = explanations.generate_rollout(input_ids=input_ids, attention_mask=attention_mask, \n",
    "                                                    index=true_class, start_layer=0)[0]\n",
    "        # Att_Gradcam\n",
    "        # Att_Gradcam_expln = explanations.generate_attn_gradcam(input_ids=input_ids, attention_mask=attention_mask, \n",
    "        #                                            index=true_class)[0]\n",
    "    \n",
    "    # return TransCAM_expln, LRP_expln, PartialLRP_expln, FullLRP_expln, Att_expln, Rollout_expln, Att_Gradcam_expln\n",
    "    return TransCAM_expln, Att_expln, PartialLRP_expln, Rollout_expln, LRP_expln"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8e25ad8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cal_aopc(original_probs, degradation_probs):\n",
    "    original_probs = np.array(original_probs)\n",
    "    degradation_probs = np.array(degradation_probs)\n",
    "    \n",
    "    diffs = []\n",
    "    for i in range(len(original_probs)):\n",
    "        diffs_k = []\n",
    "        for j in range(9):\n",
    "            diff = original_probs[i] - degradation_probs[i][j]\n",
    "            diffs_k.append(np.abs(diff))\n",
    "        diffs.append(diffs_k)\n",
    "\n",
    "    result = np.mean(diffs, axis=0)\n",
    "    aopc = np.mean(result)\n",
    "    \n",
    "    return aopc\n",
    "\n",
    "def cal_logodds(original_probs, degradation_probs):\n",
    "    original_probs = np.array(original_probs)\n",
    "    degradation_probs = np.array(degradation_probs)\n",
    "    \n",
    "    ratios = []\n",
    "    for i in range(len(original_probs)):\n",
    "        ratios_k = []\n",
    "        for j in range(9):\n",
    "            ratio = math.log(degradation_probs[i][j] / original_probs[i])\n",
    "            ratios_k.append(ratio)\n",
    "        ratios.append(ratios_k)\n",
    "\n",
    "    result = np.mean(ratios, axis=0)\n",
    "    logodds = np.mean(result)\n",
    "    \n",
    "    return logodds\n",
    "\n",
    "def cal_kendaltau(attribution1, attribution2,):\n",
    "\n",
    "    sorted_idx1 = np.argsort(-attribution1)\n",
    "    sorted_idx2 = np.argsort(-attribution2)\n",
    "\n",
    "    tau, p_value = stats.kendalltau(sorted_idx1, sorted_idx2)\n",
    "    \n",
    "    return tau"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1d2d33cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model, explanations, test_data, degrade_step = 10, seg_ids=None):\n",
    "    \n",
    "    original_probs = []\n",
    "    original_accs = [] \n",
    "    \n",
    "    degradation_probs_TransCAM = []\n",
    "    degradation_accs_TransCAM = []\n",
    "    del_probs_TransCAM = []\n",
    "    del_accs_TransCAM = []\n",
    "    \n",
    "    degradation_probs_Rawatt = []\n",
    "    degradation_accs_Rawatt = []\n",
    "    del_probs_Rawatt = []\n",
    "    del_accs_Rawatt = []\n",
    "    \n",
    "    degradation_probs_PartialLRP = []\n",
    "    degradation_accs_PartialLRP = []\n",
    "    del_probs_PartialLRP = []\n",
    "    del_accs_PartialLRP = []\n",
    "    \n",
    "    degradation_probs_Rollout = []\n",
    "    degradation_accs_Rollout = []\n",
    "    del_probs_Rollout = []\n",
    "    del_accs_Rollout = []\n",
    "    \n",
    "    degradation_probs_LRP = []\n",
    "    degradation_accs_LRP = []\n",
    "    del_probs_LRP = []\n",
    "    del_accs_LRP = []\n",
    "    \n",
    "    kendaltaus_TransCAM = []\n",
    "    kendaltaus_Rawatt = []\n",
    "    kendaltaus_PartialLRP = []\n",
    "    kendaltaus_Rollout = []\n",
    "    kendaltaus_LRP = []\n",
    "    \n",
    "    for i, test_instance in enumerate(test_data):\n",
    "        \n",
    "        text = test_instance['content']\n",
    "        target = test_instance['label'] \n",
    "\n",
    "        text_ids, att_mask, text_words = preprocess_sample(text)\n",
    "        \n",
    "        # get truc words number\n",
    "        total_len = len(text_words)\n",
    "        if total_len< 10: \n",
    "            continue\n",
    "        granularity = np.linspace(0, 1, degrade_step)\n",
    "        trunc_words_num = [int(g) for g in np.round(granularity*total_len)]\n",
    "        trunc_words_num = list(dict.fromkeys(trunc_words_num))\n",
    "        \n",
    "        original_class, original_prob = predict(model, text_ids, target)\n",
    "        \n",
    "        # get attributions\n",
    "        attribution_TransCAM, attribution_Rawatt, attribution_PartialLRP, attribution_Rollout, attribution_LRP = \\\n",
    "        generate_explns(explanations, text_ids, att_mask, start_layer=0, true_class = target, is_true = True) \n",
    "        \n",
    "        attribution_TransCAM_F, attribution_Rawatt_F, attribution_PartialLRP_F, attribution_Rollout_F, attribution_LRP_F = \\\n",
    "        generate_explns(explanations, text_ids, att_mask, start_layer=0, true_class = target, is_true = False) \n",
    "        \n",
    "        # conver to cpu numpy\n",
    "        attribution_TransCAM = attribution_TransCAM.cpu().detach().numpy()\n",
    "        attribution_Rawatt = attribution_Rawatt.cpu().detach().numpy()\n",
    "        attribution_PartialLRP = attribution_PartialLRP.cpu().detach().numpy()\n",
    "        attribution_Rollout = attribution_Rollout.cpu().detach().numpy()\n",
    "        attribution_LRP = attribution_LRP.cpu().detach().numpy()\n",
    "        \n",
    "        # conver to cpu numpy\n",
    "        attribution_TransCAM_F = attribution_TransCAM_F.cpu().detach().numpy()\n",
    "        attribution_Rawatt_F = attribution_Rawatt_F.cpu().detach().numpy()\n",
    "        attribution_PartialLRP_F = attribution_PartialLRP_F.cpu().detach().numpy()\n",
    "        attribution_Rollout_F = attribution_Rollout_F.cpu().detach().numpy()\n",
    "        attribution_LRP_F = attribution_LRP_F.cpu().detach().numpy()\n",
    "        \n",
    "        # get sorted_idx\n",
    "        sorted_idx_TransCAM = np.argsort(-attribution_TransCAM)\n",
    "        sorted_idx_Rawatt = np.argsort(-attribution_Rawatt)\n",
    "        sorted_idx_PartialLRP = np.argsort(-attribution_PartialLRP)\n",
    "        sorted_idx_Rollout = np.argsort(-attribution_Rollout)\n",
    "        sorted_idx_LRP = np.argsort(-attribution_LRP)\n",
    "        \n",
    "        instance_degradation_probs_TransCAM = []\n",
    "        instance_degradation_accs_TransCAM = []\n",
    "        instance_replace_probs_TransCAM = []\n",
    "        instance_replace_accs_TransCAM = []\n",
    "        \n",
    "        instance_degradation_probs_Rawatt = []\n",
    "        instance_degradation_accs_Rawatt = []\n",
    "        instance_replace_probs_Rawatt = []\n",
    "        instance_replace_accs_Rawatt = []\n",
    "        \n",
    "        instance_degradation_probs_PartialLRP = []\n",
    "        instance_degradation_accs_PartialLRP = []\n",
    "        instance_replace_probs_PartialLRP = []\n",
    "        instance_replace_accs_PartialLRP = []\n",
    "        \n",
    "        instance_degradation_probs_Rollout = []\n",
    "        instance_degradation_accs_Rollout = []\n",
    "        instance_replace_probs_Rollout = []\n",
    "        instance_replace_accs_Rollout = []\n",
    "        \n",
    "        instance_degradation_probs_LRP = []\n",
    "        instance_degradation_accs_LRP = []\n",
    "        instance_replace_probs_LRP = []\n",
    "        instance_replace_accs_LRP = []\n",
    "\n",
    "        for num in trunc_words_num[1:]: #exclude 0\n",
    "            \n",
    "            # TransCAM\n",
    "            truncated_text_ids_TransCAM, _, _ = truncate_words(sorted_idx_TransCAM, text_words, text_ids, \n",
    "                                                                                        num, seg_ids=seg_ids)\n",
    "            replaced_text_ids_TransCAM, _ = replace_words(sorted_idx_TransCAM, text_words, text_ids, num, mask, mask_id)\n",
    "           \n",
    "            trunc_class_TransCAM, trunc_prob_TransCAM = predict(model, truncated_text_ids_TransCAM, target, seg_ids=seg_ids)\n",
    "            rep_class_TransCAM, rep_prob_TransCAM = predict(model, replaced_text_ids_TransCAM, target, seg_ids=seg_ids)\n",
    "\n",
    "            instance_degradation_probs_TransCAM.append(trunc_prob_TransCAM)\n",
    "            instance_degradation_accs_TransCAM.append(trunc_class_TransCAM==target)\n",
    "            \n",
    "            instance_replace_probs_TransCAM.append(rep_prob_TransCAM)\n",
    "            instance_replace_accs_TransCAM.append(rep_class_TransCAM==target)\n",
    "            \n",
    "            # Rawatt\n",
    "            truncated_text_ids_Rawatt, _, _ = truncate_words(sorted_idx_Rawatt, text_words, text_ids, \n",
    "                                                                                       num, seg_ids=seg_ids)\n",
    "            replaced_text_ids_Rawatt, _ = replace_words(sorted_idx_Rawatt, text_words, text_ids, num, mask, mask_id)\n",
    "            \n",
    "            \n",
    "            trunc_class_Rawatt, trunc_prob_Rawatt = predict(model, truncated_text_ids_Rawatt, target, seg_ids=seg_ids)\n",
    "            rep_class_Rawatt, rep_prob_Rawatt = predict(model, replaced_text_ids_Rawatt, target, seg_ids=seg_ids)\n",
    "\n",
    "            instance_degradation_probs_Rawatt.append(trunc_prob_Rawatt)\n",
    "            instance_degradation_accs_Rawatt.append(trunc_class_Rawatt==target)\n",
    "            \n",
    "            instance_replace_probs_Rawatt.append(rep_prob_Rawatt)\n",
    "            instance_replace_accs_Rawatt.append(rep_class_Rawatt==target)\n",
    "            \n",
    "            # PartialLRP\n",
    "            truncated_text_ids_PartialLRP, _, _ = truncate_words(sorted_idx_PartialLRP, text_words, text_ids, \n",
    "                                                                                       num, seg_ids=seg_ids)\n",
    "            replaced_text_ids_PartialLRP, _ = replace_words(sorted_idx_PartialLRP, text_words, text_ids, num, mask, mask_id)\n",
    "            \n",
    "            trunc_class_PartialLRP, trunc_prob_PartialLRP = predict(model, truncated_text_ids_PartialLRP, target, seg_ids=seg_ids)\n",
    "            rep_class_PartialLRP, rep_prob_PartialLRP = predict(model, replaced_text_ids_PartialLRP, target, seg_ids=seg_ids)\n",
    "            \n",
    "            instance_degradation_probs_PartialLRP.append(trunc_prob_PartialLRP)\n",
    "            instance_degradation_accs_PartialLRP.append(trunc_class_PartialLRP==target)\n",
    "            \n",
    "            instance_replace_probs_PartialLRP.append(rep_prob_PartialLRP)\n",
    "            instance_replace_accs_PartialLRP.append(rep_class_PartialLRP==target)\n",
    "            \n",
    "            # Rollout\n",
    "            truncated_text_ids_Rollout, _, _ = truncate_words(sorted_idx_Rollout, text_words, text_ids, \n",
    "                                                                                       num, seg_ids=seg_ids)\n",
    "            replaced_text_ids_Rollout, _ = replace_words(sorted_idx_Rollout, text_words, text_ids, num, mask, mask_id)\n",
    "            \n",
    "            trunc_class_Rollout, trunc_prob_Rollout = predict(model, truncated_text_ids_Rollout, target, seg_ids=seg_ids)\n",
    "            rep_class_Rollout, rep_prob_Rollout = predict(model, replaced_text_ids_Rollout, target, seg_ids=seg_ids)\n",
    "            \n",
    "            instance_degradation_probs_Rollout.append(trunc_prob_Rollout)\n",
    "            instance_degradation_accs_Rollout.append(trunc_class_Rollout==target)\n",
    "            \n",
    "            instance_replace_probs_Rollout.append(rep_prob_Rollout)\n",
    "            instance_replace_accs_Rollout.append(rep_class_Rollout==target)\n",
    "            \n",
    "            # LRP\n",
    "            truncated_text_ids_LRP, _, _ = truncate_words(sorted_idx_LRP, text_words, text_ids, \n",
    "                                                                                       num, seg_ids=seg_ids)\n",
    "            replaced_text_ids_LRP, _ = replace_words(sorted_idx_LRP, text_words, text_ids, num, mask, mask_id)\n",
    "                \n",
    "            trunc_class_LRP, trunc_prob_LRP = predict(model, truncated_text_ids_LRP, target, seg_ids=seg_ids)\n",
    "            rep_class_LRP, rep_prob_LRP = predict(model, replaced_text_ids_LRP, target, seg_ids=seg_ids)\n",
    "\n",
    "            instance_degradation_probs_LRP.append(trunc_prob_LRP)\n",
    "            instance_degradation_accs_LRP.append(trunc_class_LRP==target)\n",
    "            \n",
    "            instance_replace_probs_LRP.append(rep_prob_LRP)\n",
    "            instance_replace_accs_LRP.append(rep_class_LRP==target)\n",
    "\n",
    "        original_probs.append(original_prob)\n",
    "        original_accs.append(original_class==target)\n",
    "        \n",
    "        degradation_probs_TransCAM.append(instance_degradation_probs_TransCAM)\n",
    "        degradation_accs_TransCAM.append(instance_degradation_accs_TransCAM)\n",
    "        del_probs_TransCAM.append(instance_replace_probs_TransCAM)\n",
    "        del_accs_TransCAM.append(instance_replace_accs_TransCAM)\n",
    "        \n",
    "        \n",
    "        degradation_probs_Rawatt.append(instance_degradation_probs_Rawatt)\n",
    "        degradation_accs_Rawatt.append(instance_degradation_accs_Rawatt)\n",
    "        del_probs_Rawatt.append(instance_replace_probs_Rawatt)\n",
    "        del_accs_Rawatt.append(instance_replace_accs_Rawatt)\n",
    "        \n",
    "        degradation_probs_PartialLRP.append(instance_degradation_probs_PartialLRP)\n",
    "        degradation_accs_PartialLRP.append(instance_degradation_accs_PartialLRP)\n",
    "        del_probs_PartialLRP.append(instance_replace_probs_PartialLRP)\n",
    "        del_accs_PartialLRP.append(instance_replace_accs_PartialLRP)\n",
    "        \n",
    "        degradation_probs_Rollout.append(instance_degradation_probs_Rollout)\n",
    "        degradation_accs_Rollout.append(instance_degradation_accs_Rollout)\n",
    "        del_probs_Rollout.append(instance_replace_probs_Rollout)\n",
    "        del_accs_Rollout.append(instance_replace_accs_Rollout)\n",
    "        \n",
    "        degradation_probs_LRP.append(instance_degradation_probs_LRP)\n",
    "        degradation_accs_LRP.append(instance_degradation_accs_LRP)\n",
    "        del_probs_LRP.append(instance_replace_probs_LRP)\n",
    "        del_accs_LRP.append(instance_replace_accs_LRP)\n",
    "        \n",
    "        \n",
    "        kendaltau_TransCAM = cal_kendaltau(attribution_TransCAM,attribution_TransCAM_F)\n",
    "        kendaltaus_TransCAM.append(kendaltau_TransCAM)      \n",
    "        \n",
    "        kendaltau_Rawatt = cal_kendaltau(attribution_Rawatt,attribution_Rawatt_F)\n",
    "        kendaltaus_Rawatt.append(kendaltau_Rawatt)      \n",
    "        \n",
    "        kendaltau_PartialLRP = cal_kendaltau(attribution_PartialLRP,attribution_PartialLRP_F)\n",
    "        kendaltaus_PartialLRP.append(kendaltau_PartialLRP)      \n",
    "        \n",
    "        kendaltau_Rollout = cal_kendaltau(attribution_Rollout,attribution_Rollout_F)\n",
    "        kendaltaus_Rollout.append(kendaltau_Rollout)      \n",
    "        \n",
    "        kendaltau_LRP = cal_kendaltau(attribution_LRP,attribution_LRP_F)\n",
    "        kendaltaus_LRP.append(kendaltau_LRP)      \n",
    "        \n",
    "    # get aopc score\n",
    "    aopc_TransCAM = cal_aopc(original_probs,degradation_probs_TransCAM)\n",
    "    aopc_Rawatt = cal_aopc(original_probs,degradation_probs_Rawatt)\n",
    "    aopc_PartialLRP = cal_aopc(original_probs,degradation_probs_PartialLRP)\n",
    "    aopc_Rollout = cal_aopc(original_probs,degradation_probs_Rollout)\n",
    "    aopc_LRP = cal_aopc(original_probs,degradation_probs_LRP)\n",
    "    \n",
    "    aopc_TransCAM_del = cal_aopc(original_probs,del_probs_TransCAM)\n",
    "    aopc_Rawatt_del = cal_aopc(original_probs,del_probs_Rawatt)\n",
    "    aopc_PartialLRP_del = cal_aopc(original_probs,del_probs_PartialLRP)\n",
    "    aopc_Rollout_del = cal_aopc(original_probs,del_probs_Rollout)\n",
    "    aopc_LRP_del = cal_aopc(original_probs,del_probs_LRP)\n",
    "        \n",
    "        \n",
    "    logodds_TransCAM = cal_logodds(original_probs,degradation_probs_TransCAM)\n",
    "    logodds_Rawatt = cal_logodds(original_probs,degradation_probs_Rawatt)\n",
    "    logodds_PartialLRP = cal_logodds(original_probs,degradation_probs_PartialLRP)\n",
    "    logodds_Rollout = cal_logodds(original_probs,degradation_probs_Rollout)\n",
    "    logodds_LRP = cal_logodds(original_probs,degradation_probs_LRP)\n",
    "    \n",
    "    logodds_TransCAM_del = cal_logodds(original_probs,del_probs_TransCAM)\n",
    "    logodds_Rawatt_del = cal_logodds(original_probs,del_probs_Rawatt)\n",
    "    logodds_PartialLRP_del = cal_logodds(original_probs,del_probs_PartialLRP)\n",
    "    logodds_Rollout_del = cal_logodds(original_probs,del_probs_Rollout)\n",
    "    logodds_LRP_del = cal_logodds(original_probs,del_probs_LRP)\n",
    "        \n",
    "            \n",
    "    # get kendaltau score    \n",
    "    k_TransCAM = np.mean(kendaltaus_TransCAM)\n",
    "    k_Rawatt = np.mean(kendaltaus_Rawatt)\n",
    "    k_PartialLRP = np.mean(kendaltaus_PartialLRP)\n",
    "    k_Rollout = np.mean(kendaltaus_Rollout)\n",
    "    k_LRP = np.mean(kendaltaus_LRP)\n",
    "\n",
    "        \n",
    "    return (aopc_TransCAM, aopc_Rawatt, aopc_PartialLRP, aopc_Rollout, aopc_LRP,\n",
    "            aopc_TransCAM_del, aopc_Rawatt_del, aopc_PartialLRP_del, aopc_Rollout_del, aopc_LRP_del,\n",
    "            k_TransCAM, k_Rawatt, k_PartialLRP, k_Rollout, k_LRP,\n",
    "            logodds_TransCAM, logodds_Rawatt, logodds_PartialLRP, logodds_Rollout, logodds_LRP,\n",
    "            logodds_TransCAM_del, logodds_Rawatt_del, logodds_PartialLRP_del, logodds_Rollout_del, logodds_LRP_del)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5ede31c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize the explanations generator\n",
    "explanations = Generator(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "35fae478",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test_num = 100\n",
    "to_test = np.array(dataset['test'])\n",
    "# to_test_idx = np.random.choice(len(dataset['test']), test_num, replace=False)\n",
    "# to_test = to_test[to_test_idx]\n",
    "# print(len(to_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e950483c",
   "metadata": {},
   "outputs": [],
   "source": [
    "(aopc_TransCAM, aopc_Rawatt, aopc_PartialLRP, aopc_Rollout, aopc_LRP,\n",
    "aopc_TransCAM_del, aopc_Rawatt_del, aopc_PartialLRP_del, aopc_Rollout_del, aopc_LRP_del,\n",
    "k_TransCAM, k_Rawatt, k_PartialLRP, k_Rollout, k_LRP,\n",
    "logodds_TransCAM, logodds_Rawatt, logodds_PartialLRP, logodds_Rollout, logodds_LRP,\n",
    "logodds_TransCAM_del, logodds_Rawatt_del, logodds_PartialLRP_del, logodds_Rollout_del, logodds_LRP_del) = test(model, explanations, to_test, degrade_step = 10, seg_ids=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e34b492d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(aopc_TransCAM,aopc_Rawatt,aopc_PartialLRP,aopc_Rollout,aopc_LRP)\n",
    "# print(aopc_TransCAM_del,aopc_Rawatt_del,aopc_PartialLRP_del,aopc_Rollout_del,aopc_LRP_del)\n",
    "# print(k_TransCAM,k_Rawatt,k_PartialLRP,k_Rollout,k_LRP)\n",
    "# print(logodds_TransCAM, logodds_Rawatt, logodds_PartialLRP, logodds_Rollout, logodds_LRP)\n",
    "# print(logodds_TransCAM_del, logodds_Rawatt_del, logodds_PartialLRP_del, logodds_Rollout_del, logodds_LRP_del)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dcdf73f",
   "metadata": {},
   "source": [
    "# Single Example Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d704f19b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def expln_norm(expln):\n",
    "    expln = (expln - expln.min()) / (expln.max()- expln.min())\n",
    "    return expln\n",
    "\n",
    "def show_text_attr(expln,str_list,is_relu = True):\n",
    "    if is_relu:\n",
    "        rgb = lambda x: '0,0,0' if x < 0 else '0,255,0'\n",
    "        alpha = lambda x: max(x, 0)# ** 0.5\n",
    "    else:\n",
    "        rgb = lambda x: '255,0,0' if x < 0 else '0,255,0'\n",
    "        alpha = lambda x: x * -5 if x < 0 else x * 100\n",
    "    attrs = list(expln)\n",
    "    subwords = str_list\n",
    "    \n",
    "    token_marks = [\n",
    "        f'<mark style=\"background-color:rgba({rgb(attr)},{alpha(attr)})\">{token}</mark>'\n",
    "        for token, attr in zip(subwords, attrs)\n",
    "    ]\n",
    "    \n",
    "    display(HTML('<p>' + ' '.join(token_marks) + '</p>'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "8693c740",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_ids, att_mask, text_words = preprocess_sample(to_test[4]['content'])\n",
    "label = to_test[4]['label'] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "4c37cf47",
   "metadata": {},
   "outputs": [],
   "source": [
    "#TransCAM_expln, LRP_expln, PartialLRP_expln, FullLRP_expln, Att_expln, Rollout_expln, Att_Gradcam_expln = \\\n",
    "        #generate_explns(explanations, text_ids, att_mask, start_layer=0, true_class = label, is_true = True) \n",
    "    \n",
    "TransCAM_expln, LRP_expln, PartialLRP_expln, Att_expln, Rollout_expln = \\\n",
    "        generate_explns(explanations, text_ids, att_mask, start_layer=0, true_class = label, is_true = True) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ab1f539a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# TransCAM_expln = expln_norm(TransCAM_expln)\n",
    "# show_text_attr(TransCAM_expln,text_words,is_relu = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "38658983",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LRP_expln = expln_norm(LRP_expln)\n",
    "# show_text_attr(LRP_expln,text_words,is_relu = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4bba46ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# PartialLRP_expln = expln_norm(PartialLRP_expln)\n",
    "# show_text_attr(PartialLRP_expln,text_words,is_relu = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1bb8055e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Att_expln = expln_norm(Att_expln)\n",
    "# show_text_attr(Att_expln,text_words,is_relu = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9982ec23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rollout_expln = expln_norm(Rollout_expln)\n",
    "# show_text_attr(Rollout_expln,text_words,is_relu = False)"
   ]
  }
 ],
 "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}