features_calculation_ripser_and_templates.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import itertools\n",
    "import re\n",
    "import subprocess\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm.notebook import tqdm\n",
    "from transformers import BertTokenizer, BertModel, BertForSequenceClassification\n",
    "\n",
    "from stats_count import *\n",
    "from grab_weights import grab_attention_weights, text_preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "!env | grep CUDA_VISIBLE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(42) # For reproducibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_tokens_amount  = 128 # The number of tokens to which the tokenized text is truncated / padded.\n",
    "    \n",
    "layers_of_interest = [i for i in range(12)]  # Layers for which attention matrices and features on them are \n",
    "                                             # calculated. For calculating features on all layers, leave it be\n",
    "                                             # [i for i in range(12)].\n",
    "\n",
    "model_path = tokenizer_path = \"bert-base-uncased\"  \n",
    "\n",
    "# You can use either standard or fine-tuned BERT. If you want to use fine-tuned BERT to your current task, save the\n",
    "# model and the tokenizer with the commands tokenizer.save_pretrained(output_dir); \n",
    "# bert_classifier.save_pretrained(output_dir) into the same directory and insert the path to it here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Explanation of stats_name parameter\n",
    "\n",
    "Currently, we implemented calculation of the following graphs features:\n",
    "* \"s\"    - amount of strongly connected components\n",
    "* \"w\"    - amount of weakly connected components\n",
    "* \"e\"    - amount of edges\n",
    "* \"v\"    - average vertex degree\n",
    "* \"c\"    - amount of (directed) simple cycles\n",
    "* \"b0b1\" - Betti numbers\n",
    "\n",
    "The variable stats_name contains a string with the names of the features, which you want to calculate. The format of the string is the following:\n",
    "\n",
    "\"stat_name + \"_\" + stat_name + \"_\" + stat_name + ...\"\n",
    "\n",
    "**For example**:\n",
    "\n",
    "`stats_name == \"s_w\"` means that the number of strongly and weakly connected components will be calculated\n",
    "\n",
    "`stats_name == \"b0b1\"` means that only the Betti numbers will be calculated\n",
    "\n",
    "`stats_name == \"b0b1_c\"` means that Betti numbers and the number of simple cycles will be calculated\n",
    "\n",
    "e.t.c."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset = \"test_5k\"           # .csv file with the texts, for which we count topological features\n",
    "input_dir = \"small_gpt_web/\"  # Name of the directory with .csv file\n",
    "output_dir = \"small_gpt_web/\" # Name of the directory with calculations results\n",
    "\n",
    "prefix = output_dir + subset\n",
    "\n",
    "r_file     = output_dir + 'attentions/' + subset  + \"_all_heads_\" + str(len(layers_of_interest)) + \"_layers_MAX_LEN_\" + \\\n",
    "             str(max_tokens_amount) + \"_\" + model_path.split(\"/\")[-1]\n",
    "# Name of the file for attention matrices weights\n",
    "\n",
    "barcodes_file = output_dir + 'barcodes/' + subset  + \"_all_heads_\" + str(len(layers_of_interest)) + \"_layers_MAX_LEN_\" + \\\n",
    "             str(max_tokens_amount) + \"_\" + model_path.split(\"/\")[-1]\n",
    "# Name of the file for barcodes information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'small_gpt_web/barcodes/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "barcodes_file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ".csv file must contain the column with the name **sentence** with the texts. It can also contain the column **labels**, which will be needed for testing. Any other arbitrary columns will be ignored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    data = pd.read_csv(input_dir + subset + \".csv\").reset_index(drop=True)\n",
    "except:\n",
    "    #data = pd.read_csv(input_dir + subset + \".tsv\", delimiter=\"\\t\")\n",
    "    data = pd.read_csv(input_dir + subset + \".tsv\", delimiter=\"\\t\", header=None)\n",
    "    data.columns = [\"0\", \"labels\", \"2\", \"sentence\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Unnamed: 0</th>\n",
       "      <th>id</th>\n",
       "      <th>ended</th>\n",
       "      <th>length</th>\n",
       "      <th>sentence</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>4722</td>\n",
       "      <td>259722</td>\n",
       "      <td>True</td>\n",
       "      <td>231</td>\n",
       "      <td>The Learning Co.\\n\\nDeveloped by\\n\\nThe Learni...</td>\n",
       "      <td>natural</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2757</td>\n",
       "      <td>257813</td>\n",
       "      <td>True</td>\n",
       "      <td>563</td>\n",
       "      <td>Bush doubles down on foreign policy on Saturda...</td>\n",
       "      <td>generated</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2194</td>\n",
       "      <td>257194</td>\n",
       "      <td>True</td>\n",
       "      <td>62</td>\n",
       "      <td>Here are six interesting things you need to kn...</td>\n",
       "      <td>natural</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>817</td>\n",
       "      <td>255817</td>\n",
       "      <td>True</td>\n",
       "      <td>293</td>\n",
       "      <td>Introduction\\n\\nWe would like to thank Antec f...</td>\n",
       "      <td>natural</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>3886</td>\n",
       "      <td>258886</td>\n",
       "      <td>False</td>\n",
       "      <td>1024</td>\n",
       "      <td>ELKRIDGE, Md.—A group called \"Muslims for Trum...</td>\n",
       "      <td>natural</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   Unnamed: 0      id  ended  length  \\\n",
       "0        4722  259722   True     231   \n",
       "1        2757  257813   True     563   \n",
       "2        2194  257194   True      62   \n",
       "3         817  255817   True     293   \n",
       "4        3886  258886  False    1024   \n",
       "\n",
       "                                            sentence      label  \n",
       "0  The Learning Co.\\n\\nDeveloped by\\n\\nThe Learni...    natural  \n",
       "1  Bush doubles down on foreign policy on Saturda...  generated  \n",
       "2  Here are six interesting things you need to kn...    natural  \n",
       "3  Introduction\\n\\nWe would like to thank Antec f...    natural  \n",
       "4  ELKRIDGE, Md.—A group called \"Muslims for Trum...    natural  "
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average amount of words in example: 2723.5122\n",
      "Max. amount of words in example: 6151\n",
      "Min. amount of words in example: 34\n"
     ]
    }
   ],
   "source": [
    "sentences = data['sentence']\n",
    "print(\"Average amount of words in example:\", \\\n",
    "      np.mean(list(map(len, map(lambda x: re.sub('\\w', ' ', x).split(\" \"), data['sentence'])))))\n",
    "print(\"Max. amount of words in example:\", \\\n",
    "      np.max(list(map(len, map(lambda x: re.sub('\\w', ' ', x).split(\" \"), data['sentence'])))))\n",
    "print(\"Min. amount of words in example:\", \\\n",
    "      np.min(list(map(len, map(lambda x: re.sub('\\w', ' ', x).split(\" \"), data['sentence'])))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_token_length(batch_texts):\n",
    "    inputs = tokenizer.batch_encode_plus(batch_texts,\n",
    "       return_tensors='pt',\n",
    "       add_special_tokens=True,\n",
    "       max_length=MAX_LEN,             # Max length to truncate/pad\n",
    "       pad_to_max_length=True,         # Pad sentence to max length\n",
    "       truncation=True\n",
    "    )\n",
    "    inputs = inputs['input_ids'].numpy()\n",
    "    n_tokens = []\n",
    "    indexes = np.argwhere(inputs == tokenizer.pad_token_id)\n",
    "    for i in range(inputs.shape[0]):\n",
    "        ids = indexes[(indexes == i)[:, 0]]\n",
    "        if not len(ids):\n",
    "            n_tokens.append(MAX_LEN)\n",
    "        else:\n",
    "            n_tokens.append(ids[0, 1])\n",
    "    return n_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_LEN = max_tokens_amount\n",
    "tokenizer = BertTokenizer.from_pretrained(tokenizer_path, do_lower_case=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['tokenizer_length'] = get_token_length(data['sentence'].values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Unnamed: 0</th>\n",
       "      <th>id</th>\n",
       "      <th>ended</th>\n",
       "      <th>length</th>\n",
       "      <th>sentence</th>\n",
       "      <th>label</th>\n",
       "      <th>tokenizer_length</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>4722</td>\n",
       "      <td>259722</td>\n",
       "      <td>True</td>\n",
       "      <td>231</td>\n",
       "      <td>The Learning Co.\\n\\nDeveloped by\\n\\nThe Learni...</td>\n",
       "      <td>natural</td>\n",
       "      <td>128</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2757</td>\n",
       "      <td>257813</td>\n",
       "      <td>True</td>\n",
       "      <td>563</td>\n",
       "      <td>Bush doubles down on foreign policy on Saturda...</td>\n",
       "      <td>generated</td>\n",
       "      <td>128</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2194</td>\n",
       "      <td>257194</td>\n",
       "      <td>True</td>\n",
       "      <td>62</td>\n",
       "      <td>Here are six interesting things you need to kn...</td>\n",
       "      <td>natural</td>\n",
       "      <td>71</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>817</td>\n",
       "      <td>255817</td>\n",
       "      <td>True</td>\n",
       "      <td>293</td>\n",
       "      <td>Introduction\\n\\nWe would like to thank Antec f...</td>\n",
       "      <td>natural</td>\n",
       "      <td>128</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>3886</td>\n",
       "      <td>258886</td>\n",
       "      <td>False</td>\n",
       "      <td>1024</td>\n",
       "      <td>ELKRIDGE, Md.—A group called \"Muslims for Trum...</td>\n",
       "      <td>natural</td>\n",
       "      <td>128</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4995</th>\n",
       "      <td>1472</td>\n",
       "      <td>256472</td>\n",
       "      <td>False</td>\n",
       "      <td>1024</td>\n",
       "      <td>Occasionally, we come across interesting scena...</td>\n",
       "      <td>natural</td>\n",
       "      <td>128</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4996</th>\n",
       "      <td>326</td>\n",
       "      <td>255337</td>\n",
       "      <td>False</td>\n",
       "      <td>1024</td>\n",
       "      <td>Providing insight not only into the memes that...</td>\n",
       "      <td>generated</td>\n",
       "      <td>128</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4997</th>\n",
       "      <td>3862</td>\n",
       "      <td>258862</td>\n",
       "      <td>True</td>\n",
       "      <td>339</td>\n",
       "      <td>Each year, MONEY digs into enrollment data and...</td>\n",
       "      <td>natural</td>\n",
       "      <td>128</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4998</th>\n",
       "      <td>2862</td>\n",
       "      <td>257862</td>\n",
       "      <td>False</td>\n",
       "      <td>1024</td>\n",
       "      <td>Grounding of the Queen Elizabeth 2 (response) ...</td>\n",
       "      <td>natural</td>\n",
       "      <td>128</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4999</th>\n",
       "      <td>2439</td>\n",
       "      <td>257487</td>\n",
       "      <td>True</td>\n",
       "      <td>506</td>\n",
       "      <td>Route 128\\n\\nTaxi Use\\n\\nCan 2,245 miles by 2....</td>\n",
       "      <td>generated</td>\n",
       "      <td>128</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5000 rows × 7 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      Unnamed: 0      id  ended  length  \\\n",
       "0           4722  259722   True     231   \n",
       "1           2757  257813   True     563   \n",
       "2           2194  257194   True      62   \n",
       "3            817  255817   True     293   \n",
       "4           3886  258886  False    1024   \n",
       "...          ...     ...    ...     ...   \n",
       "4995        1472  256472  False    1024   \n",
       "4996         326  255337  False    1024   \n",
       "4997        3862  258862   True     339   \n",
       "4998        2862  257862  False    1024   \n",
       "4999        2439  257487   True     506   \n",
       "\n",
       "                                               sentence      label  \\\n",
       "0     The Learning Co.\\n\\nDeveloped by\\n\\nThe Learni...    natural   \n",
       "1     Bush doubles down on foreign policy on Saturda...  generated   \n",
       "2     Here are six interesting things you need to kn...    natural   \n",
       "3     Introduction\\n\\nWe would like to thank Antec f...    natural   \n",
       "4     ELKRIDGE, Md.—A group called \"Muslims for Trum...    natural   \n",
       "...                                                 ...        ...   \n",
       "4995  Occasionally, we come across interesting scena...    natural   \n",
       "4996  Providing insight not only into the memes that...  generated   \n",
       "4997  Each year, MONEY digs into enrollment data and...    natural   \n",
       "4998  Grounding of the Queen Elizabeth 2 (response) ...    natural   \n",
       "4999  Route 128\\n\\nTaxi Use\\n\\nCan 2,245 miles by 2....  generated   \n",
       "\n",
       "      tokenizer_length  \n",
       "0                  128  \n",
       "1                  128  \n",
       "2                   71  \n",
       "3                  128  \n",
       "4                  128  \n",
       "...                ...  \n",
       "4995               128  \n",
       "4996               128  \n",
       "4997               128  \n",
       "4998               128  \n",
       "4999               128  \n",
       "\n",
       "[5000 rows x 7 columns]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "ntokens_array = data['tokenizer_length'].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import ceil\n",
    "\n",
    "batch_size = 10 # batch size\n",
    "number_of_batches = ceil(len(data['sentence']) / batch_size)\n",
    "DUMP_SIZE = 100 # number of batches to be dumped"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculating Ripser features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Format: \"h{dim}\\_{type}\\_{args}\"\n",
    "\n",
    "Dimension: 0, 1, etc.; homology dimension\n",
    "\n",
    "Types: \n",
    "    \n",
    "    1. s: sum of lengths; example: \"h1_s\".\n",
    "    2. m: mean of lengths; example: \"h1_m\"\n",
    "    3. v: variance of lengths; example \"h1_v\"\n",
    "    4. n: number of barcodes with time of birth/death more/less then threshold.\n",
    "        4.1. b/d: birth or death\n",
    "        4.2. m/l: more or less than threshold\n",
    "        4.2. t: threshold value\n",
    "       example: \"h0_n_d_m_t0.5\", \"h1_n_b_l_t0.75\"\n",
    "    5. t: time of birth/death of the longest barcode (not incl. inf).\n",
    "        3.1. b/d: birth of death\n",
    "        example: \"h0_t_d\", \"h1_t_b\"\n",
    "    6. nb: number of barcodes in dim\n",
    "       example: h0_nb\n",
    "    7. e: entropy; example: \"h1_e\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part1of5.npy',\n",
       " 'small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part2of5.npy',\n",
       " 'small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part3of5.npy',\n",
       " 'small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part4of5.npy',\n",
       " 'small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part5of5.npy']"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import timeit\n",
    "import ripser_count\n",
    "\n",
    "adj_filenames = [\n",
    "    output_dir + 'attentions/' + filename \n",
    "    for filename in os.listdir(output_dir + 'attentions/') if r_file in (output_dir + 'attentions/' + filename)\n",
    "]\n",
    "# sorted by part number\n",
    "adj_filenames = sorted(adj_filenames, key = lambda x: int(x.split('_')[-1].split('of')[0][4:].strip())) \n",
    "adj_filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "dim = 1\n",
    "lower_bound = 1e-3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculating and saving barcodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from multiprocessing import Process, Queue\n",
    "\n",
    "def subprocess_wrap(queue, function, args):\n",
    "    queue.put(function(*args))\n",
    "#     print(\"Putted in Queue\")\n",
    "    queue.close()\n",
    "    exit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import itertools\n",
    "from collections import defaultdict\n",
    "\n",
    "def get_only_barcodes(adj_matricies, ntokens_array, dim, lower_bound):\n",
    "    \"\"\"Get barcodes from adj matricies for each layer, head\"\"\"\n",
    "    barcodes = {}\n",
    "    layers, heads = range(adj_matricies.shape[1]), range(adj_matricies.shape[2])\n",
    "    for (layer, head) in itertools.product(layers, heads):\n",
    "        matricies = adj_matricies[:, layer, head, :, :]\n",
    "        barcodes[(layer, head)] = ripser_count.get_barcodes(matricies, ntokens_array, dim, lower_bound, (layer, head))\n",
    "    return barcodes\n",
    "\n",
    "def format_barcodes(barcodes):\n",
    "    \"\"\"Reformat barcodes to json-compatible format\"\"\"\n",
    "    return [{d: b[d].tolist() for d in b} for b in barcodes]\n",
    "\n",
    "def save_barcodes(barcodes, filename):\n",
    "    \"\"\"Save barcodes to file\"\"\"\n",
    "    formatted_barcodes = defaultdict(dict)\n",
    "    for layer, head in barcodes:\n",
    "        formatted_barcodes[layer][head] = format_barcodes(barcodes[(layer, head)])\n",
    "    json.dump(formatted_barcodes, open(filename, 'w'))\n",
    "    \n",
    "def unite_barcodes(barcodes, barcodes_part):\n",
    "    \"\"\"Unite 2 barcodes\"\"\"\n",
    "    for (layer, head) in barcodes_part:\n",
    "        barcodes[(layer, head)].extend(barcodes_part[(layer, head)])\n",
    "    return barcodes\n",
    "\n",
    "def split_matricies_and_lengths(adj_matricies, ntokens, number_of_splits):\n",
    "    splitted_ids = np.array_split(np.arange(ntokens.shape[0]), number_of_splits) \n",
    "    splitted = [(adj_matricies[ids], ntokens[ids]) for ids in splitted_ids]\n",
    "    return splitted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4e24586201cf4088b96b040aecdc694d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Calculating barcodes', max=5.0, style=ProgressStyle(descr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Matricies loaded from: small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part1of5.npy\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=2.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Matricies loaded from: small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part2of5.npy\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=2.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Matricies loaded from: small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part3of5.npy\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=2.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Matricies loaded from: small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part4of5.npy\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=2.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Matricies loaded from: small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part5of5.npy\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=2.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "queue = Queue()\n",
    "number_of_splits = 2\n",
    "for i, filename in enumerate(tqdm(adj_filenames, desc='Calculating barcodes')):\n",
    "    barcodes = defaultdict(list)\n",
    "    adj_matricies = np.load(filename, allow_pickle=True) # samples X \n",
    "    print(f\"Matricies loaded from: {filename}\")\n",
    "    ntokens = ntokens_array[i*batch_size*DUMP_SIZE : (i+1)*batch_size*DUMP_SIZE]\n",
    "    splitted = split_matricies_and_lengths(adj_matricies, ntokens, number_of_splits)\n",
    "    for matricies, ntokens in tqdm(splitted, leave=False):\n",
    "        p = Process(\n",
    "            target=subprocess_wrap,\n",
    "            args=(\n",
    "                queue,\n",
    "                get_only_barcodes,\n",
    "                (matricies, ntokens, dim, lower_bound)\n",
    "            )\n",
    "        )\n",
    "        p.start()\n",
    "        barcodes_part = queue.get() # block until putted and get barcodes from the queue\n",
    "#         print(\"Features got.\")\n",
    "        p.join() # release resources\n",
    "#         print(\"The process is joined.\")\n",
    "        p.close() # releasing resources of ripser\n",
    "#         print(\"The proccess is closed.\")\n",
    "        \n",
    "        barcodes = unite_barcodes(barcodes, barcodes_part)\n",
    "    part = filename.split('_')[-1].split('.')[0]\n",
    "    save_barcodes(barcodes, barcodes_file + '_' + part + '.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculating features of saved barcodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "ripser_feature_names=[\n",
    "    'h0_s', \n",
    "    'h0_e',\n",
    "    'h0_t_d', \n",
    "    'h0_n_d_m_t0.75',\n",
    "    'h0_n_d_m_t0.5',\n",
    "    'h0_n_d_l_t0.25',\n",
    "    'h1_t_b',\n",
    "    'h1_n_b_m_t0.25',\n",
    "    'h1_n_b_l_t0.95', \n",
    "    'h1_n_b_l_t0.70',  \n",
    "    'h1_s',\n",
    "    'h1_e',\n",
    "    'h1_v',\n",
    "    'h1_nb'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['small_gpt_web/barcodes/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part1of5.json',\n",
       " 'small_gpt_web/barcodes/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part2of5.json',\n",
       " 'small_gpt_web/barcodes/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part3of5.json',\n",
       " 'small_gpt_web/barcodes/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part4of5.json',\n",
       " 'small_gpt_web/barcodes/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part5of5.json']"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import timeit\n",
    "import ripser_count\n",
    "import json\n",
    "\n",
    "adj_filenames = [\n",
    "    output_dir + 'barcodes/' + filename \n",
    "    for filename in os.listdir(output_dir + 'barcodes/') if r_file.split('/')[-1] == filename.split('_part')[0]\n",
    "]\n",
    "adj_filenames = sorted(adj_filenames, key = lambda x: int(x.split('_')[-1].split('of')[0][4:].strip())) \n",
    "adj_filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reformat_barcodes(barcodes):\n",
    "    \"\"\"Return barcodes to their original format\"\"\"\n",
    "    formatted_barcodes = []\n",
    "    for barcode in barcodes:\n",
    "        formatted_barcode = {}\n",
    "        for dim in barcode:\n",
    "            formatted_barcode[int(dim)] = np.asarray(\n",
    "                [(b, d) for b,d in barcode[dim]], dtype=[('birth', '<f4'), ('death', '<f4')]\n",
    "            )\n",
    "        formatted_barcodes.append(formatted_barcode)\n",
    "    return formatted_barcodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ad37fa4c10a2480db1659d4ce208e72b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Calculating ripser++ features', max=5.0, style=ProgressSt…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Barcodes loaded from: small_gpt_web/barcodes/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part1of5.json\n",
      "Barcodes loaded from: small_gpt_web/barcodes/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part2of5.json\n",
      "Barcodes loaded from: small_gpt_web/barcodes/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part3of5.json\n",
      "Barcodes loaded from: small_gpt_web/barcodes/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part4of5.json\n",
      "Barcodes loaded from: small_gpt_web/barcodes/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part5of5.json\n",
      "\n"
     ]
    }
   ],
   "source": [
    "features_array = []\n",
    "\n",
    "for filename in tqdm(adj_filenames, desc='Calculating ripser++ features'):\n",
    "    barcodes = json.load(open(filename))\n",
    "    print(f\"Barcodes loaded from: {filename}\", flush=True)\n",
    "    features_part = []\n",
    "    for layer in barcodes:\n",
    "        features_layer = []\n",
    "        for head in barcodes[layer]:\n",
    "            ref_barcodes = reformat_barcodes(barcodes[layer][head])\n",
    "            features = ripser_count.count_ripser_features(ref_barcodes, ripser_feature_names)\n",
    "            features_layer.append(features)\n",
    "        features_part.append(features_layer)\n",
    "    features_array.append(np.asarray(features_part))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'small_gpt_web/features/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_ripser.npy'"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ripser_file = output_dir + 'features/' + subset + \"_all_heads_\" + str(len(layers_of_interest)) + \"_layers\" \\\n",
    "             + \"_MAX_LEN_\" + str(max_tokens_amount) + \\\n",
    "             \"_\" + model_path.split(\"/\")[-1] + \"_ripser\" + '.npy'\n",
    "ripser_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(12, 12, 5000, 14)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "features = np.concatenate(features_array, axis=2)\n",
    "features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(ripser_file, features)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculating template features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from multiprocessing import Pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained(\"bert-base-uncased\", do_lower_case=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def matrix_distance(matricies, template, broadcast=True):\n",
    "    \"\"\"\n",
    "    Calculates the distance between the list of matricies and the template matrix.\n",
    "    Args:\n",
    "    \n",
    "    -- matricies: np.array of shape (n_matricies, dim, dim)\n",
    "    -- template: np.array of shape (dim, dim) if broadcast else (n_matricies, dim, dim)\n",
    "    \n",
    "    Returns:\n",
    "    -- diff: np.array of shape (n_matricies, )\n",
    "    \"\"\"\n",
    "    diff = np.linalg.norm(matricies-template, ord='fro', axis=(1, 2))\n",
    "    div = np.linalg.norm(matricies, ord='fro', axis=(1, 2))**2\n",
    "    if broadcast:\n",
    "        div += np.linalg.norm(template, ord='fro')**2\n",
    "    else:\n",
    "        div += np.linalg.norm(template, ord='fro', axis=(1, 2))**2\n",
    "    return diff/np.sqrt(div)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "attention_dir = 'small_gpt_web/attentions/'\n",
    "attention_name = 'test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased'\n",
    "\n",
    "texts_name = 'small_gpt_web/test_5k.csv'\n",
    "\n",
    "MAX_LEN = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "def attention_to_self(matricies):\n",
    "    \"\"\"\n",
    "    Calculates the distance between input matricies and identity matrix, \n",
    "    which representes the attention to the same token.\n",
    "    \"\"\"\n",
    "    _, n, m = matricies.shape\n",
    "    assert n == m, f\"Input matrix has shape {n} x {m}, but the square matrix is expected\"\n",
    "    template_matrix = np.eye(n)\n",
    "    return matrix_distance(matricies, template_matrix)\n",
    "\n",
    "def attention_to_next_token(matricies):\n",
    "    \"\"\"\n",
    "    Calculates the distance between input and E=(i, i+1) matrix, \n",
    "    which representes the attention to the next token.\n",
    "    \"\"\"\n",
    "    _, n, m = matricies.shape\n",
    "    assert n == m, f\"Input matrix has shape {n} x {m}, but the square matrix is expected\"\n",
    "    template_matrix = np.triu(np.tri(n, k=1, dtype=matricies.dtype), k=1)\n",
    "    return matrix_distance(matricies, template_matrix)\n",
    "\n",
    "def attention_to_prev_token(matricies):\n",
    "    \"\"\"\n",
    "    Calculates the distance between input and E=(i+1, i) matrix, \n",
    "    which representes the attention to the previous token.\n",
    "    \"\"\"\n",
    "    _, n, m = matricies.shape\n",
    "    assert n == m, f\"Input matrix has shape {n} x {m}, but the square matrix is expected\"\n",
    "    template_matrix = np.triu(np.tri(n, k=-1, dtype=matricies.dtype), k=-1)\n",
    "    return matrix_distance(matricies, template_matrix)\n",
    "\n",
    "def attention_to_beginning(matricies):\n",
    "    \"\"\"\n",
    "    Calculates the distance between input and E=(i+1, i) matrix, \n",
    "    which representes the attention to [CLS] token (beginning).\n",
    "    \"\"\"\n",
    "    _, n, m = matricies.shape\n",
    "    assert n == m, f\"Input matrix has shape {n} x {m}, but the square matrix is expected\"\n",
    "    template_matrix = np.zeros((n, n))\n",
    "    template_matrix[:, 0] = 1.0\n",
    "    return matrix_distance(matricies, template_matrix)\n",
    "\n",
    "def attention_to_ids(matricies, list_of_ids, token_id):\n",
    "    \"\"\"\n",
    "    Calculates the distance between input and ids matrix, \n",
    "    which representes the attention to some particular tokens,\n",
    "    which ids are in the `list_of_ids` (commas, periods, separators).\n",
    "    \"\"\"\n",
    "   \n",
    "    batch_size, n, m = matricies.shape\n",
    "    EPS = 1e-7\n",
    "    assert n == m, f\"Input matrix has shape {n} x {m}, but the square matrix is expected\"\n",
    "#     assert len(list_of_ids) == batch_size, f\"List of ids length doesn't match the dimension of the matrix\"\n",
    "    template_matrix = np.zeros_like(matricies)\n",
    "    ids = np.argwhere(list_of_ids == token_id)\n",
    "    if len(ids):\n",
    "        batch_ids, row_ids = zip(*ids)\n",
    "        template_matrix[np.array(batch_ids), :, np.array(row_ids)] = 1.0\n",
    "        template_matrix /= (np.sum(template_matrix, axis=-1, keepdims=True) + EPS)\n",
    "    return matrix_distance(matricies, template_matrix, broadcast=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_template_features(matricies, feature_list=['self', 'beginning', 'prev', 'next', 'comma', 'dot'], ids=None):\n",
    "    features = []\n",
    "    comma_id = 1010\n",
    "    dot_id = 1012\n",
    "    for feature in feature_list:\n",
    "        if feature == 'self':\n",
    "            features.append(attention_to_self(matricies))\n",
    "        elif feature == 'beginning':\n",
    "            features.append(attention_to_beginning(matricies))\n",
    "        elif feature == 'prev':\n",
    "            features.append(attention_to_prev_token(matricies))\n",
    "        elif feature == 'next':\n",
    "            features.append(attention_to_next_token(matricies))\n",
    "        elif feature == 'comma':\n",
    "            features.append(attention_to_ids(matricies, ids, comma_id))\n",
    "        elif feature == 'dot':\n",
    "            features.append(attention_to_ids(matricies, ids, dot_id))\n",
    "    return np.array(features)\n",
    "\n",
    "def calculate_features_t(adj_matricies, template_features, ids=None):\n",
    "    \"\"\"Calculate template features for adj_matricies\"\"\"\n",
    "    features = []\n",
    "    for layer in range(adj_matricies.shape[1]):\n",
    "        features.append([])\n",
    "        for head in range(adj_matricies.shape[2]):\n",
    "            matricies = adj_matricies[:, layer, head, :, :]\n",
    "            lh_features = count_template_features(matricies, template_features, ids) # samples X n_features\n",
    "            features[-1].append(lh_features)\n",
    "    return np.asarray(features) # layer X head X n_features X samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "# '.' id 1012\n",
    "# ',' id 1010\n",
    "def get_list_of_ids(sentences, tokenizer):\n",
    "    inputs = tokenizer.batch_encode_plus([text_preprocessing(s) for s in sentences],\n",
    "                                       add_special_tokens=True,\n",
    "                                       max_length=MAX_LEN,             # Max length to truncate/pad\n",
    "                                       pad_to_max_length=True,         # Pad sentence to max length)\n",
    "                                       truncation=True\n",
    "                                      )\n",
    "    return np.array(inputs['input_ids'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_of_workers = 20 \n",
    "pool = Pool(num_of_workers)\n",
    "feature_list = ['self', 'beginning', 'prev', 'next', 'comma', 'dot']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part1of5.npy',\n",
       " 'small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part2of5.npy',\n",
       " 'small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part3of5.npy',\n",
       " 'small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part4of5.npy',\n",
       " 'small_gpt_web/attentions/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_part5of5.npy']"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adj_filenames = [\n",
    "    attention_dir + filename \n",
    "    for filename in os.listdir(attention_dir) \n",
    "    if attention_name == filename.split(\"_part\")[0]\n",
    "]\n",
    "# sorted by part number\n",
    "adj_filenames = sorted(adj_filenames, key = lambda x: int(x.split('_')[-1].split('of')[0][4:].strip())) \n",
    "adj_filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased'"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attention_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = pd.read_csv(texts_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ce7f8566dd76435c902cf850f77a2603",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Features calc', max=5.0, style=ProgressStyle(description_…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1f802967dfe342edb56245c2c98150c1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Calculating token ids on iter 0 from 5', max=20.0, style=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1ccd46d9f04543ea832f5d123fe4e51e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Calculating token ids on iter 1 from 5', max=20.0, style=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "781880442c194c71b9896f8470607593",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Calculating token ids on iter 2 from 5', max=20.0, style=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a73ed6a58e014436bab1b8382eb4e9d3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Calculating token ids on iter 3 from 5', max=20.0, style=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2dd92288d1be466e8272fe6aba2a2f96",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Calculating token ids on iter 4 from 5', max=20.0, style=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "features_array = []\n",
    "\n",
    "for i, filename in tqdm(list(enumerate(adj_filenames)), desc='Features calc'):\n",
    "    adj_matricies = np.load(filename, allow_pickle=True)\n",
    "    batch_size = adj_matricies.shape[0]\n",
    "    sentences = texts['sentence'].values[i*batch_size:(i+1)*batch_size]\n",
    "    splitted_indexes = np.array_split(np.arange(batch_size), num_of_workers)\n",
    "    splitted_list_of_ids = [\n",
    "        get_list_of_ids(sentences[indx], tokenizer) \n",
    "        for indx in tqdm(splitted_indexes, desc=f\"Calculating token ids on iter {i} from {len(adj_filenames)}\")\n",
    "    ]\n",
    "    splitted_adj_matricies = [adj_matricies[indx] for indx in splitted_indexes]\n",
    "    \n",
    "    args = [(m, feature_list, list_of_ids) for m, list_of_ids in zip(splitted_adj_matricies, splitted_list_of_ids)]\n",
    "    \n",
    "    features_array_part = pool.starmap(\n",
    "        calculate_features_t, args\n",
    "    )\n",
    "    features_array.append(np.concatenate([_ for _ in features_array_part], axis=3))\n",
    "features_array = np.concatenate(features_array, axis=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'small_gpt_web/features/test_5k_all_heads_12_layers_MAX_LEN_128_bert-base-uncased_template.npy'"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"small_gpt_web/features/\" + attention_name + \"_template.npy\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(\"small_gpt_web/features/\" + attention_name + \"_template.npy\", features_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}