distilbert_sst.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package punkt to /home/dehghani/nltk_data...\n",
      "[nltk_data]   Package punkt is already up-to-date!\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "from transformers import DistilBertTokenizer, TFDistilBertForSequenceClassification\n",
    "from transformers import *\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "from attention_graph_util import *\n",
    "from notebooks.notebook_utils import *\n",
    "from util import inflect\n",
    "\n",
    "from tqdm import tqdm\n",
    "import math\n",
    "import pandas as pd\n",
    "\n",
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"   # see issue #152\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')\n",
    "model = TFDistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased-finetuned-sst-2-english')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:absl:Overwrite dataset info from restored data version.\n",
      "INFO:absl:Reusing dataset glue (/home/dehghani/tensorflow_datasets/glue/sst2/1.0.0)\n",
      "INFO:absl:Constructing tf.data.Dataset for split None, from /home/dehghani/tensorflow_datasets/glue/sst2/1.0.0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 14/100 [===>..........................] - ETA: 15s - loss: 0.3824 - accuracy: 0.9106WARNING:tensorflow:Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches (in this case, 100 batches). You may need to use the repeat() function when building your dataset.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches (in this case, 100 batches). You may need to use the repeat() function when building your dataset.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 14/100 [===>..........................] - ETA: 16s - loss: 0.3605 - accuracy: 0.9106"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[0.05353218950331211, 0.9105505]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = tfds.load('glue/sst2')\n",
    "\n",
    "# Prepare dataset for GLUE as a tf.data.Dataset instance\n",
    "train_dataset = glue_convert_examples_to_features(data['train'], tokenizer, max_length=128, task='sst-2')\n",
    "valid_dataset = glue_convert_examples_to_features(data['validation'], tokenizer, max_length=128, task='sst-2')\n",
    "train_dataset = train_dataset.shuffle(100).batch(32).repeat(2)\n",
    "valid_dataset = valid_dataset.batch(64)\n",
    "test_dataset = glue_convert_examples_to_features(data['test'], tokenizer, max_length=128, task='sst-2')\n",
    "test_dataset = test_dataset.batch(1)\n",
    "\n",
    "# Prepare training: Compile tf.keras model with optimizer, loss and learning rate schedule\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0)\n",
    "loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')\n",
    "model.compile(optimizer=optimizer, loss=loss, metrics=[metric])\n",
    "\n",
    "model.evaluate(valid_dataset, steps=100)\n",
    "# Train and evaluate using tf.keras.Model.fit()\n",
    "# history = model.fit(train_dataset, epochs=2, steps_per_epoch=115,\n",
    "#                     validation_data=valid_dataset, validation_steps=7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([-3.6605854], shape=(1,), dtype=float32)\n",
      "tf.Tensor([-3.6501346], shape=(1,), dtype=float32)\n",
      "tf.Tensor([-3.5526536], shape=(1,), dtype=float32)\n",
      "tf.Tensor([4.650803], shape=(1,), dtype=float32)\n",
      "tf.Tensor([-3.1746943], shape=(1,), dtype=float32)\n",
      "tf.Tensor([-2.030062], shape=(1,), dtype=float32)\n",
      "tf.Tensor([4.5237985], shape=(1,), dtype=float32)\n",
      "tf.Tensor([-3.2844446], shape=(1,), dtype=float32)\n",
      "tf.Tensor([4.5983357], shape=(1,), dtype=float32)\n",
      "tf.Tensor([-3.4095068], shape=(1,), dtype=float32)\n",
      "tf.Tensor([2.0042334], shape=(1,), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "model.config.output_attentions = True\n",
    "model.config.output_hidden_states = True\n",
    "\n",
    "model.distilbert.transformer.output_attentions = True\n",
    "model.distilbert.transformer.output_hidden_states = True\n",
    "\n",
    "for layer in model.distilbert.transformer.layer:\n",
    "    layer.output_attentions = True    \n",
    "    layer.output_hidden_states = True\n",
    "    layer.attention.output_attentions = True\n",
    "    layer.attention.output_hidden_states = True\n",
    "  \n",
    "all_examples_grads = []\n",
    "all_examples_attentions = []\n",
    "all_examples_x = []\n",
    "n_examples = 10\n",
    "for x,y in test_dataset:\n",
    "    with tf.GradientTape() as tape:\n",
    "        inputs_embeds = model.distilbert.embeddings(x['input_ids'])\n",
    "        tape.watch(inputs_embeds)\n",
    "        logits, hidden_states, attentions = model({'attention_mask':x['attention_mask'], \n",
    "                                                    'inputs_embeds':inputs_embeds,\n",
    "                                                    'token_type_ids': x['token_type_ids']}, training=False\n",
    "                                            )\n",
    "        true_logits = logits[:,y[0]]\n",
    "        print(true_logits)\n",
    "    grads = tape.gradient(true_logits, inputs_embeds)[0]  \n",
    "    \n",
    "    length = tf.reduce_sum(x['attention_mask'], axis=-1)[0]    \n",
    "    all_examples_grads.append(tf.reduce_sum(grads, -1)[:length])\n",
    "    \n",
    "    _attentions = [att.numpy() for att in attentions]\n",
    "    attentions_mat = np.asarray(_attentions)[:,0]\n",
    "        \n",
    "    cropped_input = x['input_ids'][0, :length]\n",
    "    all_examples_x.append(cropped_input)\n",
    "    cropped_attention_mat = attentions_mat[:,:,:length,:length]\n",
    "    all_examples_attentions.append(cropped_attention_mat)\n",
    "        \n",
    "    if n_examples == 0:\n",
    "        break\n",
    "    n_examples -= 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 11/11 [00:00<00:00, 666.99it/s]\n",
      "100%|██████████| 11/11 [00:00<00:00, 731.42it/s]\n",
      "100%|██████████| 11/11 [00:00<00:00, 731.57it/s]\n",
      "100%|██████████| 11/11 [00:00<00:00, 745.09it/s]\n",
      "100%|██████████| 11/11 [00:00<00:00, 725.72it/s]\n",
      "100%|██████████| 11/11 [00:00<00:00, 716.94it/s]\n",
      "100%|██████████| 11/11 [00:00<00:00, 612.48it/s]\n",
      "100%|██████████| 11/11 [00:00<00:00, 622.65it/s]\n",
      "100%|██████████| 11/11 [00:00<00:00, 636.09it/s]\n",
      "100%|██████████| 11/11 [00:00<00:00, 604.55it/s]\n",
      "100%|██████████| 11/11 [00:00<00:00, 576.44it/s]\n",
      "  0%|          | 0/11 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "compute raw relevance scores ...\n",
      "compute joint relevance scores ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 11/11 [00:00<00:00, 594.15it/s]\n",
      "  0%|          | 0/11 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "compute flow relevance scores ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 11/11 [07:29<00:00, 40.89s/it]\n"
     ]
    }
   ],
   "source": [
    "print(\"compute raw relevance scores ...\")\n",
    "all_examples_raw_relevance = {}\n",
    "for l in np.arange(0,6):\n",
    "    all_examples_raw_relevance[l] = []\n",
    "    for i in tqdm(np.arange(len(all_examples_x))):\n",
    "        tokens = tokens = tokenizer.convert_ids_to_tokens(all_examples_x[i])\n",
    "        length = len(tokens)\n",
    "        attention_relevance = get_raw_att_relevance(all_examples_attentions[i], tokens, layer=l)\n",
    "        all_examples_raw_relevance[l].append(np.asarray(attention_relevance))\n",
    "\n",
    "print(\"compute joint relevance scores ...\")\n",
    "all_examples_joint_relevance = {}\n",
    "for l in np.arange(0,6):\n",
    "    all_examples_joint_relevance[l] = []\n",
    "    for i in tqdm(np.arange(len(all_examples_x))):\n",
    "        tokens = tokenizer.convert_ids_to_tokens(all_examples_x[i])\n",
    "        length = len(tokens)\n",
    "        attention_relevance = get_joint_relevance(all_examples_attentions[i], tokens, layer=l)\n",
    "        all_examples_joint_relevance[l].append(np.asarray(attention_relevance))\n",
    "    \n",
    "print(\"compute flow relevance scores ...\")\n",
    "all_examples_flow_relevance = {}\n",
    "for l in np.arange(5,6):\n",
    "    all_examples_flow_relevance[l] = []\n",
    "    for i in tqdm(np.arange(len(all_examples_x))):\n",
    "        tokens = tokenizer.convert_ids_to_tokens(all_examples_x[i])\n",
    "        length = len(tokens)            \n",
    "        attention_relevance = get_flow_relevance(all_examples_attentions[i], tokens, layer=l)\n",
    "        all_examples_flow_relevance[l].append(np.asarray(attention_relevance))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "###############Layer  0 #############\n",
      "raw grad\n",
      "(10,) (10,)\n",
      "0.11334062720565909 0.19622961505440625\n",
      "joint grad\n",
      "(10,) (10,)\n",
      "0.11334062720565909 0.19622961505440625\n",
      "flow grad\n",
      "###############Layer  1 #############\n",
      "raw grad\n",
      "(10,) (10,)\n",
      "0.13438242197187858 0.2093550383622229\n",
      "joint grad\n",
      "(10,) (10,)\n",
      "0.06895459052629656 0.2719340046989568\n",
      "flow grad\n",
      "###############Layer  2 #############\n",
      "raw grad\n",
      "(10,) (10,)\n",
      "0.10122507476342474 0.16864664813463295\n",
      "joint grad\n",
      "(10,) (10,)\n",
      "0.10001791316979691 0.2615892774898635\n",
      "flow grad\n",
      "###############Layer  3 #############\n",
      "raw grad\n",
      "(10,) (10,)\n",
      "0.03112317207684452 0.2250614444605738\n",
      "joint grad\n",
      "(10,) (10,)\n",
      "0.11786395932036378 0.2698312459821461\n",
      "flow grad\n",
      "###############Layer  4 #############\n",
      "raw grad\n",
      "(10,) (10,)\n",
      "0.02943448485732545 0.20305380256410632\n",
      "joint grad\n",
      "(10,) (10,)\n",
      "0.11878310925735795 0.26319519062457825\n",
      "flow grad\n",
      "###############Layer  5 #############\n",
      "raw grad\n",
      "(10,) (10,)\n",
      "0.08192566848896693 0.1834930264223889\n",
      "joint grad\n",
      "(10,) (10,)\n",
      "0.11808525406936811 0.25113535147660354\n",
      "flow grad\n",
      "(10,) (10,)\n",
      "0.11445567686664294 0.2042271147571669\n"
     ]
    }
   ],
   "source": [
    "for l in np.arange(0,6):\n",
    "    print(\"###############Layer \",l, \"#############\")\n",
    "\n",
    "    print('raw grad')\n",
    "    print(all_examples_raw_relevance[l][0].shape, all_examples_grads[0].shape)\n",
    "    raw_sps_grad = []\n",
    "    for i in np.arange(len(all_examples_x)):\n",
    "        sp = spearmanr(all_examples_raw_relevance[l][i],all_examples_grads[i])\n",
    "        raw_sps_grad.append(sp[0])\n",
    " \n",
    "        \n",
    "    print(np.mean(raw_sps_grad), np.std(raw_sps_grad))\n",
    "\n",
    "    \n",
    "    print('joint grad')\n",
    "    print(all_examples_joint_relevance[l][0].shape, all_examples_grads[0].shape)\n",
    "    joint_sps_grad = []\n",
    "    for i in np.arange(len(all_examples_x)):\n",
    "        sp = spearmanr(all_examples_joint_relevance[l][i],all_examples_grads[i])\n",
    "        joint_sps_grad.append(sp[0])\n",
    "\n",
    "        \n",
    "    print(np.mean(joint_sps_grad), np.std(joint_sps_grad))\n",
    "\n",
    "  \n",
    "    print('flow grad')\n",
    "for l in np.arange(5,6):\n",
    "    print(all_examples_joint_relevance[l][0].shape, all_examples_grads[0].shape)\n",
    "    flow_sps_grad = []\n",
    "    for i in np.arange(len(all_examples_x)):\n",
    "        sp = spearmanr(all_examples_flow_relevance[l][i],all_examples_grads[i])\n",
    "        flow_sps_grad.append(sp[0])\n",
    "        \n",
    "    print(np.mean(flow_sps_grad), np.std(flow_sps_grad))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.mean([spearmanr(all_examples_flow_relevance[l][i], all_examples_joint_relevance[l][i]) for i in np.arange(len(all_examples_grads))])\n",
    "np.mean([spearmanr(all_examples_flow_relevance[l][i], all_examples_grads[i]) for i in np.arange(len(all_examples_grads))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{5: [array([3.45120337, 0.38171321, 0.35614553, 0.64314935, 0.32497267,\n",
       "         0.50573498, 0.26461444, 0.29924043, 0.50490437, 3.26832165]),\n",
       "  array([7.03851756, 0.388758  , 0.46693231, 0.50656893, 0.39827641,\n",
       "         0.42258721, 0.57095143, 0.5175971 , 0.91077591, 0.53382325,\n",
       "         0.59075615, 0.59748871, 0.43935374, 0.41571566, 0.37244633,\n",
       "         0.4351109 , 0.44730335, 0.36060066, 0.43673612, 0.51734508,\n",
       "         0.37707071, 1.01195414, 0.41534465, 0.37809714, 0.32798143,\n",
       "         0.34877272, 1.65494549, 7.11818891]),\n",
       "  array([5.35423478, 0.39700633, 0.3638119 , 0.29394217, 0.50870478,\n",
       "         0.38152168, 0.36061633, 0.39126104, 0.44283174, 0.44666919,\n",
       "         0.57516408, 0.42065887, 0.52866337, 0.46851681, 0.47301116,\n",
       "         0.65925707, 0.35952456, 0.41882319, 1.65683918, 5.49894176]),\n",
       "  array([5.04530681, 0.36488562, 0.36786416, 0.39413226, 0.36986755,\n",
       "         0.65043878, 0.3624081 , 0.6024357 , 0.29888739, 0.33457911,\n",
       "         0.4688398 , 0.42846398, 0.41640085, 0.30840039, 0.29583143,\n",
       "         0.41145822, 1.73671532, 5.14308452]),\n",
       "  array([7.93927715, 0.71707407, 0.55912663, 0.69661904, 0.54852698,\n",
       "         0.47366068, 0.48465751, 0.34445064, 0.47464883, 0.33593645,\n",
       "         0.51946857, 0.464743  , 0.46124487, 0.42985705, 0.35072307,\n",
       "         0.54984454, 0.34670596, 0.46389571, 0.31539356, 0.41831647,\n",
       "         0.38536297, 0.76472136, 0.57486073, 0.51557016, 0.78161147,\n",
       "         0.75721984, 0.54917004, 0.87964322, 0.43120062, 0.40572649,\n",
       "         0.54238677, 0.53916153, 0.38657017, 0.52076028, 2.11519445,\n",
       "         8.95666911]),\n",
       "  array([6.73951753, 0.40552154, 0.44995539, 0.60573714, 0.3943868 ,\n",
       "         0.58799708, 0.61202964, 0.53801207, 0.43536331, 0.52699661,\n",
       "         0.42583338, 0.55088881, 0.49205596, 0.64598235, 0.61653948,\n",
       "         0.54639475, 0.43502779, 0.77978579, 0.35408509, 0.51239276,\n",
       "         0.51140659, 0.43198992, 0.47017167, 0.41608801, 0.34495327,\n",
       "         1.06595688, 0.45973611, 0.53935415, 1.75766464, 7.34817549]),\n",
       "  array([5.25601543, 0.38631331, 0.4996868 , 0.6597432 , 0.38510249,\n",
       "         0.47786637, 0.55633058, 0.42738891, 0.28936914, 0.46370399,\n",
       "         0.36451777, 0.75197385, 0.53779786, 0.52480288, 0.37870706,\n",
       "         0.44588949, 0.75645365, 0.34258937, 0.33258024, 1.3146573 ,\n",
       "         5.84851029]),\n",
       "  array([3.95033654, 0.36103394, 0.360615  , 0.37232464, 0.42318143,\n",
       "         0.56605338, 0.42702302, 0.42011093, 0.44010926, 0.30426292,\n",
       "         0.33838114, 1.26942951, 3.76713829]),\n",
       "  array([2.56218301, 0.29280962, 0.24395742, 0.22546701, 0.32175789,\n",
       "         0.94301634, 2.4108087 ]),\n",
       "  array([5.28307581, 0.4689703 , 0.53951735, 0.45982038, 0.37960949,\n",
       "         0.30815811, 0.45808326, 0.57066118, 0.56973379, 0.48054546,\n",
       "         0.87151148, 0.36546688, 0.57195295, 0.45902365, 0.52140576,\n",
       "         0.47392913, 0.40279393, 0.30780658, 0.56527162, 1.38408735,\n",
       "         5.55857554]),\n",
       "  array([6.2870986 , 0.54938468, 0.38703129, 0.38752758, 0.48246849,\n",
       "         0.37719979, 0.47733407, 0.41758239, 0.8536703 , 0.49971211,\n",
       "         0.49836299, 0.46741681, 0.67659933, 0.40751686, 0.41214355,\n",
       "         0.76351658, 0.41735714, 0.4447447 , 0.59143934, 0.45274485,\n",
       "         0.460276  , 0.47008077, 0.76522686, 0.46593257, 0.44556672,\n",
       "         1.86599826, 6.67606733])]}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_examples_joint_relevance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def spearmanr(x, y):\n",
    "    \"\"\" `x`, `y` --> pd.Series\"\"\"\n",
    "    x = pd.Series(x)\n",
    "    y = pd.Series(y)\n",
    "    assert x.shape == y.shape\n",
    "    rx = x.rank(method='dense')\n",
    "    ry = y.rank(method='dense')\n",
    "    d = rx - ry\n",
    "    dsq = np.sum(np.square(d))\n",
    "    n = x.shape[0]\n",
    "    coef = 1. - (6. * dsq) / (n * (n**2 - 1.))\n",
    "    return [coef]\n",
    "\n",
    "def get_raw_att_relevance(full_att_mat, input_tokens, layer=-1):\n",
    "    cls_index = 0\n",
    "    return full_att_mat[layer].sum(axis=0).sum(axis=0)\n",
    "    \n",
    "\n",
    "def compute_node_flow(G, labels_to_index, input_nodes, output_nodes,length):\n",
    "    number_of_nodes = len(labels_to_index)\n",
    "    flow_values=np.zeros((number_of_nodes,number_of_nodes))\n",
    "    for key in output_nodes:\n",
    "        if key not in input_nodes:\n",
    "            current_layer = int(labels_to_index[key] / length)\n",
    "            pre_layer = current_layer - 1\n",
    "            u = labels_to_index[key]\n",
    "            for inp_node_key in input_nodes:\n",
    "                v = labels_to_index[inp_node_key]\n",
    "                flow_value = nx.maximum_flow_value(G,u,v)\n",
    "                flow_values[u][pre_layer*length+v] = flow_value\n",
    "            #normalize flow values\n",
    "            flow_values[u] /= flow_values[u].sum()\n",
    "            \n",
    "    return flow_values\n",
    "\n",
    "def get_flow_relevance(full_att_mat, input_tokens, layer):\n",
    "    \n",
    "    res_att_mat = full_att_mat.sum(axis=1)/full_att_mat.shape[1]\n",
    "    res_att_mat = res_att_mat + np.eye(res_att_mat.shape[1])[None,...]\n",
    "    res_att_mat = res_att_mat / res_att_mat.sum(axis=-1)[...,None]\n",
    "\n",
    "    res_adj_mat, res_labels_to_index = get_adjmat(mat=res_att_mat, input_tokens=input_tokens)\n",
    "    \n",
    "    A = res_adj_mat\n",
    "    res_G=nx.from_numpy_matrix(A, create_using=nx.DiGraph())\n",
    "    for i in np.arange(A.shape[0]):\n",
    "        for j in np.arange(A.shape[1]):\n",
    "            nx.set_edge_attributes(res_G, {(i,j): A[i,j]}, 'capacity')\n",
    "\n",
    "    output_nodes = []\n",
    "    input_nodes = []\n",
    "    for key in res_labels_to_index:\n",
    "        if key.startswith('L'+str(layer+1)+'_'):\n",
    "            output_nodes.append(key)\n",
    "        if res_labels_to_index[key] < full_att_mat.shape[-1]:\n",
    "            input_nodes.append(key)\n",
    "    \n",
    "    flow_values = compute_node_flow(res_G, res_labels_to_index, input_nodes, output_nodes, length=full_att_mat.shape[-1])\n",
    "    \n",
    "    n_layers = full_att_mat.shape[0]\n",
    "    length = full_att_mat.shape[-1]\n",
    "    final_layer_attention_raw = flow_values[(layer+1)*length: (layer+2)*length,layer*length: (layer+1)*length]\n",
    "    relevance_attention_raw = final_layer_attention_raw.sum(axis=0)\n",
    "\n",
    "    return relevance_attention_raw\n",
    "    \n",
    "    \n",
    "def get_joint_relevance(full_att_mat, input_tokens, layer):\n",
    "    att_sum_heads =  full_att_mat.sum(axis=1)/full_att_mat.shape[1]\n",
    "    joint_attentions = compute_joint_attention(att_sum_heads, add_residual=True)\n",
    "    relevance_attentions = joint_attentions[layer].sum(axis=0)\n",
    "    return relevance_attentions\n",
    "\n",
    "\n",
    "def get_nores_joint_relevance(full_att_mat, input_tokens, layer):\n",
    "    att_sum_heads =  full_att_mat.sum(axis=1)/full_att_mat.shape[1]\n",
    "    joint_attentions = compute_joint_attention(att_sum_heads, add_residual=False)\n",
    "    relevance_attentions = joint_attentions[layer].sum(axis=0)\n",
    "    return relevance_attentions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}