Head_selection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Head Selection\n",
    "**_BERT_** is a **_Multi-layer_ _Multi-Head_** Transformer architecture. As discuss in many of the current reseachers, different Attention heads captures different lingustic patterns. For a better deletion of words using Attention mechanism we need to choose a head which **captures pattern useful for classification.**\n",
    "\n",
    "To do this we are using a Brute force mechanism to seach through all the possible heads. We are deleting TopK words attended by different heads from the sentence and measuring the new classification score. In case of sentiments, removing sentiments related words makes the sentence neutral. The heads are sorted by the amount to which it is able to make the sentences from dev set to Neutral."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "import logging\n",
    "import os\n",
    "import random\n",
    "import sys\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,\n",
    "                              TensorDataset)\n",
    "from torch.utils.data.distributed import DistributedSampler\n",
    "from tqdm import tqdm, trange\n",
    "\n",
    "from pytorch_pretrained_bert.file_utils import PYTORCH_PRETRAINED_BERT_CACHE\n",
    "from pytorch_pretrained_bert.modeling import BertForSequenceClassification, BertConfig, WEIGHTS_NAME, CONFIG_NAME\n",
    "#from pytorch_pretrained_bert.tokenization import BertTokenizer\n",
    "from pytorch_pretrained_bert.optimization import BertAdam, warmup_linear\n",
    "\n",
    "from bertviz.bertviz import attention, visualization\n",
    "from bertviz.bertviz.pytorch_pretrained_bert import BertModel, BertTokenizer\n",
    "\n",
    "logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',\n",
    "                    datefmt = '%m/%d/%Y %H:%M:%S',\n",
    "                    level = logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "07/01/2022 21:28:58 - INFO - __main__ -   device: cuda, n_gpu 1\n"
     ]
    }
   ],
   "source": [
    "logger = logging.getLogger(__name__)\n",
    "bert_classifier_model_dir = \"../models/BERT/\" ## Path of BERT classifier model path\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "n_gpu = torch.cuda.device_count()\n",
    "logger.info(\"device: {}, n_gpu {}\".format(device, n_gpu))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "07/01/2022 21:28:58 - INFO - pytorch_pretrained_bert.modeling -   loading archive file ../models/BERT/\n",
      "07/01/2022 21:28:58 - INFO - pytorch_pretrained_bert.modeling -   Model config {\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"max_position_embeddings\": 512,\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_hidden_layers\": 12,\n",
      "  \"type_vocab_size\": 2,\n",
      "  \"vocab_size\": 30522\n",
      "}\n",
      "\n",
      "07/01/2022 21:29:00 - INFO - bertviz.bertviz.pytorch_pretrained_bert.tokenization -   loading vocabulary file https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-vocab.txt from cache at /home/jack/.pytorch_pretrained_bert/26bc1ad6c0ac742e9b52263248f6d0f00068293b33709fae12320c0e35ccfbbb.542ce4285a40d23a559526243235df47c5f75c197f04f37d1a0c124c32c9a084\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BertForSequenceClassification(\n",
       "  (bert): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(30522, 768)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): BertLayerNorm()\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (1): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (2): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (3): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (4): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (5): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (6): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (7): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (8): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (9): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (10): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (11): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (pooler): BertPooler(\n",
       "      (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "      (activation): Tanh()\n",
       "    )\n",
       "  )\n",
       "  (dropout): Dropout(p=0.1, inplace=False)\n",
       "  (classifier): Linear(in_features=768, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## Model for performing Classification\n",
    "model_cls = BertForSequenceClassification.from_pretrained(bert_classifier_model_dir, num_labels=2)\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)\n",
    "model_cls.to(device)\n",
    "model_cls.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "07/01/2022 21:29:01 - INFO - bertviz.bertviz.pytorch_pretrained_bert.modeling -   loading archive file ../models/BERT/\n",
      "07/01/2022 21:29:01 - INFO - bertviz.bertviz.pytorch_pretrained_bert.modeling -   Model config {\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"max_position_embeddings\": 512,\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_hidden_layers\": 12,\n",
      "  \"type_vocab_size\": 2,\n",
      "  \"vocab_size\": 30522\n",
      "}\n",
      "\n",
      "07/01/2022 21:29:02 - INFO - bertviz.bertviz.pytorch_pretrained_bert.tokenization -   loading vocabulary file https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-vocab.txt from cache at /home/jack/.pytorch_pretrained_bert/26bc1ad6c0ac742e9b52263248f6d0f00068293b33709fae12320c0e35ccfbbb.542ce4285a40d23a559526243235df47c5f75c197f04f37d1a0c124c32c9a084\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BertModel(\n",
       "  (embeddings): BertEmbeddings(\n",
       "    (word_embeddings): Embedding(30522, 768)\n",
       "    (position_embeddings): Embedding(512, 768)\n",
       "    (token_type_embeddings): Embedding(2, 768)\n",
       "    (LayerNorm): BertLayerNorm()\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (encoder): BertEncoder(\n",
       "    (layer): ModuleList(\n",
       "      (0): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): BertLayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (1): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): BertLayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (2): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): BertLayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (3): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): BertLayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (4): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): BertLayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (5): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): BertLayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (6): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): BertLayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (7): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): BertLayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (8): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): BertLayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (9): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): BertLayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (10): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): BertLayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (11): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): BertLayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (pooler): BertPooler(\n",
       "    (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "    (activation): Tanh()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## Model to get the attention weights of all the heads\n",
    "model = BertModel.from_pretrained(bert_classifier_model_dir)\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)\n",
    "model.to(device)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_seq_len=70 # Maximum sequence length \n",
    "sm = torch.nn.Softmax(dim=-1) ## Softmax over the batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_multiple_examples(input_sentences, bs=32):\n",
    "    \"\"\"\n",
    "    This fucntion returns classification predictions for batch of sentences.\n",
    "    input_sentences: list of strings\n",
    "    bs : batch_size : int\n",
    "    \"\"\"\n",
    "    \n",
    "    ## Prepare data for classification\n",
    "    ids = []\n",
    "    segment_ids = []\n",
    "    input_masks = []\n",
    "    pred_lt = []\n",
    "    for sen in input_sentences:\n",
    "        text_tokens = tokenizer.tokenize(sen)\n",
    "        tokens = [\"[CLS]\"] + text_tokens + [\"[SEP]\"]\n",
    "        temp_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "        input_mask = [1] * len(temp_ids)\n",
    "        segment_id = [0] * len(temp_ids)\n",
    "        padding = [0] * (max_seq_len - len(temp_ids))\n",
    "\n",
    "        temp_ids += padding\n",
    "        input_mask += padding\n",
    "        segment_id += padding\n",
    "        \n",
    "        ids.append(temp_ids)\n",
    "        input_masks.append(input_mask)\n",
    "        segment_ids.append(segment_id)\n",
    "    \n",
    "    ## Convert input lists to Torch Tensors\n",
    "    ids = torch.tensor(ids)\n",
    "    segment_ids = torch.tensor(segment_ids)\n",
    "    input_masks = torch.tensor(input_masks)\n",
    "    \n",
    "    steps = len(ids) // bs\n",
    "    \n",
    "    for i in range(steps+1):\n",
    "        if i == steps:\n",
    "            temp_ids = ids[i * bs : len(ids)]\n",
    "            temp_segment_ids = segment_ids[i * bs: len(ids)]\n",
    "            temp_input_masks = input_masks[i * bs: len(ids)]\n",
    "        else:\n",
    "            temp_ids = ids[i * bs : i * bs + bs]\n",
    "            temp_segment_ids = segment_ids[i * bs: i * bs + bs]\n",
    "            temp_input_masks = input_masks[i * bs: i * bs + bs]\n",
    "        \n",
    "        temp_ids = temp_ids.to(device)\n",
    "        temp_segment_ids = temp_segment_ids.to(device)\n",
    "        temp_input_masks = temp_input_masks.to(device)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            preds = sm(model_cls(temp_ids, temp_segment_ids, temp_input_masks))\n",
    "        pred_lt.extend(preds.tolist())\n",
    "    \n",
    "    return pred_lt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_file(path,size):\n",
    "    with open(path) as fp:\n",
    "        data = fp.read().splitlines()[:size]\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_attention_for_batch(input_sentences, bs=32):\n",
    "    \"\"\"\n",
    "    This function calculates attention weights of all the heads and\n",
    "    returns it along with the encoded sentence for further processing.\n",
    "    \n",
    "    input sentence: list of strings\n",
    "    bs : batch_size\n",
    "    \"\"\"\n",
    "    \n",
    "    ## Preprocessing for BERT \n",
    "    ids = []\n",
    "    segment_ids = []\n",
    "    input_masks = []\n",
    "    pred_lt = []\n",
    "    ids_for_decoding = []\n",
    "    for sen in input_sentences:\n",
    "        text_tokens = tokenizer.tokenize(sen)\n",
    "        tokens = [\"[CLS]\"] + text_tokens + [\"[SEP]\"]\n",
    "        temp_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "        input_mask = [1] * len(temp_ids)\n",
    "        segment_id = [0] * len(temp_ids)\n",
    "        padding = [0] * (max_seq_len - len(temp_ids))\n",
    "        \n",
    "        ids_for_decoding.append(tokenizer.convert_tokens_to_ids(tokens))\n",
    "        temp_ids += padding\n",
    "        input_mask += padding\n",
    "        segment_id += padding\n",
    "        \n",
    "        ids.append(temp_ids)\n",
    "        input_masks.append(input_mask)\n",
    "        segment_ids.append(segment_id)\n",
    "    ## Convert the list of int ids to Torch Tensors\n",
    "    ids = torch.tensor(ids)\n",
    "    segment_ids = torch.tensor(segment_ids)\n",
    "    input_masks = torch.tensor(input_masks)\n",
    "    \n",
    "    steps = len(ids) // bs\n",
    "    \n",
    "    for i in trange(steps+1):\n",
    "        if i == steps:\n",
    "            temp_ids = ids[i * bs : len(ids)]\n",
    "            temp_segment_ids = segment_ids[i * bs: len(ids)]\n",
    "            temp_input_masks = input_masks[i * bs: len(ids)]\n",
    "        else:\n",
    "            temp_ids = ids[i * bs : i * bs + bs]\n",
    "            temp_segment_ids = segment_ids[i * bs: i * bs + bs]\n",
    "            temp_input_masks = input_masks[i * bs: i * bs + bs]\n",
    "        \n",
    "        temp_ids = temp_ids.to(device)\n",
    "        temp_segment_ids = temp_segment_ids.to(device)\n",
    "        temp_input_masks = temp_input_masks.to(device)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            _, _, attn = model(temp_ids, temp_segment_ids, temp_input_masks)\n",
    "        \n",
    "        # Add all the Attention Weights to CPU memory\n",
    "        # Attention weights for each layer is stored in a dict 'attn_prob'\n",
    "        for k in range(12):\n",
    "            attn[k]['attn_probs'] = attn[k]['attn_probs'].to('cpu')\n",
    "        \n",
    "        '''\n",
    "        attention weights are stored in this way:\n",
    "        att_lt[layer]['attn_probs']['input_sentence']['head']['length_of_sentence']\n",
    "        '''\n",
    "        # Concate Attention weights for all the examples in the list att_lt[layer_no]['attn_probs']\n",
    "        \n",
    "        if i == 0:\n",
    "            att_lt = attn\n",
    "            heads = len(att_lt)\n",
    "        else:\n",
    "            for j in range(heads):\n",
    "                att_lt[j]['attn_probs'] = torch.cat((att_lt[j]['attn_probs'],attn[j]['attn_probs']),0)\n",
    "        \n",
    "    \n",
    "    return att_lt, ids_for_decoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_sentences(input_sentences, att, decoding_ids, threshold=0.25):\n",
    "    \"\"\"\n",
    "    This function processes each input sentence by removing the top tokens defined threshold value.\n",
    "    Each sentence is processed for each head.\n",
    "    \n",
    "    input_ids: list of strings\n",
    "    decoding_ids: indexed input_sentnces thus len(input_sentences) == len(decoding_ids)\n",
    "    threshold: Percentage of the top indexes to be removed\n",
    "    \"\"\"\n",
    "    # List of None of num_of_layers * num_of_heads to save the results of each head for input_sentences\n",
    "    \n",
    "    lt = [None for x in range(len(att) * len(att[0]['attn_probs'][0]))]\n",
    "    #print(len(lt))\n",
    "    \n",
    "    inx = 0\n",
    "    for i in trange(len(att)): #  For all the layers\n",
    "        for j in range(len(att[i]['attn_probs'][0])): # For all the heads in the ith Layer\n",
    "            processed_sen = [None for q in decoding_ids] # List of len(decoding_ids)\n",
    "            for k in range(len(input_sentences)): # For all the senteces \n",
    "                _, topi = att[i]['attn_probs'][k][j][0].topk(len(decoding_ids[k])) # Get top attended ids\n",
    "                topi = topi.tolist()\n",
    "                topi = topi[:int(len(topi) * threshold)] \n",
    "                ## Decode the sentece after removing the topk indexes\n",
    "                final_indexes = []\n",
    "                count = 0\n",
    "                count1 = 0\n",
    "                tokens = [\"[CLS]\"] + tokenizer.tokenize(input_sentences[k]) + [\"[SEP]\"]\n",
    "                while count < len(decoding_ids[k]):\n",
    "                    if count in topi: # Remove index if present in topk\n",
    "                        while (count + count1 + 1) < len(decoding_ids[k]):\n",
    "                            if \"##\" in tokens[count + count1 + 1]:\n",
    "                                count1 += 1\n",
    "                            else:\n",
    "                                break\n",
    "                        count += count1\n",
    "                        count1 = 0\n",
    "                    else: # Else add to the decoded sentence\n",
    "                        final_indexes.append(decoding_ids[k][count])\n",
    "                    count += 1\n",
    "                tmp = tokenizer.convert_ids_to_tokens(final_indexes) # Convert ids to token\n",
    "                # Convert toknes to sentence\n",
    "                processed_sen[k] = \" \".join(tmp).replace(\" ##\", \"\").replace(\"[CLS]\",\"\").replace(\"[SEP]\",\"\").strip()\n",
    "            lt[inx] = processed_sen # Store sentences for inxth head\n",
    "            inx += 1\n",
    "    \n",
    "    return lt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_block_head(processed_sentence_list, lmbd = 0.1):\n",
    "    \"\"\"\n",
    "    This function calculate classification scores for sentences generated by each head\n",
    "    and sort them from best to worst.\n",
    "    score = min(pred) + lmbd / max(pred) + lmbd, lmbd is smoothing param\n",
    "    pred is list of probability score for each class, for best case pred = [0.5, 0.5] ==> score = 1\n",
    "    \n",
    "    it returns sorted list of (Layer, Head, Score)\n",
    "    \"\"\"\n",
    "    scores = {}\n",
    "    #scores_1 = {}\n",
    "    for i in trange(len(processed_sentence_list)): # sentences by each head\n",
    "        pred = np.array(run_multiple_examples(processed_sentence_list[i]))\n",
    "        scores[i] = np.mean([(min(x[0], x[1])+lmbd)/(max(x[0], x[1])+lmbd) for x in pred])\n",
    "        #scores_1[i] = np.mean([abs(max(x[0],x[1]) - min(x[0],x[1])) for x in pred])\n",
    "    temp = sorted(scores.items(), key=lambda kv: kv[1], reverse=True)\n",
    "    #temp1 = sorted(scores_1.items(), key=lambda kv: kv[1], reverse=False)\n",
    "    score_lt = [(x // 12, x - (12 * (x // 12)),y) for x,y in temp]\n",
    "    #score1_lt = [(x // 12, x - (12 * (x // 12)),y) for x,y in temp1]\n",
    "    return score_lt  #score1_lt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos_examples_file = \"/home/jack/Desktop/NN/clean/datasets/yelp/sentiment.dev.1\"\n",
    "neg_examples_file = \"/home/jack/Desktop/NN/clean/datasets/yelp/sentiment.dev.0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "100 examples from each class worked good, the bottlenack is the run_multiple_examples() function,\n",
    "with higher memory (either with cpu of gpu) one can reduce the processing time by incresing batch_size.\n",
    "With batch_size of 32 it takes around 24 mins for 100 example on cpu.\n",
    "'''\n",
    "pos_data = read_file(pos_examples_file,100)\n",
    "neg_data = read_file(neg_examples_file,100)\n",
    "data = pos_data + neg_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100 100 200\n"
     ]
    }
   ],
   "source": [
    "print(len(pos_data), len(neg_data), len(data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00,  7.71it/s]\n"
     ]
    }
   ],
   "source": [
    "att, decoding_ids = get_attention_for_batch(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 12/12 [00:02<00:00,  4.84it/s]\n"
     ]
    }
   ],
   "source": [
    "sen_list = process_sentences(data, att, decoding_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 144/144 [00:36<00:00,  3.89it/s]\n"
     ]
    }
   ],
   "source": [
    "scores = get_block_head(sen_list)\n",
    "#(Layer, Head, Score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(4, 2, 0.2756469215333084),\n",
       " (8, 4, 0.2730110615638815),\n",
       " (8, 8, 0.26503659479341135),\n",
       " (8, 7, 0.2580905416106284),\n",
       " (6, 0, 0.25572399244530375),\n",
       " (10, 0, 0.2516887464487021),\n",
       " (9, 2, 0.24970683977038988),\n",
       " (9, 0, 0.24708689472474177),\n",
       " (9, 5, 0.24283926598519417),\n",
       " (8, 9, 0.24082947585836592),\n",
       " (9, 7, 0.22945535364116207),\n",
       " (9, 8, 0.22486696327745462),\n",
       " (8, 1, 0.22455436359670308),\n",
       " (5, 4, 0.2182670475177032),\n",
       " (10, 1, 0.21443260326638997),\n",
       " (11, 3, 0.21268681342838108),\n",
       " (11, 5, 0.21230563698846822),\n",
       " (10, 4, 0.2107462382280202),\n",
       " (11, 11, 0.20866078782370057),\n",
       " (6, 10, 0.20679038950079726),\n",
       " (10, 10, 0.20516895511245303),\n",
       " (8, 0, 0.20359565654894207),\n",
       " (7, 0, 0.2035650194900011),\n",
       " (7, 5, 0.20354271620052608),\n",
       " (7, 4, 0.20350667180131837),\n",
       " (8, 11, 0.20287416935696242),\n",
       " (6, 5, 0.20194093099473318),\n",
       " (9, 10, 0.199527478723871),\n",
       " (5, 7, 0.19399639300568355),\n",
       " (7, 9, 0.1939808844555725),\n",
       " (10, 11, 0.190830167371757),\n",
       " (4, 9, 0.1881700012574875),\n",
       " (7, 1, 0.18792038601234803),\n",
       " (4, 6, 0.1861797650228878),\n",
       " (3, 2, 0.18608476132807175),\n",
       " (6, 4, 0.18298562759313378),\n",
       " (11, 8, 0.1826706109354017),\n",
       " (8, 10, 0.1819432196957768),\n",
       " (8, 3, 0.18113500319960898),\n",
       " (8, 5, 0.1809133016276752),\n",
       " (6, 2, 0.1789835936178072),\n",
       " (9, 1, 0.17437386312161662),\n",
       " (0, 9, 0.17175563553478176),\n",
       " (4, 1, 0.1716945683966091),\n",
       " (8, 6, 0.17082037525886684),\n",
       " (7, 6, 0.17020695761301866),\n",
       " (9, 3, 0.16949805354725764),\n",
       " (5, 8, 0.16884294976655553),\n",
       " (6, 7, 0.16795098669555938),\n",
       " (6, 1, 0.1675820877129961),\n",
       " (5, 3, 0.16740198483746332),\n",
       " (6, 8, 0.16721768003755316),\n",
       " (4, 7, 0.16582199182914031),\n",
       " (4, 4, 0.16329861023993586),\n",
       " (5, 0, 0.16220945672733844),\n",
       " (4, 11, 0.16030610389602168),\n",
       " (8, 2, 0.15538780481448938),\n",
       " (4, 0, 0.15520076178923922),\n",
       " (5, 10, 0.153964055909925),\n",
       " (5, 2, 0.1533835036528971),\n",
       " (4, 10, 0.15226037968646788),\n",
       " (9, 9, 0.14848619241305813),\n",
       " (11, 4, 0.14842676130403326),\n",
       " (5, 5, 0.1460130568045256),\n",
       " (2, 5, 0.14595487680385807),\n",
       " (7, 11, 0.1448076645043975),\n",
       " (3, 8, 0.144134388201107),\n",
       " (9, 11, 0.14355609442974532),\n",
       " (7, 8, 0.14323358258336788),\n",
       " (6, 9, 0.14300252226995086),\n",
       " (3, 5, 0.14228475674072494),\n",
       " (7, 7, 0.14220034039717347),\n",
       " (7, 10, 0.1417832148725962),\n",
       " (4, 3, 0.14172131318156547),\n",
       " (11, 9, 0.14085346436331897),\n",
       " (6, 11, 0.13967998229589676),\n",
       " (9, 4, 0.13821565796659938),\n",
       " (4, 5, 0.13771306806503905),\n",
       " (5, 11, 0.1370568755981653),\n",
       " (10, 9, 0.13545649869622087),\n",
       " (5, 9, 0.13532165130153082),\n",
       " (7, 2, 0.13499977711292663),\n",
       " (6, 6, 0.134703916732059),\n",
       " (6, 3, 0.13336750608965023),\n",
       " (5, 1, 0.1332873588795881),\n",
       " (1, 1, 0.1331049258093971),\n",
       " (10, 5, 0.133095962208896),\n",
       " (10, 2, 0.13272208469647992),\n",
       " (7, 3, 0.1324545213051558),\n",
       " (9, 6, 0.13208760581084442),\n",
       " (2, 3, 0.1311315918149723),\n",
       " (3, 6, 0.1302094966453778),\n",
       " (0, 7, 0.1298280990452591),\n",
       " (3, 10, 0.12952286895560738),\n",
       " (3, 1, 0.12940201514454736),\n",
       " (10, 6, 0.12904942016823792),\n",
       " (11, 2, 0.12886996316353186),\n",
       " (1, 10, 0.12870273891559564),\n",
       " (2, 11, 0.1278857755361972),\n",
       " (1, 9, 0.12759637037406374),\n",
       " (2, 7, 0.12748650166308845),\n",
       " (0, 4, 0.12735220876706638),\n",
       " (1, 11, 0.1268229261940193),\n",
       " (11, 0, 0.1265662633921614),\n",
       " (0, 8, 0.12625533310747547),\n",
       " (11, 7, 0.1262047103001089),\n",
       " (3, 7, 0.12611685442599435),\n",
       " (1, 4, 0.12573918663595307),\n",
       " (0, 11, 0.12541034862948952),\n",
       " (4, 8, 0.125239776849847),\n",
       " (1, 5, 0.12508178917865023),\n",
       " (11, 6, 0.12496996652018037),\n",
       " (0, 5, 0.1245927901038978),\n",
       " (3, 3, 0.12448210788250919),\n",
       " (0, 3, 0.12429224875736442),\n",
       " (0, 6, 0.12405193719452651),\n",
       " (2, 2, 0.1239752700384091),\n",
       " (0, 10, 0.12392765415055607),\n",
       " (0, 2, 0.12369223572287172),\n",
       " (2, 4, 0.1231232815467306),\n",
       " (2, 8, 0.12281130789320209),\n",
       " (1, 3, 0.12253041239980682),\n",
       " (10, 3, 0.1224780532404444),\n",
       " (1, 2, 0.12220900354646876),\n",
       " (2, 6, 0.12186335514388802),\n",
       " (1, 6, 0.12140319606172849),\n",
       " (3, 11, 0.12102828802613776),\n",
       " (2, 0, 0.12092575410270077),\n",
       " (1, 0, 0.12082219476467579),\n",
       " (2, 1, 0.12069109998098757),\n",
       " (3, 4, 0.12033974903264447),\n",
       " (3, 9, 0.12012907355489456),\n",
       " (11, 1, 0.11975716057926389),\n",
       " (10, 7, 0.11909246481199168),\n",
       " (3, 0, 0.11891762229211236),\n",
       " (2, 9, 0.11890569217660021),\n",
       " (0, 0, 0.11888871944842112),\n",
       " (1, 8, 0.1180775155345885),\n",
       " (1, 7, 0.11778166245106096),\n",
       " (2, 10, 0.11723498123857495),\n",
       " (10, 8, 0.11631687444284161),\n",
       " (0, 1, 0.11494290120536185),\n",
       " (5, 6, 0.11317138749833287),\n",
       " (11, 10, 0.1111817445215301)]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('cleantransfer')",
   "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.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "1262644d422a82105623d802af89e6e9ed007275a479308d900e7081e9b63df8"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}