Linguistic_minimal_pairs.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "We used a modification of Ripser++ that includes additional optimization for faster (up to 55%) computation of RTD barcodes. Normal Ripser++ can also be used instead from our modification."
      ],
      "metadata": {
        "id": "Q-ERHtMXSElC"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QBLF8bO_f-QI"
      },
      "outputs": [],
      "source": [
        "!pip install git+https://github.com/ArGintum/RipserZeros.git"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "If you don't have Transformers package installed, please do it now --- it is  required to work with BERT"
      ],
      "metadata": {
        "id": "6dpo9pI6mvpF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#!pip install transformers"
      ],
      "metadata": {
        "id": "Y1ZUXvUHZd4i"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "ulGHUP8of_hl"
      },
      "outputs": [],
      "source": [
        "import itertools\n",
        "import re\n",
        "import subprocess\n",
        "import json\n",
        "import torch\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "from collections import defaultdict\n",
        "from tqdm.notebook import tqdm\n",
        "from transformers import BertTokenizer, BertModel, BertConfig, BertForSequenceClassification\n",
        "from transformers import XLMRobertaTokenizer, XLMRobertaForSequenceClassification, XLMRobertaConfig\n",
        "\n",
        "import ripserplusplus as rpp_py\n",
        "\n",
        "# Fix random seed for reproduciability \n",
        "np.random.seed(42)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1. Computing RTD scores on BLIMP data"
      ],
      "metadata": {
        "id": "tnI4ouUMnTPu"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Download pretrained BERT model"
      ],
      "metadata": {
        "id": "B3-NsAMmWaak"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 252,
          "referenced_widgets": [
            "862a36f51649445dbd240f95ada2a5ad",
            "34f395fa37044bbe99beeb7360518c17",
            "51cc6f14d2734b5e8d65ede4c9ddb3d4",
            "74997a64156b41deb796a17ea723c201",
            "81a874f137f44faaa43be284530607d6",
            "3a39b709749343a2ab5895c593240b86",
            "9d2e5699d2394b3ebb4a98ad73a4f06d",
            "b3a65127c72c4d288e25b24a7a89ca1b",
            "532796128b334eb986fa69dcafa754b8",
            "03990fd85c69437fa89970d340818e96",
            "10a9e7d0768f4e10bd5e40f4b6586692",
            "d5bba8f5baa34f37b373a4e138a85847",
            "dffc9b9127504709a4d54103971959d1",
            "df3fa7a020a64369b052928eb20919ea",
            "9dd0f6e7830e4ef48ec83b7a312e0220",
            "0140a4dc47f24a3487d80e0d4b7b9b6f",
            "eb4edee7dfaf441ea9b3c08d06a91aa1",
            "df6d1685ed5a490abc82f6dae6f54615",
            "4a751133491f4bc096a1c348047dfd1c",
            "e2b3490ff3a04e9db1cddbd062ac31c2",
            "26dcd9ac60ab48aba05c38a8d3a6dd4e",
            "a8c509aef293412e96aeb48e8fd40fe5",
            "473c0a01091e4388872430b8c0c44d35",
            "703339200c4347ba8f218541630f704a",
            "a23188c1dd1a40279f5f6aa9069172f9",
            "4b8d21c07afd40ccb8855eebca0c7a41",
            "a6abd989b82747eea590ef6cace9707a",
            "a513662430504674acecfd65ae79b25a",
            "150fb1efff694044b362a9ed73d591c6",
            "b464ae82d45e4d7a881225d52930b088",
            "0b04d9f8940c423c9cc996bfe9f168b0",
            "e0582498fa9049d49d5230577e1c78a3",
            "40f0492a5f36452bbd46277ca5f11d65",
            "2552dd1401f44610929f4b4759e4d94e",
            "162e4b82696e4da3af6d96f75d35f57d",
            "b83eaa17ea60438ca1900b2bcda556ec",
            "915f4e18bf9b4666b3d549d6663a7937",
            "60f03b26131448ea9ad9fa13f78eb467",
            "f73164e38aa64ea9a2d6487ee8f5d0de",
            "47243decf36946b28574eee5214870eb",
            "7b72032555b542e9b1a31e0d772e3de3",
            "93f70d89da4040e690e554d936626adc",
            "00d465bda1794159b06b176fa61bc684",
            "56a043b6985e42e996a9bf5067d34699"
          ]
        },
        "id": "YoxRMyGBgAxV",
        "outputId": "b7c3b932-6a72-40a5-a808-eb316695534e"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading:   0%|          | 0.00/570 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "862a36f51649445dbd240f95ada2a5ad"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading:   0%|          | 0.00/416M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "d5bba8f5baa34f37b373a4e138a85847"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Some weights of the model checkpoint at bert-base-cased were not used when initializing BertForSequenceClassification: ['cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.bias', 'cls.predictions.decoder.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.bias']\n",
            "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
            "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading:   0%|          | 0.00/208k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "473c0a01091e4388872430b8c0c44d35"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading:   0%|          | 0.00/29.0 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "2552dd1401f44610929f4b4759e4d94e"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "device='cuda'\n",
        "model_path = tokenizer_path = \"bert-base-cased\"\n",
        "\n",
        "config = BertConfig.from_pretrained('bert-base-cased', output_hidden_states=False, output_attentions=True)\n",
        "model = BertForSequenceClassification.from_pretrained(model_path, config=config).to(device)\n",
        "tokenizer = BertTokenizer.from_pretrained(tokenizer_path, do_lower_case=False)\n",
        "\n",
        "MAX_LEN = max_tokens_amount = 128"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Set dataset structure: dictionary of pairs `category name' : list of filenames containing data about phenomenons of that category"
      ],
      "metadata": {
        "id": "XXAPMpHzWg3U"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "eGRAGYcygFJl"
      },
      "outputs": [],
      "source": [
        "blimp_structure = {\n",
        "    'Anaphor agreement' : ['anaphor_gender_agreement.json', 'anaphor_number_agreement.json'],\n",
        "    'Argument structure' : ['animate_subject_passive.json', 'animate_subject_trans.json', 'causative.json', 'drop_argument.json', 'inchoative.json', 'intransitive.json', 'passive_1.json', 'passive_2.json', 'transitive.json'],\n",
        "    'Binding' : ['principle_A_c_command.json', 'principle_A_case_1.json', 'principle_A_case_2.json', 'principle_A_domain_1.json', 'principle_A_domain_2.json', 'principle_A_domain_3.json', 'principle_A_reconstruction.json'],\n",
        "    'Control / raising' : ['existential_there_object_raising.json', 'existential_there_subject_raising.json', 'expletive_it_object_raising.json', 'tough_vs_raising_1.json', 'tough_vs_raising_2.json'],\n",
        "    'Determiner-Noun agr.' : ['determiner_noun_agreement_1.json', 'determiner_noun_agreement_2.json', 'determiner_noun_agreement_irregular_1.json', 'determiner_noun_agreement_irregular_2.json', \n",
        "                              'determiner_noun_agreement_with_adjective_1.json', 'determiner_noun_agreement_with_adj_2.json', 'determiner_noun_agreement_with_adj_irregular_1.json', \n",
        "                              'determiner_noun_agreement_with_adj_irregular_2.json'],\n",
        "    'Ellipsis' : ['ellipsis_n_bar_1.json', 'ellipsis_n_bar_2.json'],\n",
        "    'Filler gap' : ['wh_questions_object_gap.json', 'wh_questions_subject_gap.json', 'wh_questions_subject_gap_long_distance.json', 'wh_vs_that_no_gap.json', 'wh_vs_that_no_gap_long_distance.json',\n",
        "                    'wh_vs_that_with_gap.json', 'wh_vs_that_with_gap_long_distance.json'],\n",
        "    'Irregular forms' : ['irregular_past_participle_adjectives.json', 'irregular_past_participle_verbs.json'],\n",
        "    'Island effects' : ['adjunct_island.json', 'complex_NP_island.json', 'coordinate_structure_constraint_complex_left_branch.json', 'coordinate_structure_constraint_object_extraction.json',\n",
        "                        'left_branch_island_echo_question.json', 'left_branch_island_simple_question.json', 'sentential_subject_island.json', 'wh_island.json'],\n",
        "    'NPI licensing' : ['matrix_question_npi_licensor_present.json', 'npi_present_1.json', 'npi_present_2.json', 'only_npi_licensor_present.json', 'only_npi_scope.json', 'sentential_negation_npi_licensor_present.json',\n",
        "                       'sentential_negation_npi_scope.json'],\n",
        "    'Quantifiers' : ['existential_there_quantifiers_1.json', 'existential_there_quantifiers_2.json', 'superlative_quantifiers_1.json', 'superlative_quantifiers_2.json'],\n",
        "    'Subject-verb agr.' : ['distractor_agreement_relational_noun.json', 'distractor_agreement_relative_clause.json', 'irregular_plural_subject_verb_agreement_1.json', 'irregular_plural_subject_verb_agreement_2.json',\n",
        "                           'regular_plural_subject_verb_agreement_1.json', 'regular_plural_subject_verb_agreement_2.json']\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Function for computing RTD of two graphs given as adjacency matrices\n"
      ],
      "metadata": {
        "id": "4ae4u_XyWkF6"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "id": "SoW0Z6U_gH5N"
      },
      "outputs": [],
      "source": [
        "def rtd(a, b, dim=1, mode='max'):\n",
        "  # Parameters:\n",
        "  #           a    - adjacency matrix (symmetric) of the first graph\n",
        "  #           b    - adjacency matrix (symmetric) of the first graph\n",
        "  #           dim  - dimension of RTD. Default: 1 (in all our experiments we used this value only)\n",
        "  #           mode - variant of RTD to use. Default: 'max' (in all our experiments we used this value only)\n",
        "\n",
        "  # Returns: list of bars of RTD(<a>, <b>) in dimension <dim>\n",
        "\n",
        "  if mode != 'max' and mode != 'min':\n",
        "     raise ValueError(\"Wrong Value for 'mode' parameter. RTD mode can be either 'min' or 'max' \")\n",
        "\n",
        "  n = a.shape[0]\n",
        "  d = np.zeros((2 * n, 2 * n))\n",
        "  \n",
        "  if mode == 'max':\n",
        "    d[n:, :n] = torch.maximum(a.cpu(), b.cpu()).detach().numpy()  \n",
        "    d[:n, n:] = torch.maximum(a.cpu(), b.cpu()).detach().numpy()  \n",
        "    d[n:, n:] = b.cpu().detach().numpy()\n",
        "  else:\n",
        "    d[n:, :n] = a.cpu().detach().numpy()\n",
        "    d[:n, n:] = a.cpu().detach().numpy()\n",
        "    d[n:, n:] = torch.minimum(a.cpu(), b.cpu()).detach().numpy() \n",
        "\n",
        "  m = d[n:, :n].mean()\n",
        "  d[d < m*(1e-6)] = 0\n",
        "  results = rpp_py.run(\"--format distance --mode rtd --dim \" + str(dim), d)['dgms'][dim]\n",
        "  return results"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Main function for scoring Rules on individual heads"
      ],
      "metadata": {
        "id": "yZ0uJXSrIoHP"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 56,
      "metadata": {
        "id": "cWQd8FMrtkxi"
      },
      "outputs": [],
      "source": [
        "def attentive_work_in_class(model, phenomenons, root_dir=\"\", metric=\"rtd\", finish_pos=1000, skip=0, attention_layers=12, attention_heads=12):\n",
        "  # Parameters:\n",
        "  #           model       - linguistic model for scoring\n",
        "  #           phenomenons - dictionary structure of the dataset\n",
        "  #           root_dir    - path to the directory with data files\n",
        "  #           metric      - topology metric to be used for scoring. Can be either 'rtd' for RTD or 'h0s' for H0_s.\n",
        "  #           finish_pos   \n",
        "  #           skip        - last two arguments gives the range of strings from data files to use : [skip; finish_pos)\n",
        "  #                         Default: skip = 0, finish_pos = 1000 --- this gives all strings from original BLIMP files\n",
        "\n",
        "  # Returns: matrix of size (finish_pos - slip)*<total number of phenomenons> X <total number of heads>\n",
        "\n",
        "  if metric != 'rtd' and metric != 'h0s':\n",
        "     raise ValueError(\"Wrong Value for 'metric' parameter. Metric to be used can be either 'rtd' or 'h0s' \")\n",
        "\n",
        "  samples = finish_pos - skip\n",
        "  full_matrix = np.zeros((67 * samples, attention_layers * attention_heads))\n",
        "  pos = 0\n",
        "\n",
        "  model.eval()\n",
        "\n",
        "  for name in phenomenons:\n",
        "    for json_name in phenomenons[name]:\n",
        "      step_cnt = 0  \n",
        "\n",
        "      fle = open(root_dir + json_name + 'l', 'r')\n",
        " \n",
        "      for s in fle.readlines():\n",
        "        if step_cnt < skip:\n",
        "          step_cnt += 1\n",
        "          continue\n",
        "\n",
        "        if step_cnt >= finish_pos:\n",
        "          break\n",
        "\n",
        "        data = json.loads(s)\n",
        "        inputs = tokenizer(data['sentence_good'], max_length=128, return_tensors=\"pt\").to(device)\n",
        "        inputs2 = tokenizer(data['sentence_bad'], max_length=128, return_tensors=\"pt\").to(device)\n",
        "\n",
        "        step_cnt += 1\n",
        "\n",
        "        # To ensure one-to-one correspondence between tokens when sequences are unequal we take only the `tail' from the longest\n",
        "        n_len = min(len(inputs['input_ids'][0]), len(inputs2['input_ids'][0]))\n",
        "        goods_list = torch.stack(model(**inputs)[1], dim=0)[:attention_layers, 0, :attention_heads, -n_len:, -n_len:]\n",
        "        bads_list = torch.stack(model(**inputs2)[1], dim=0)[:attention_layers, 0, :attention_heads, -n_len:, -n_len:]\n",
        "        # ---------------------------------------\n",
        "\n",
        "        for i in range(attention_layers):\n",
        "          for j in range(attention_heads):       \n",
        "            # First, we symmetrize attention matrices\n",
        "            d_good = 1.0 - (torch.triu(goods_list[i][j], 1) + torch.transpose(torch.triu(goods_list[i][j], 1), 0, 1))\n",
        "            d_good -= torch.diag(torch.diag(d_good))\n",
        "            d_bad = 1.0 - (torch.triu(bads_list[i][j], 1) + torch.transpose(torch.triu(bads_list[i][j], 1), 0, 1))\n",
        "            d_bad -= torch.diag(torch.diag(d_bad)) \n",
        "            # ---------------------------------------\n",
        "\n",
        "            if metric == \"rtd\":\n",
        "              results = rtd(d_good, d_bad, mode='max')\n",
        "            else:\n",
        "              results = rpp_py.run(\"--format distance --mode original --dim 0\", d_good)['dgms'][0]\n",
        "\n",
        "            if len(results) > 0:\n",
        "              term_left = np.mean([tmp_x[1] - tmp_x[0] for tmp_x in results])\n",
        "            else:\n",
        "              term_left = 0.0\n",
        "            \n",
        "            if metric == \"rtd\":\n",
        "              results = rtd(d_bad, d_good, mode='max')\n",
        "            else:\n",
        "              results = rpp_py.run(\"--format distance --mode original --dim 0\", d_bad)['dgms'][0]\n",
        "\n",
        "            if len(results) > 0:\n",
        "              term_right = np.mean([tmp_x[1] - tmp_x[0] for tmp_x in results])\n",
        "            else:\n",
        "              term_right = 0.0\n",
        "\n",
        "            # If Rule I gives the correct result for this pair, corresponding element of the matrix is set to 1\n",
        "            # Otherwise, Rule II is correct (since for A != B holds RTD(A, B) != RTD(B, A) ) and full_matrix[i, j] set to 0\n",
        "            if term_left < term_right:\n",
        "              full_matrix[pos, attention_heads * i + j] = 1\n",
        "            # ---------------------------------------\n",
        "        pos += 1\n",
        "\n",
        "      fle.close()\n",
        "      \n",
        "      smol_floppa = np.mean(full_matrix[pos - samples : pos, :], axis=0)\n",
        "      arg = np.argmax(smol_floppa)\n",
        "      arg2 = np.argmin(smol_floppa)\n",
        "    \n",
        "      print('phenomenon ', json_name[:-5], 'Rule I: best head (', arg // attention_heads, ',', arg % attention_heads, ') acc.: ', 100 * np.max(smol_floppa),  '/// Rule II: best head (', arg2 // attention_heads, ',', arg2 % attention_heads, ') acc.: ', 100 * (1 - np.min(smol_floppa)) )\n",
        "    \n",
        "    big_floppa = np.mean(full_matrix[pos - samples * len(phenomenons[name]) : pos, :], axis=0)\n",
        "    arg = np.argmax(big_floppa)\n",
        "    arg2 = np.argmin(big_floppa)\n",
        "    \n",
        "    print('category ', name, 'Rule I: best head (', arg // attention_heads, ',', arg % attention_heads, ') acc.: ', 100 * np.max(big_floppa),  '/// Rule II: best head (', arg2 // attention_heads, ',', arg2 % attention_heads, ') acc.: ', 100 * (1 - np.min(big_floppa)) )\n",
        "    print('\\n')\n",
        "    \n",
        "  \n",
        "  total_floppa = np.mean(full_matrix, axis=0)\n",
        "  arg = np.argmax(total_floppa)\n",
        "  arg2 = np.argmin(total_floppa)\n",
        "  print('Total. Rule I: best head (', arg // attention_heads, ',', arg % attention_heads, ') acc.: ', 100 * np.max(total_floppa),  '/// Rule II: best head (', arg2 // attention_heads, ',', arg2 % attention_heads, ') acc.: ', 100 * (1 - np.min(total_floppa)) )\n",
        "\n",
        "  return full_matrix"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bE-2MtKCUi7s"
      },
      "outputs": [],
      "source": [
        "fullBERT_rtd_matrix = attentive_work_in_class(model, blimp_structure, root_dir=\"data/BLIMP/\", metric=\"rtd\", finish_pos=10)\n",
        "#np.save('fullBERT_rtd.dat', fullBERT_rtd_matrix)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The same but using H0_S:"
      ],
      "metadata": {
        "id": "_j3GiF9RWSOK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "fullBERT_h0s_matrix = attentive_work_in_class(model, blimp_structure, root_dir=\"data/BLIMP/\", metric=\"h0s\")\n",
        "#np.save('fullBERT_h0s.dat', fullBERT_rtd_matrix)"
      ],
      "metadata": {
        "id": "o4RCVepLWRw7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Compute scores on generated auxillary data to perform ensemble selection"
      ],
      "metadata": {
        "id": "Rp8zwdHcWKvU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "smallBERT_rtd_matrix = attentive_work_in_class(model, blimp_structure, root_dir=\"data/BLIMP_auxillary/\", metric=\"rtd\")\n",
        "#np.save('auxBERT_rtd.dat', smallBERT_rtd_matrix)"
      ],
      "metadata": {
        "id": "gxZJbyJzWHz8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "smallBERT_h0s_matrix = attentive_work_in_class(model, blimp_structure, root_dir=\"data/BLIMP_auxillary/\", metric=\"h0s\")"
      ],
      "metadata": {
        "id": "ZZsXwv9Pdwyn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Head ensamble selection"
      ],
      "metadata": {
        "id": "YKEDdLtrZX3K"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import heapq\n",
        "import copy"
      ],
      "metadata": {
        "id": "KLOZip8eZWzj"
      },
      "execution_count": 60,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Auxillary step: expand matrices to be able to work with Rule II the same way as with Rule I"
      ],
      "metadata": {
        "id": "ptjySEjwgxp9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "full_matrix = np.load('SciWork/test_04_sentence_pairs/full_matrix_cased.dat.npy')\n",
        "generated_matrix = np.load('SciWork/test_04_sentence_pairs/generated/smoll_matrix_cased.dat.npy')\n",
        "\n",
        "full_matrix = np.hstack([full_matrix, 1- full_matrix])\n",
        "generated_matrix = np.hstack([generated_matrix, 1- generated_matrix])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "igTqt_5BizLW",
        "outputId": "255c3dfb-e8c1-47b3-916a-77754280859e"
      },
      "execution_count": 61,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(67000, 288) (6700, 288)\n",
            "(67000, 288)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "full_matrix = np.hstack([fullBERT_rtd_matrix, 1- fullBERT_rtd_matrix])\n",
        "generated_matrix = np.hstack([smallBERT_rtd_matrix, 1- smallBERT_rtd_matrix])\n",
        "\n",
        "print(full_matrix.shape, generated_matrix.shape)\n",
        "print(full_matrix.shape)"
      ],
      "metadata": {
        "id": "YF-PllQSgtLX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Function for selecting best ensamble of heads"
      ],
      "metadata": {
        "id": "e5CB05O8egCX"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 78,
      "metadata": {
        "id": "XIAPr1vq93m1"
      },
      "outputs": [],
      "source": [
        "def calc_proc(full_matrix, generated_matrix, phenomenons, max_len=51, max_queue_size=40):\n",
        "  # Parameters:\n",
        "  #           full_matrix      - dataset to score on (a.k.a. test data)\n",
        "  #           generated_matrix - dataset to select ensembles on (a.k.a. train data)\n",
        "  #           phenomenons      - dictionary structure of the dataset\n",
        "  #           max_len          - maximal number of heads in the ensemble\n",
        "  #           max_queue_size   - regulates the `width' of the beam search\n",
        "\n",
        "\n",
        "  # Output: \n",
        "  #           Dictionary with keys <acc> - Accuracy of the best ensemble\n",
        "  #                                <Rule I vertices>  - List of heads as pairs (# of layer, # of head within layer) that are used with Rule I in the best ensemble\n",
        "  #                                <Rule II vertices> - List of heads as pairs (# of layer, # of head within layer) that are used with Rule II in the best ensemble\n",
        "\n",
        "  small_matrix = generated_matrix\n",
        "\n",
        "  entities = [ [np.mean(small_matrix[:, x] > 0.5), -x, [x]] for x in range(288)]\n",
        "  cnt = 0\n",
        "  best_score = -1\n",
        "  best_vertex_list = []\n",
        "  print(\"Length 1 --- calculations completed\")\n",
        "  for i in range(max_len // 2):\n",
        "      h = []\n",
        "      for elem in entities:\n",
        "        for new_vert in range(288):\n",
        "          for new_vert_2 in range(new_vert):\n",
        "            if new_vert in elem[2] or new_vert_2 in elem[2] or new_vert - new_vert_2 == 144:\n",
        "              continue\n",
        "          \n",
        "            msk = (np.sum(small_matrix[:, elem[2]], axis=1) + small_matrix[:, new_vert] + small_matrix[:, new_vert_2]) / (len(elem[2]) + 2)\n",
        "            sum = np.mean(msk > 0.5)\n",
        "            lst = elem[2] + [new_vert, new_vert_2]\n",
        "\n",
        "            if sum >= elem[0]:\n",
        "              if len(h) >= max_queue_size:\n",
        "                heapq.heapreplace(h, [sum, cnt, lst])\n",
        "              else:\n",
        "                heapq.heappush(h, [sum, cnt, lst])\n",
        "\n",
        "              cnt += 1\n",
        "\n",
        "      entities = copy.deepcopy(sorted(h))\n",
        "\n",
        "      if len(entities) == 0:\n",
        "        continue\n",
        "      best = sorted(entities)[-1]\n",
        "\n",
        "      sm = np.zeros(full_matrix.shape[0])\n",
        "      vert_list = []\n",
        "\n",
        "      for vert in best[2]:\n",
        "        sm += full_matrix[: , vert]\n",
        "        vert_list.append((vert // 12, vert % 12))\n",
        "\n",
        "      sm /= len(best[2])\n",
        "      sm = (sm > 0.5).astype(int)\n",
        "\n",
        "      if best_score >= 100 * np.mean(sm):\n",
        "        break\n",
        "      else:\n",
        "        best_score =  100 * np.mean(sm)\n",
        "        best_vertex_list = vert_list\n",
        "      print(\"Length\", i * 2 + 3, \" --- calculations completed\")\n",
        "\n",
        "  result_dict = {'acc' : best_score, 'Rule I vertices' : [], 'Rule II vertices' : []}\n",
        "  for vert in vert_list:\n",
        "    if vert[0] < 12:\n",
        "      result_dict['Rule I vertices'].append(vert)\n",
        "    else:\n",
        "      vert[0] -= 12\n",
        "      result_dict['Rule II vertices'].append(vert)\n",
        "\n",
        "  return result_dict"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "A simple example:"
      ],
      "metadata": {
        "id": "ppPSyz3inYiU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "result_dict = calc_proc(full_matrix, generated_matrix, blimp_structure, max_len=3, max_queue_size=2)\n",
        "print(\"Results: \", result_dict)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HEsWiZVRjNHl",
        "outputId": "8f1e1dc8-0324-462d-b26a-fa7f4b4cdfa6"
      },
      "execution_count": 80,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Length 1 --- calculations completed\n",
            "Length 3  --- calculations completed\n",
            "Results:  {'acc': 55.159701492537316, 'Rule I vertices': [(0, 7), (0, 9), (0, 1)], 'Rule II vertices': []}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Setup as it was used in our work:"
      ],
      "metadata": {
        "id": "BmGAjI7jnaYD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "calc_proc(full_matrix, generated_matrix, blimp_structure, max_len=51, max_queue_size=40)"
      ],
      "metadata": {
        "id": "sgDcJX50nVcj"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "Linguistic_minimal_pairs.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "862a36f51649445dbd240f95ada2a5ad": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_34f395fa37044bbe99beeb7360518c17",
              "IPY_MODEL_51cc6f14d2734b5e8d65ede4c9ddb3d4",
              "IPY_MODEL_74997a64156b41deb796a17ea723c201"
            ],
            "layout": "IPY_MODEL_81a874f137f44faaa43be284530607d6"
          }
        },
        "34f395fa37044bbe99beeb7360518c17": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_3a39b709749343a2ab5895c593240b86",
            "placeholder": "​",
            "style": "IPY_MODEL_9d2e5699d2394b3ebb4a98ad73a4f06d",
            "value": "Downloading: 100%"
          }
        },
        "51cc6f14d2734b5e8d65ede4c9ddb3d4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b3a65127c72c4d288e25b24a7a89ca1b",
            "max": 570,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_532796128b334eb986fa69dcafa754b8",
            "value": 570
          }
        },
        "74997a64156b41deb796a17ea723c201": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_03990fd85c69437fa89970d340818e96",
            "placeholder": "​",
            "style": "IPY_MODEL_10a9e7d0768f4e10bd5e40f4b6586692",
            "value": " 570/570 [00:00&lt;00:00, 12.5kB/s]"
          }
        },
        "81a874f137f44faaa43be284530607d6": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3a39b709749343a2ab5895c593240b86": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9d2e5699d2394b3ebb4a98ad73a4f06d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "b3a65127c72c4d288e25b24a7a89ca1b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "532796128b334eb986fa69dcafa754b8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "03990fd85c69437fa89970d340818e96": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "10a9e7d0768f4e10bd5e40f4b6586692": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "d5bba8f5baa34f37b373a4e138a85847": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_dffc9b9127504709a4d54103971959d1",
              "IPY_MODEL_df3fa7a020a64369b052928eb20919ea",
              "IPY_MODEL_9dd0f6e7830e4ef48ec83b7a312e0220"
            ],
            "layout": "IPY_MODEL_0140a4dc47f24a3487d80e0d4b7b9b6f"
          }
        },
        "dffc9b9127504709a4d54103971959d1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_eb4edee7dfaf441ea9b3c08d06a91aa1",
            "placeholder": "​",
            "style": "IPY_MODEL_df6d1685ed5a490abc82f6dae6f54615",
            "value": "Downloading: 100%"
          }
        },
        "df3fa7a020a64369b052928eb20919ea": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_4a751133491f4bc096a1c348047dfd1c",
            "max": 435779157,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_e2b3490ff3a04e9db1cddbd062ac31c2",
            "value": 435779157
          }
        },
        "9dd0f6e7830e4ef48ec83b7a312e0220": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_26dcd9ac60ab48aba05c38a8d3a6dd4e",
            "placeholder": "​",
            "style": "IPY_MODEL_a8c509aef293412e96aeb48e8fd40fe5",
            "value": " 416M/416M [00:12&lt;00:00, 49.1MB/s]"
          }
        },
        "0140a4dc47f24a3487d80e0d4b7b9b6f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "eb4edee7dfaf441ea9b3c08d06a91aa1": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "df6d1685ed5a490abc82f6dae6f54615": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "4a751133491f4bc096a1c348047dfd1c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e2b3490ff3a04e9db1cddbd062ac31c2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "26dcd9ac60ab48aba05c38a8d3a6dd4e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a8c509aef293412e96aeb48e8fd40fe5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "473c0a01091e4388872430b8c0c44d35": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_703339200c4347ba8f218541630f704a",
              "IPY_MODEL_a23188c1dd1a40279f5f6aa9069172f9",
              "IPY_MODEL_4b8d21c07afd40ccb8855eebca0c7a41"
            ],
            "layout": "IPY_MODEL_a6abd989b82747eea590ef6cace9707a"
          }
        },
        "703339200c4347ba8f218541630f704a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a513662430504674acecfd65ae79b25a",
            "placeholder": "​",
            "style": "IPY_MODEL_150fb1efff694044b362a9ed73d591c6",
            "value": "Downloading: 100%"
          }
        },
        "a23188c1dd1a40279f5f6aa9069172f9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b464ae82d45e4d7a881225d52930b088",
            "max": 213450,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_0b04d9f8940c423c9cc996bfe9f168b0",
            "value": 213450
          }
        },
        "4b8d21c07afd40ccb8855eebca0c7a41": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_e0582498fa9049d49d5230577e1c78a3",
            "placeholder": "​",
            "style": "IPY_MODEL_40f0492a5f36452bbd46277ca5f11d65",
            "value": " 208k/208k [00:00&lt;00:00, 613kB/s]"
          }
        },
        "a6abd989b82747eea590ef6cace9707a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a513662430504674acecfd65ae79b25a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "150fb1efff694044b362a9ed73d591c6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "b464ae82d45e4d7a881225d52930b088": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0b04d9f8940c423c9cc996bfe9f168b0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "e0582498fa9049d49d5230577e1c78a3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "40f0492a5f36452bbd46277ca5f11d65": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "2552dd1401f44610929f4b4759e4d94e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_162e4b82696e4da3af6d96f75d35f57d",
              "IPY_MODEL_b83eaa17ea60438ca1900b2bcda556ec",
              "IPY_MODEL_915f4e18bf9b4666b3d549d6663a7937"
            ],
            "layout": "IPY_MODEL_60f03b26131448ea9ad9fa13f78eb467"
          }
        },
        "162e4b82696e4da3af6d96f75d35f57d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f73164e38aa64ea9a2d6487ee8f5d0de",
            "placeholder": "​",
            "style": "IPY_MODEL_47243decf36946b28574eee5214870eb",
            "value": "Downloading: 100%"
          }
        },
        "b83eaa17ea60438ca1900b2bcda556ec": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7b72032555b542e9b1a31e0d772e3de3",
            "max": 29,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_93f70d89da4040e690e554d936626adc",
            "value": 29
          }
        },
        "915f4e18bf9b4666b3d549d6663a7937": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_00d465bda1794159b06b176fa61bc684",
            "placeholder": "​",
            "style": "IPY_MODEL_56a043b6985e42e996a9bf5067d34699",
            "value": " 29.0/29.0 [00:00&lt;00:00, 542B/s]"
          }
        },
        "60f03b26131448ea9ad9fa13f78eb467": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f73164e38aa64ea9a2d6487ee8f5d0de": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "47243decf36946b28574eee5214870eb": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "7b72032555b542e9b1a31e0d772e3de3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "93f70d89da4040e690e554d936626adc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "00d465bda1794159b06b176fa61bc684": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "56a043b6985e42e996a9bf5067d34699": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}