Word sense disambiguation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('/home/felix/Workspaces/PhD/lm-inspector/.vn')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!which python\n",
    "!which pip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "import pandas\n",
    "import torch\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch import nn\n",
    "\n",
    "from transformers import AutoConfig, AutoTokenizer, AutoModel\n",
    "\n",
    "class BERTEncoder(nn.Module):\n",
    "\n",
    "    def __init__(self, model_name, device, sample_size):\n",
    "        super().__init__()\n",
    "        config = AutoConfig.from_pretrained(model_name,\n",
    "                                            output_hidden_states=True,\n",
    "                                            output_attentions=True\n",
    "                                            )\n",
    "        self.bert = AutoModel.from_pretrained(model_name, config=config)\n",
    "        self.output_size = self.bert.config.hidden_size\n",
    "        self.device = device\n",
    "\n",
    "        # For logging purposes\n",
    "        self.sample_idx = 0\n",
    "        self.sample_size = sample_size\n",
    "\n",
    "    def forward(self, documents):\n",
    "        # Convert the encoded sequences to a tensor.\n",
    "        docs = torch.tensor([d['doc'] for d in documents]).to(self.device)\n",
    "        # Get the word representations of all sequences.\n",
    "        bert_outputs = self.bert(docs)\n",
    "        # Use the top layer of BERT\n",
    "        top_layer_output = bert_outputs[0]\n",
    "        # word_outputs will store the word representation of the ambigious word.\n",
    "        word_outputs = torch.empty( len(documents), self.bert.config.hidden_size).to(self.device)\n",
    "        #word_outputs = top_layer_output[:,0,:]\n",
    "        for idx, d in enumerate(documents):\n",
    "          pos = d['pos']\n",
    "          word_outputs[idx] = top_layer_output[idx, pos]\n",
    "\n",
    "        # For logging purposes, since this can take a long time.\n",
    "        print('\\r' + str(self.sample_idx) + \" / \" + str(self.sample_size), end='')\n",
    "        sys.stdout.flush()\n",
    "        self.sample_idx += len(documents)\n",
    "        if self.sample_idx >= self.sample_size:\n",
    "          self.sample_idx = 0\n",
    "\n",
    "        return word_outputs\n",
    "\n",
    "encoder = BERTEncoder('distilbert-base-cased', 'cpu', 10000)\n",
    "\n",
    "def context(doc, pos, window_size):\n",
    "\n",
    "  new_pos = pos\n",
    "  if pos < window_size:\n",
    "    left = 0\n",
    "  else:\n",
    "    left = pos - window_size\n",
    "    new_pos = window_size\n",
    "\n",
    "  window_doc = doc[left: pos + window_size]\n",
    "\n",
    "  # Sanity check that new_pos still refers to the ambigious word.\n",
    "  assert doc[pos] == window_doc[new_pos]\n",
    "\n",
    "  return {'doc': window_doc, 'pos': new_pos}\n",
    "\n",
    "def read_data(filename, window_size):\n",
    "  column_names = ['sense_key', 'lemma', 'pos', 'text']\n",
    "  df = pandas.read_csv(filename, sep='\\t', names=column_names)\n",
    "  X = []\n",
    "  Y = []\n",
    "  for idx, row in df.iterrows():\n",
    "    pos, text, sense_key = row['pos'], row['text'].split(), row['sense_key']\n",
    "    c = context(text, pos, window_size)\n",
    "    X.append(c)\n",
    "    Y.append(sense_key)\n",
    "  return X, Y\n",
    "\n",
    "def bert_tokenize_and_encode(tokenizer, X, max_len):\n",
    "  for x in X:\n",
    "    x['doc'] = tokenizer.encode(x['doc'], max_length=max_len, pad_to_max_length=True, add_special_tokens=False)\n",
    "  return X\n",
    "\n",
    "X, Y = read_data('examples/wsd_train.txt', 64)\n",
    "\n",
    "seq = torch.nn.Sequential(\n",
    "            encoder,\n",
    "            torch.nn.Dropout(0.2),\n",
    "            torch.nn.Linear(encoder.output_size, out_features=222)\n",
    "        )\n",
    "state_dict = torch.load('models/pre-trained/wsd-clf.pt', map_location='cpu')\n",
    "seq.load_state_dict(state_dict)\n",
    "\n",
    "_, Xval, _, Yval = train_test_split(X, Y, test_size=0.2, random_state=0)\n",
    "config = AutoConfig.from_pretrained('distilbert-base-cased',\n",
    "                                    output_hidden_states=True,\n",
    "                                    output_attentions=True\n",
    "                                    )\n",
    "\n",
    "tokenizer = tokenizer = AutoTokenizer.from_pretrained('distilbert-base-cased', config=config)\n",
    "Xval = bert_tokenize_and_encode(tokenizer, Xval, 128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lm_inspect.inspect import LanguageModelInspector\n",
    "inspector = LanguageModelInspector(seq, Xval[:100], Yval[:100], tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <body>\n",
       "          <div id=\"container\">\n",
       "            <table id=\"model\">\n",
       "            </table>\n",
       "            <table id=\"topk\">\n",
       "            </table>\n",
       "          <div/>\n",
       "        </body>\n",
       "        <script>\n",
       "            results = {'indices': [[[0, 1103, 100, 1104, 117]], [[0, 1103, 100, 1104, 117]], [[0, 1103, 100, 1104, 117]]], 'values': [[[3119.21728515625, 485.5113220214844, 423.2947082519531, 311.3454284667969, 280.5463562011719]], [[2866.37841796875, 483.1680603027344, 470.1274108886719, 319.6204833984375, 282.1422424316406]], [[2963.950439453125, 492.59637451171875, 432.6274108886719, 312.6864929199219, 283.74169921875]]]}\n",
       "            tokenizer = {0: '[ P A D ]', 1103: 't h e', 100: '[ U N K ]', 1104: 'o f', 117: ','}\n",
       "            require.config({paths: {d3: \"http://localhost:8888/notebooks/d3/d3.v4.min\"}});\n",
       "require([\"d3\"], function(d3) {\n",
       "  var margin = {top: 1, right: 1, bottom: 1, left: 1},\n",
       "      width = 300 - margin.left - margin.right,\n",
       "      height = 145 - margin.top - margin.bottom;\n",
       "\n",
       "    function round(value, decimals) {\n",
       "      return Number(Math.round(value+'e'+decimals)+'e-'+decimals);\n",
       "    }\n",
       "\n",
       "    window.fetchByLabels = function(labels, layers, heads, type) {\n",
       "      return new Promise((resolve, reject) => {\n",
       "          var callbacks = {\n",
       "              iopub: {\n",
       "                  output: (data) => resolve(data.content.text.trim())\n",
       "              }\n",
       "          };\n",
       "          Jupyter.notebook.kernel.execute(`print(${python})`, callbacks);\n",
       "      });\n",
       "    };\n",
       "\n",
       "    var tmpTokenizer = {8960: 'servants', 2049: 'force', 9602: 'Lithuania', 1669: 'period', 1425: 'age', 1556: 'With', 1817: 'leave', 6391: 'pregnant', 11037: 'Insurance', 1316: 'held', 2601: 'Law', 170: 'a', 1578: 'career', 8108: 'servant', 1201: 'years', 2740: 'cases', 6714: 'dismissed', 1210: 'three', 1470: 'public', 1219: 'during', 1482: 'children', 2250: 'Republic', 5965: 'reaches', 1103: 'the', 1104: 'of', 1105: 'and', 1235: 'until', 2516: 'Service', 1111: 'for', 1113: 'on',\n",
       "    5986: 'insurance', 1250: 'work', 100: '[UNK]', 6245: 'benefits', 2790: 'provides', 1126: 'an', 2027: 'child', 113: '(', 114: ')', 1137: 'or', 117: ',', 24950: 'paragraph', 119: '.', 9467: 'specified', 17663: 'sickness',\n",
       "  100: '[UNK]', 0: '[PAD]', 117: ',', 3854: 'Nations', 24950: 'paragraph', 123: '2', 1244: 'United', 1413: 'line', 2373: 'read', 6606: 'Fund', 13563: 'Property', 6534: 'Organization', 1367: '12', 4288: 'Children', 124: '3', 1291: 'World', 8500: 'Page', 1370: 'For'}\n",
       "\n",
       "\n",
       "    // Temporary hack\n",
       "    var data = results;\n",
       "    var tmpTokenizer = tokenizer;\n",
       "\n",
       "    var uniqueIndices = data.indices\n",
       "      .reduce((prev, curr) => prev.concat(curr), [])\n",
       "      .reduce((prev, curr) => prev.concat(curr), [])\n",
       "      .filter((item, i, arr) => arr.indexOf(item) === i);\n",
       "\n",
       "\n",
       "\n",
       "    var myColor = d3.scaleOrdinal().domain(uniqueIndices).range(d3.schemeCategory20b);//([\"gold\", \"blue\", \"green\", \"yellow\", \"black\", \"grey\", \"darkgreen\", \"pink\", \"brown\", \"slateblue\", \"grey1\", \"orange\"])\n",
       "\n",
       "    window.decodeWord = function(index) {\n",
       "      return tmpTokenizer[index];\n",
       "    };\n",
       "\n",
       "    var modelTable = d3.select('#model');\n",
       "    var topkList = d3.select('#topk');\n",
       "\n",
       "    var update = function(nofLayers, nofHeads, nofWords = 10) {\n",
       "\n",
       "      var svg = modelTable\n",
       "        .selectAll(\"tr\")\n",
       "        .data([...Array(nofHeads).keys()])\n",
       "        .enter()\n",
       "        .append(\"tr\")\n",
       "        .attr(\"head\", function(d) {\n",
       "          return d;\n",
       "        })\n",
       "        .selectAll(\"th\")\n",
       "        .data([...Array(nofLayers).keys()])\n",
       "        .enter()\n",
       "        .append(\"th\")\n",
       "        .attr(\"layer\", function(d) {\n",
       "          return d;\n",
       "        })\n",
       "        .append(\"svg\");\n",
       "\n",
       "      var svgData =\n",
       "        svg\n",
       "        .selectAll(\"rect\")\n",
       "        .data(function(d) {\n",
       "          var head = parseInt(this.parentNode.parentNode.getAttribute(\"head\"));\n",
       "          var layer = d;\n",
       "          var top = data.indices[layer][head].map(function(d, i) {\n",
       "            return {\"id\": d, \"value\": data.values[layer][head][i]};\n",
       "          });\n",
       "          top.push({\"id\": \"root\"});\n",
       "          var topWords = data.indices[layer][head];\n",
       "          var topValues = data.values[layer][head];\n",
       "          topWords.push(\"root\");\n",
       "          topValues.push(\"root\");\n",
       "          var root = d3.stratify()\n",
       "            .id(function(d) { return d.id; })   // Name of the entity (column name is name in csv)\n",
       "            .parentId(function(d, i) {\n",
       "              return d.id === \"root\" ? \"\" : \"root\";\n",
       "            })   // Name of the parent (column name is parent in csv)\n",
       "          (top);\n",
       "          root.sum(function(d,i) {\n",
       "            return +d.value;\n",
       "          });\n",
       "          d3.treemap()\n",
       "            .size([width, height])\n",
       "            .padding(4)\n",
       "            (root);\n",
       "          return root.leaves();\n",
       "        })\n",
       "        .enter();\n",
       "        svgData.append(\"rect\")\n",
       "          .attr('x', function (d) { return d.x0; })\n",
       "          .attr('y', function (d) { return d.y0; })\n",
       "          .attr('width', function (d) { return d.x1 - d.x0; })\n",
       "          .attr('height', function (d) { return d.y1 - d.y0; })\n",
       "          .style(\"stroke\", \"black\")\n",
       "          //.style(\"fill\", \"#69b3a2\")\n",
       "          .style(\"fill\", function(d) {return myColor(d.id); })\n",
       "        //.select(function() {return this.parentNode;})\n",
       "        svgData.append(\"text\")\n",
       "          .attr('x', function (d) { return d.x0+10; })\n",
       "          .attr('y', function (d) { return d.y0+20; })\n",
       "          .text(function(d) {return window.decodeWord(d.id);})\n",
       "          .attr(\"font-size\", \"15px\")\n",
       "          .attr(\"fill\", \"white\");\n",
       "        svgData.append(\"text\")\n",
       "          .attr('x', function (d) { return d.x0+10; })\n",
       "          .attr('y', function (d) { return d.y0+40; })\n",
       "          .text(function(d) {\n",
       "            var prob = d.value;\n",
       "            var percentage = round(prob, 3);\n",
       "            console.log(percentage);\n",
       "            return `${percentage}`;\n",
       "          })\n",
       "          .attr(\"font-size\", \"15px\")\n",
       "          .attr(\"fill\", \"white\");\n",
       "\n",
       "\n",
       "      // Head labels\n",
       "      modelTable\n",
       "        .insert(\"tr\", \"tr\")\n",
       "        .selectAll(\"th\")\n",
       "        .data([...Array(nofLayers)])\n",
       "        .enter()\n",
       "        .append(\"th\")\n",
       "        .html(function(_, d) {\n",
       "          return `l<sub>${d}</sub>`;\n",
       "        });\n",
       "\n",
       "     //Layer labels\n",
       "     modelTable\n",
       "      .selectAll(\"tr\")\n",
       "      .insert(\"th\", \"th\")\n",
       "      .html(function() {\n",
       "        var head = this.parentNode.getAttribute(\"head\");\n",
       "        if (head) {\n",
       "          return `h<sub>${head}</sub>`;\n",
       "        }\n",
       "      });\n",
       "\n",
       "      topkList\n",
       "        .selectAll('tr')\n",
       "        .data(data['indices'][nofLayers-1][nofHeads-1])\n",
       "        .enter()\n",
       "        .append('tr')\n",
       "        .append('th')\n",
       "        .text(function(d) {\n",
       "          return window.decodeWord(d);\n",
       "        })\n",
       "        .style(\"background-color\", function(d) {\n",
       "          return myColor(d);\n",
       "        })\n",
       "        .style(\"color\", function(d) {\n",
       "          return \"white\";\n",
       "        })\n",
       "        .select(function() {\n",
       "          return this.parentNode;\n",
       "        })\n",
       "        .append('th')\n",
       "        .text(function(d, i) {\n",
       "          var prob =  data['values'][nofLayers-1][nofHeads-1][i]\n",
       "          var percentage = round(prob, 3) * 100\n",
       "          return `${percentage}`;\n",
       "        });\n",
       "    };\n",
       "    /*\n",
       "\n",
       "    JSON response type (v1):\n",
       "\n",
       "    {\n",
       "      topk: [\n",
       "        {\n",
       "          \"word\": <word\"\n",
       "          \"prob\": 0-1\n",
       "        },\n",
       "        ...\n",
       "      ]\n",
       "\n",
       "\n",
       "    }\n",
       "\n",
       "    JSON response type (v2):\n",
       "\n",
       "    {\n",
       "      layers: [\n",
       "        0: {\n",
       "            \"word\": <word>,\n",
       "            \"prob\": 0 ... 1\n",
       "          },\n",
       "          { ... }\n",
       "        1: { ... }\n",
       "        ...\n",
       "        12:\n",
       "          {\n",
       "            \"word\": <word>,\n",
       "            \"prob:\" 0 ... 1\n",
       "          }\n",
       "      ],\n",
       "\n",
       "      agg: [\n",
       "        sum: [\n",
       "          {...}\n",
       "        ],\n",
       "        last_layer: [\n",
       "          {...}\n",
       "        ]\n",
       "      ],\n",
       "\n",
       "\n",
       "    }\n",
       "    */\n",
       "\n",
       "    var nofLayers = data.indices.length;\n",
       "    var nofHeads = data.indices[0].length;\n",
       "\n",
       "    update(nofLayers, nofHeads);\n",
       "\n",
       "    // Generate and display some random table data\n",
       "\n",
       "});\n",
       "\n",
       "        </script>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[[[2.0676e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [4.9419e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [3.6749e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [9.1717e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.1893e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.2239e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]],\n",
       "\n",
       "\n",
       "         [[[3.4331e-01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [1.2636e-01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [1.2395e-02, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [6.7487e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [6.8048e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [6.5916e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]],\n",
       "\n",
       "\n",
       "         [[[1.5636e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [2.0119e-01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [2.6601e-01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [4.6155e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [4.6161e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [4.6063e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]]],\n",
       "\n",
       "\n",
       "\n",
       "        [[[[2.2206e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [7.1261e-01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [7.2945e-01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [9.1420e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.2418e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.0078e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]],\n",
       "\n",
       "\n",
       "         [[[1.0558e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [7.0999e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [3.7756e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [7.3137e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [8.8599e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [8.8875e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]],\n",
       "\n",
       "\n",
       "         [[[1.5883e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [1.1443e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [2.0434e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [7.4936e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [8.6858e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [8.7499e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]]],\n",
       "\n",
       "\n",
       "\n",
       "        [[[[1.7842e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [6.2401e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.3302e-01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [9.2162e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.2383e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.2643e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]],\n",
       "\n",
       "\n",
       "         [[[5.8321e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [2.5571e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [1.9570e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [8.7172e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [8.7825e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [8.6893e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]],\n",
       "\n",
       "\n",
       "         [[[1.8572e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [3.9459e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [3.4923e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [8.3755e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [8.3796e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [8.3963e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]]],\n",
       "\n",
       "\n",
       "\n",
       "        ...,\n",
       "\n",
       "\n",
       "\n",
       "        [[[[1.5337e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [4.4155e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [1.0454e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [9.2090e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.2314e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.2601e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]],\n",
       "\n",
       "\n",
       "         [[[3.1191e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [1.3577e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [3.4713e-01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [8.3364e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [8.3345e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [8.2366e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]],\n",
       "\n",
       "\n",
       "         [[[8.0666e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [1.9980e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [2.7427e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [7.7769e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [7.7850e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [7.7993e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]]],\n",
       "\n",
       "\n",
       "\n",
       "        [[[[4.3758e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [1.6828e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [1.6917e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [9.2176e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.2351e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.2623e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]],\n",
       "\n",
       "\n",
       "         [[[5.2864e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [4.2502e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [3.2500e-01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [9.0413e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.0434e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.0233e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]],\n",
       "\n",
       "\n",
       "         [[[1.8088e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [3.0570e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [4.1569e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [8.5453e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [8.5544e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [8.5603e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]]],\n",
       "\n",
       "\n",
       "\n",
       "        [[[[2.5760e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [4.4290e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [6.4028e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [9.1986e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.2090e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [9.2428e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]],\n",
       "\n",
       "\n",
       "         [[[3.3516e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [2.3473e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [1.8964e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [7.5111e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [7.5770e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [7.4771e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]],\n",
       "\n",
       "\n",
       "         [[[1.1316e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [4.9548e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [6.3307e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           ...,\n",
       "           [5.7358e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [5.7219e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00],\n",
       "           [5.7356e+01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
       "            0.0000e+00, 0.0000e+00]]]]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inspector.topk_most_attended(k=2, label=\"case\", layer=[0,3,5], head=1, visualize=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".vn",
   "language": "python",
   "name": ".vn"
  },
  "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}