2_1_Topological_and_template_features_calculation.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Notebook for topological and template features calculation.\n",
        "Based on [this code](https://github.com/danchern97/tda4atd/blob/main/features_calculation_by_thresholds.ipynb).   \n",
        "**Running this code does not require GPU, however, execution time heavily depends on the availiable num_of_workers.**   \n",
        "\n",
        "Execution info (for calculating features on ```en-dev``` subset (527 entities)):\n",
        " * Average feature calculation time given **8** num_of_workers: ~4 minutes for topological features and ~1 minute for template features.  \n",
        " * Storing the features requires ~12Mb of free memory."
      ],
      "metadata": {
        "id": "1VRAl4FTmS8y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## Uncomment if running in colab\n",
        "# from google.colab import drive\n",
        "# drive.mount('/content/gdrive')\n",
        "# %cd /content/gdrive/My Drive/test-la-tda\n",
        "# %cd la-tda\n",
        "# !pip install -r requirements.txt -q"
      ],
      "metadata": {
        "id": "hymVFQRPexut"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gn6BEg5TevK5"
      },
      "source": [
        "# Topological features"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "VgzystbCevK6"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoModelForSequenceClassification,AutoTokenizer\n",
        "from multiprocessing import Pool\n",
        "import os\n",
        "from time import sleep\n",
        "from tqdm.notebook import tqdm\n",
        "import logging\n",
        "import pandas as pd\n",
        "from math import ceil\n",
        "import numpy as np\n",
        "import torch\n",
        "import gzip\n",
        "import networkx as nx\n",
        "from pathlib import Path"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "8I-McQ02evK7"
      },
      "outputs": [],
      "source": [
        "def cutoff_matrix(matrix, ntokens):\n",
        "    \"\"\"Return normalized submatrix of first n_tokens\"\"\"\n",
        "    matrix = matrix[:ntokens, :ntokens]\n",
        "    matrix /= matrix.sum(axis=1, keepdims=True)\n",
        "    return matrix"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "0ixCuTWeevK7"
      },
      "outputs": [],
      "source": [
        "def get_filtered_mat_list(adj_matrix, thresholds_array, ntokens):\n",
        "    \"\"\"\n",
        "    Converts adjacency matrix with real weights into list of binary matrices.\n",
        "    For each threshold, those weights of adjacency matrix, which are less than\n",
        "    threshold, get \"filtered out\" (set to 0), remained weights are set to ones.\n",
        "\n",
        "    Args:\n",
        "        adj_matrix (np.array[float, float])\n",
        "        thresholds_array (iterable[float])\n",
        "        ntokens (int)\n",
        "\n",
        "    Returns:\n",
        "        filtered_matricies (list[np.array[int, int]])\n",
        "    \"\"\"\n",
        "    filtered_matrices = []\n",
        "    for thr in thresholds_array:\n",
        "        filtered_matrix = adj_matrix.copy()\n",
        "        filtered_matrix = cutoff_matrix(filtered_matrix, ntokens)\n",
        "        filtered_matrix[filtered_matrix < thr] = 0\n",
        "        filtered_matrix[filtered_matrix >= thr] = 1\n",
        "        filtered_matrices.append(filtered_matrix.astype(np.int8))\n",
        "    return filtered_matrices\n",
        "\n",
        "\n",
        "def adj_m_to_nx_list(adj_matrix, thresholds_array, ntokens, no_mat_output=False):\n",
        "    \"\"\"\n",
        "    Converts adjacency matrix into list of unweighted digraphs, using filtering\n",
        "    process from previous function.\n",
        "\n",
        "    Args:\n",
        "        adj_matrix (np.array[float, float])\n",
        "        thresholds_array (iterable[float])\n",
        "        ntokens (int)\n",
        "\n",
        "    Returns:\n",
        "        nx_graphs_list (list[nx.MultiDiGraph])\n",
        "        filt_mat_list(list[np.array[int, int]])\n",
        "\n",
        "    \"\"\"\n",
        "    #     adj_matrix = adj_matrix[:length,:length]\n",
        "    filt_mat_list = get_filtered_mat_list(adj_matrix, thresholds_array, ntokens)\n",
        "    nx_graphs_list = []\n",
        "    for mat in filt_mat_list:\n",
        "        nx_graphs_list.append(nx.from_numpy_matrix(np.array(mat),\n",
        "                                                   create_using=nx.MultiDiGraph()))\n",
        "    if no_mat_output:\n",
        "        return nx_graphs_list, []\n",
        "    else:\n",
        "        return nx_graphs_list, filt_mat_list\n",
        "\n",
        "\n",
        "def adj_ms_to_nx_lists(adj_matricies, \\\n",
        "                       thresholds_array, \\\n",
        "                       ntokens_array, \\\n",
        "                       verbose=False, \\\n",
        "                       no_mat_output=False):\n",
        "    \"\"\"\n",
        "    Executes adj_m_to_nx_list for each matrix in adj_matricies array, arranges\n",
        "    the results. If verbose==True, shows progress bar.\n",
        "\n",
        "    Args:\n",
        "        adj_matricies (np.array[float, float])\n",
        "        thresholds_array (iterable[float])\n",
        "        verbose (bool)\n",
        "\n",
        "    Returns:\n",
        "        nx_graphs_list (list[nx.MultiDiGraph])\n",
        "        filt_mat_lists (list[list[np.array[int,int]]])\n",
        "    \"\"\"\n",
        "    graph_lists = []\n",
        "    filt_mat_lists = []\n",
        "\n",
        "    iterable = range(len(adj_matricies))\n",
        "    if verbose:\n",
        "        iterable = tqdm(range(len(adj_matricies)),\n",
        "                        desc=\"Calc graphs list\")\n",
        "    for i in iterable:\n",
        "        g_list, filt_mat_list = adj_m_to_nx_list(adj_matricies[i], \\\n",
        "                                                 thresholds_array, \\\n",
        "                                                 ntokens_array[i], \\\n",
        "                                                 no_mat_output=no_mat_output)\n",
        "        graph_lists.append(g_list)\n",
        "        filt_mat_lists.append(filt_mat_lists)\n",
        "\n",
        "    return graph_lists, filt_mat_lists\n",
        "\n",
        "\n",
        "def count_stat(g_listt_j, function=nx.weakly_connected_components, cap=500):\n",
        "    stat_amount = 0\n",
        "    for _ in function(g_listt_j):\n",
        "        stat_amount += 1\n",
        "        if stat_amount >= cap:\n",
        "            break\n",
        "    return stat_amount\n",
        "\n",
        "\n",
        "def count_weak_components(g_listt_j, cap=500):\n",
        "    return count_stat(g_listt_j, function=nx.weakly_connected_components, cap=cap)\n",
        "\n",
        "\n",
        "def count_strong_components(g_listt_j, cap=500):\n",
        "    return count_stat(g_listt_j, function=nx.strongly_connected_components, cap=cap)\n",
        "\n",
        "\n",
        "def count_simple_cycles(g_listt_j, cap=500):\n",
        "    return count_stat(g_listt_j, function=nx.simple_cycles, cap=cap)\n",
        "\n",
        "\n",
        "def count_b1(g_listt_j, cap=500):\n",
        "    return count_stat(g_listt_j, function=nx.cycle_basis, cap=cap)\n",
        "\n",
        "\n",
        "def dim_connected_components(graph_lists, strong=False, verbose=False, cap=500):\n",
        "    \"\"\"\n",
        "    Calculates number of connected components for each graph in list\n",
        "    of lists of digraphs. If strong==True, calculates strongly connected\n",
        "    components, otherwise calculates weakly connected components.\n",
        "    If verbose==True, shows progress bar.\n",
        "\n",
        "    Args:\n",
        "        graph_lists (list[list[nx.MultiDiGraph]])\n",
        "        strong (bool)\n",
        "        verbose (bool)\n",
        "\n",
        "    Returns:\n",
        "        w_lists (list[list[int])\n",
        "    \"\"\"\n",
        "    w_lists = []  # len == len(w_graph_lists)\n",
        "    iterable = range(len(graph_lists))\n",
        "    if verbose:\n",
        "        iterable = tqdm(range(len(graph_lists)),\n",
        "                        desc=\"Calc weak comp\")\n",
        "    for i in iterable:\n",
        "        g_list = graph_lists[i]\n",
        "        w_cmp = []\n",
        "        for j in range(len(g_list)):\n",
        "            if strong:\n",
        "                w_cmp.append(count_strong_components(g_list[j], cap=cap))\n",
        "            else:\n",
        "                w_cmp.append(count_weak_components(g_list[j], cap=cap))\n",
        "        w_lists.append(w_cmp)\n",
        "    return w_lists\n",
        "\n",
        "\n",
        "def dim_simple_cycles(graph_lists, verbose, cap=500):\n",
        "    \"\"\"\n",
        "    Calculates number of simple cycles for each graph in list\n",
        "    of lists of digraphs. If verbose==True, shows progress bar.\n",
        "\n",
        "    Args:\n",
        "        graph_lists (list[list[nx.MultiDiGraph]])\n",
        "        verbose (bool)\n",
        "\n",
        "    Returns:\n",
        "        c_lists (list[list[int])\n",
        "    \"\"\"\n",
        "    c_lists = []  # len == len(pos_w_graph_lists)\n",
        "    iterable = range(len(graph_lists))\n",
        "    if verbose:\n",
        "        iterable = tqdm(range(len(graph_lists)),\n",
        "                        desc=\"Calc cycles\")\n",
        "    for i in iterable:\n",
        "        g_list = graph_lists[i]\n",
        "        c = []\n",
        "        for j in range(len(g_list)):\n",
        "            c.append(count_simple_cycles(g_list[j], cap=cap))\n",
        "        c_lists.append(c)\n",
        "    if verbose:\n",
        "        logger = logging.getLogger()\n",
        "        flat = [x for l in c_lists for x in l]\n",
        "        logger.debug('dim_simple_cycles: min=%f mean=%f max=%f,  ratio of cap = %d / %d' %\n",
        "                     (\n",
        "                         np.min(c_lists), np.mean(c_lists), np.max(c_lists), len([x for x in flat if x == cap]),\n",
        "                         len(flat)))\n",
        "\n",
        "    return c_lists\n",
        "\n",
        "\n",
        "def dim_b1(graph_lists, verbose, cap=500):\n",
        "    b1_lists = []  # len == len(pos_w_graph_lists)\n",
        "    iterable = range(len(graph_lists))\n",
        "    if verbose:\n",
        "        iterable = tqdm(range(len(graph_lists)),\n",
        "                        desc=\"Calc b1 (undirected graphs)\")\n",
        "    for i in iterable:\n",
        "        g_list = graph_lists[i]\n",
        "        b1 = []\n",
        "        for j in range(len(g_list)):\n",
        "            b1.append(count_b1(nx.Graph(g_list[j].to_undirected()), cap=cap))\n",
        "        b1_lists.append(b1)\n",
        "    return b1_lists\n",
        "\n",
        "\n",
        "def b0_b1(graph_lists, verbose):\n",
        "    b0_lists = []\n",
        "    b1_lists = []  # len == len(pos_w_graph_lists)\n",
        "    iterable = range(len(graph_lists))\n",
        "    if verbose:\n",
        "        iterable = tqdm(range(len(graph_lists)),\n",
        "                        desc=\"Calc b0, b1\")\n",
        "    for i in iterable:\n",
        "        g_list = graph_lists[i]\n",
        "        b0 = []\n",
        "        b1 = []\n",
        "        for j in range(len(g_list)):\n",
        "            g = nx.Graph(g_list[j].to_undirected())\n",
        "            w = nx.number_connected_components(g)\n",
        "            e = g.number_of_edges()\n",
        "            v = g.number_of_nodes()\n",
        "            b0.append(w)\n",
        "            b1.append(e - v + w)\n",
        "        b0_lists.append(b0)\n",
        "        b1_lists.append(b1)\n",
        "    return b0_lists, b1_lists\n",
        "\n",
        "\n",
        "def edges_f(graph_lists, verbose):\n",
        "    \"\"\"\n",
        "    Calculates number of edges for each graph in list\n",
        "    of lists of digraphs. If verbose==True, shows progress bar.\n",
        "\n",
        "    Args:\n",
        "        graph_lists (list[list[nx.MultiDiGraph]])\n",
        "        verbose (bool)\n",
        "\n",
        "    Returns:\n",
        "        e_lists (list[list[int])\n",
        "    \"\"\"\n",
        "    e_lists = []  # len == len(pos_w_graph_lists)\n",
        "    iterable = range(len(graph_lists))\n",
        "    if verbose > 2:\n",
        "        iterable = tqdm(range(len(graph_lists)),\n",
        "                        desc=\"Calc edges number\")\n",
        "    for i in iterable:\n",
        "        g_list = graph_lists[i]\n",
        "        e = []\n",
        "        for j in range(len(g_list)):\n",
        "            e.append(g_list[j].number_of_edges())\n",
        "        e_lists.append(e)\n",
        "    return e_lists\n",
        "\n",
        "\n",
        "def v_degree_f(graph_lists, verbose):\n",
        "    \"\"\"\n",
        "    Calculates number of edges for each graph in list\n",
        "    of lists of digraphs. If verbose==True, shows progress bar.\n",
        "\n",
        "    Args:\n",
        "        graph_lists (list[list[nx.MultiDiGraph]])\n",
        "        verbose (bool)\n",
        "\n",
        "    Returns:\n",
        "        v_lists (list[list[int])\n",
        "    \"\"\"\n",
        "    v_lists = []  # len == len(pos_w_graph_lists)\n",
        "    iterable = range(len(graph_lists))\n",
        "    if verbose > 2:\n",
        "        iterable = tqdm(range(len(graph_lists)),\n",
        "                        desc=\"Calc average vertex degree\")\n",
        "    for i in iterable:\n",
        "        g_list = graph_lists[i]\n",
        "        v = []\n",
        "        for j in range(len(g_list)):\n",
        "            degrees = g_list[j].degree()\n",
        "            degree_values = [v for k, v in degrees]\n",
        "            sum_of_edges = sum(degree_values) / float(len(degree_values))\n",
        "            v.append(sum_of_edges)\n",
        "        v_lists.append(v)\n",
        "    return v_lists\n",
        "\n",
        "\n",
        "def chordality_f(graph_lists, verbose):\n",
        "    \"\"\"\n",
        "    Checks whether the graph is chordal or not for each graph in list of lists of graphs.\n",
        "    If verbose==True, shows progress bar.\n",
        "\n",
        "    Args:\n",
        "        graph_lists (list[list[nx.MultiDiGraph]])\n",
        "        verbose (bool)\n",
        "\n",
        "    Returns:\n",
        "        ch_lists (list[list[int])\n",
        "    \"\"\"\n",
        "    ch_lists = []  # len == len(pos_w_graph_lists)\n",
        "    iterable = range(len(graph_lists))\n",
        "    for i in iterable:\n",
        "        g_list = graph_lists[i]\n",
        "        ch = []\n",
        "        for j in range(len(g_list)):\n",
        "            g = nx.Graph(g_list[j].to_undirected())\n",
        "            # print(g)\n",
        "            # print(g.edges())\n",
        "            g.remove_edges_from(nx.selfloop_edges(g))\n",
        "            ch_i = nx.is_chordal(g)\n",
        "            ch.append(int(ch_i))\n",
        "        ch_lists.append(ch)\n",
        "    return ch_lists\n",
        "def max_matching_f(graph_lists, verbose):\n",
        "    \"\"\"\n",
        "    Calculates max matching size for each graph in list\n",
        "    of lists of graphs. If verbose==True, shows progress bar.\n",
        "\n",
        "    Args:\n",
        "        graph_lists (list[list[nx.MultiDiGraph]])\n",
        "        verbose (bool)\n",
        "\n",
        "    Returns:\n",
        "        max_m_lists (list[list[int])\n",
        "    \"\"\"\n",
        "    max_m_lists = []  # len == len(pos_w_graph_lists)\n",
        "    iterable = range(len(graph_lists))\n",
        "    if verbose > 2:\n",
        "        iterable = tqdm(range(len(graph_lists)),\n",
        "                        desc=\"Calc max matching of the graph\")\n",
        "    for i in iterable:\n",
        "        g_list = graph_lists[i]\n",
        "        max_m = []\n",
        "        for j in range(len(g_list)):\n",
        "            g = nx.Graph(g_list[j].to_undirected())\n",
        "            m_i = nx.maximal_matching(g)\n",
        "            max_m.append(len(m_i))\n",
        "        max_m_lists.append(max_m)\n",
        "    return max_m_lists\n",
        "\n",
        "def count_top_stats(adj_matricies,\n",
        "                    thresholds_array,\n",
        "                    ntokens_array,\n",
        "                    stats_to_count={\"s\", \"e\", \"c\", \"v\", \"b0b1\"},\n",
        "                    stats_cap=500,\n",
        "                    verbose=False):\n",
        "    \"\"\"\n",
        "    The main function for calculating topological invariants. Unites the\n",
        "    functional of all functions above.\n",
        "    Args:\n",
        "        adj_matricies (np.array[float, float, float, float, float])\n",
        "        thresholds_array (list[float])\n",
        "        stats_to_count (str)\n",
        "        stats_cap (int)\n",
        "        verbose (bool)\n",
        "    Returns:\n",
        "        stats_tuple_lists_array (np.array[float, float, float, float, float])\n",
        "    \"\"\"\n",
        "    stats_tuple_lists_array = []\n",
        "\n",
        "    for layer_of_interest in tqdm(range(adj_matricies.shape[1])):\n",
        "        stats_tuple_lists_array.append([])\n",
        "        for head_of_interest in range(adj_matricies.shape[2]):\n",
        "            adj_ms = adj_matricies[:, layer_of_interest, head_of_interest, :, :]\n",
        "            g_lists, _ = adj_ms_to_nx_lists(adj_ms,\n",
        "                                            thresholds_array=thresholds_array,\n",
        "                                            ntokens_array=ntokens_array,\n",
        "                                            verbose=verbose)\n",
        "            feat_lists = []\n",
        "            if \"s\" in stats_to_count:\n",
        "                feat_lists.append(dim_connected_components(g_lists,\n",
        "                                                           strong=True,\n",
        "                                                           verbose=verbose,\n",
        "                                                           cap=stats_cap))\n",
        "            if \"w\" in stats_to_count:\n",
        "                feat_lists.append(dim_connected_components(g_lists,\n",
        "                                                           strong=False,\n",
        "                                                           verbose=verbose,\n",
        "                                                           cap=stats_cap))\n",
        "            if \"e\" in stats_to_count:\n",
        "                feat_lists.append(edges_f(g_lists, verbose=verbose))\n",
        "            if \"v\" in stats_to_count:\n",
        "                feat_lists.append(v_degree_f(g_lists, verbose=verbose))\n",
        "            if \"c\" in stats_to_count:\n",
        "                feat_lists.append(dim_simple_cycles(g_lists,\n",
        "                                                    verbose=verbose,\n",
        "                                                    cap=50))\n",
        "\n",
        "            if \"b0b1\" in stats_to_count:\n",
        "                b0_lists, b1_lists = b0_b1(g_lists, verbose=verbose)\n",
        "                feat_lists.append(b0_lists)\n",
        "                feat_lists.append(b1_lists)\n",
        "            if \"m\" in stats_to_count:\n",
        "                feat_lists.append(max_matching_f(g_lists,\n",
        "                                                    verbose=verbose))\n",
        "            if \"k\" in stats_to_count:\n",
        "                feat_lists.append(chordality_f(g_lists,\n",
        "                                                    verbose=verbose))\n",
        "            stats_tuple_lists_array[-1].append(tuple(feat_lists))\n",
        "\n",
        "    stats_tuple_lists_array = np.asarray(stats_tuple_lists_array,\n",
        "                                         dtype=np.float16)\n",
        "    return stats_tuple_lists_array\n",
        "\n",
        "\n",
        "def function_for_v(list_of_v_degrees_of_graph):\n",
        "    return sum(map(lambda x: np.sqrt(x * x), list_of_v_degrees_of_graph))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "jMnrS4q4SkqF"
      },
      "outputs": [],
      "source": [
        "subset = \"dev\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "73YlQ-nmmzwt"
      },
      "outputs": [],
      "source": [
        "model_path = \"./bert-base-cased-en-cola_32_3e-05_lr_0.01_decay_balanced/\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "z-3gZpVz2lxX"
      },
      "outputs": [],
      "source": [
        "input_dir = output_dir = './'\n",
        "layers_of_interest = list(range(24))\n",
        "data = pd.read_csv(\"./data/en-cola/\" + subset + '.csv')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "s5f1AsPn2rmp"
      },
      "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": 9,
      "metadata": {
        "id": "-PmbEU-jXsZc"
      },
      "outputs": [],
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
        "max_tokens_amount  = 64\n",
        "MAX_LEN = max_tokens_amount"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "KFtjQMDmXrd5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b41c35fc-7ae8-41d2-cb16-300fdc50350d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/transformers/tokenization_utils_base.py:2302: FutureWarning: The `pad_to_max_length` argument is deprecated and will be removed in a future version, use `padding=True` or `padding='longest'` to pad to the longest sequence in the batch, or use `padding='max_length'` to pad to a max length. In this case, you can give a specific length with `max_length` (e.g. `max_length=45`) or leave max_length to None to pad to the maximal input size of the model (e.g. 512 for Bert).\n",
            "  warnings.warn(\n"
          ]
        }
      ],
      "source": [
        "data['tokenizer_length'] = get_token_length(list(data['sentence'].values))\n",
        "ntokens_array = data['tokenizer_length'].values"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "hMjuDds1zndK",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 36
        },
        "outputId": "e93d4275-fcf9-4e75-ee87-75585865bf78"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'./bert-base-cased-en-cola_32_3e-05_lr_0.01_decay_balanced/features/devs_e_v_c_b0b1_m_k_array_6.npy'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 11
        }
      ],
      "source": [
        "stats_name = \"s_e_v_c_b0b1_m_k\"\n",
        "thresholds_array = [0.025, 0.05, 0.1, 0.25, 0.5, 0.75]\n",
        "thrs = len(thresholds_array)\n",
        "layers_of_interest = [i for i in range(12)]\n",
        "stats_file = model_path + 'features/' + subset + \"\" + stats_name + \"_array_\" + str(thrs) + '.npy'\n",
        "stats_file"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "zC4XVIYKevK-"
      },
      "outputs": [],
      "source": [
        "def order_files(path, subset):\n",
        "    files_path = Path(path)\n",
        "    files = list(filter(lambda y: (y.is_file() and subset in str(y)), files_path.iterdir()))\n",
        "    files = [str(_) for _ in files]\n",
        "    files = sorted(files, key=lambda x: int(x.split('_')[-1].split('of')[0][4:].strip()))\n",
        "    return files"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OUzWlQ_JevK-",
        "outputId": "f55c9a54-e034-4ba4-bed2-de4ffe32b2ba"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['bert-base-cased-en-cola_32_3e-05_lr_0.01_decay_balanced/attentions/dev_part1of1.npy']"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ],
      "source": [
        "output_dir=model_path\n",
        "attn_dir = model_path + \"/attentions/\"\n",
        "adj_filenames = order_files(path=attn_dir, subset=subset)\n",
        "adj_filenames"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "Y122PsBY0NVn"
      },
      "outputs": [],
      "source": [
        "def split_matricies_and_lengths(adj_matricies, ntokens_array, num_of_workers):\n",
        "    splitted_adj_matricies = np.array_split(adj_matricies, num_of_workers)\n",
        "    splitted_ntokens = np.array_split(ntokens_array, num_of_workers)\n",
        "    assert all([len(m)==len(n) for m, n in zip(splitted_adj_matricies, splitted_ntokens)]), \"Split is not valid!\"\n",
        "    return zip(splitted_adj_matricies, splitted_ntokens)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "pMVNj_LX0RCX",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e43a2942-c134-4ce0-9701-02e6347573c1"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "8"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ],
      "source": [
        "num_of_workers = os.cpu_count()\n",
        "pool = Pool(num_of_workers)\n",
        "num_of_workers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "xqGXgFxC0Scf"
      },
      "outputs": [],
      "source": [
        "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\n",
        "batched_sentences = np.array_split(data['sentence'].values, number_of_batches)\n",
        "number_of_files = ceil(number_of_batches / DUMP_SIZE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "iIzGtPKNevK_"
      },
      "outputs": [],
      "source": [
        "debug=False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "mYmEU9SL0UwI",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "d393dc43207b4e7cbdf4fe6cbc454983",
            "ba16b2d9d4424f3088d921621fa8e006",
            "403377afd6e343258fa2744653ed88a9",
            "e432ce298758494aa5251feb5a3b4a44",
            "1e94f2f586ea431caea6afeb21a4ff21",
            "15e456956579490e9b3bceda282c7aab",
            "b7bfcdefb3374faeb9874dfe1a41ab22",
            "d919d4f2a46d430f97a0c0b3c2294226",
            "13a32f41fa294472bad73a99186b2d9b",
            "895206c698af486094e0caba08568f41",
            "ca553fc7e7d041e9903cd8d8f20849ce"
          ]
        },
        "outputId": "984b8b2c-09d2-4e24-bd9e-ed479a5278c1"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Calculating topological features:   0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "d393dc43207b4e7cbdf4fe6cbc454983"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "stats_name = \"s_e_v_c_b0b1_m_k\" # \"c\"\n",
        "stats_cap = 500\n",
        "stats_tuple_lists_array = []\n",
        "for i, filename in enumerate(tqdm(adj_filenames,\n",
        "                                  desc='Calculating topological features')):\n",
        "    if \"gz\" in filename:\n",
        "        with gzip.GzipFile(filename, 'rb') as f:\n",
        "            adj_matricies = np.load(f, allow_pickle=True)\n",
        "    else:\n",
        "        with open(filename, 'rb') as f:\n",
        "            adj_matricies = np.load(f, allow_pickle=True)\n",
        "    ntokens = ntokens_array[i*batch_size*DUMP_SIZE : (i+1)*batch_size*DUMP_SIZE]\n",
        "    if debug:\n",
        "        stats_tuple_lists_array_part = count_top_stats(adj_matricies, thresholds_array, ntokens, stats_name.split(\"_\"), stats_cap, verbose=2)\n",
        "    else:\n",
        "        splitted = split_matricies_and_lengths(adj_matricies, ntokens, num_of_workers)\n",
        "        args = [(m, thresholds_array, ntokens, stats_name.split(\"_\"), stats_cap) for m, ntokens in splitted]\n",
        "        stats_tuple_lists_array_part = pool.starmap(\n",
        "            count_top_stats, args\n",
        "        )\n",
        "    stats_tuple_lists_array.append(np.concatenate([_ for _ in stats_tuple_lists_array_part], axis=3))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "CG3s3Vaw0d5w"
      },
      "outputs": [],
      "source": [
        "stats_tuple_lists_array = np.concatenate(stats_tuple_lists_array, axis=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yvcgp6TU0sdX"
      },
      "outputs": [],
      "source": [
        "%cd $output_dir\n",
        "!mkdir features"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd .."
      ],
      "metadata": {
        "id": "R0-3kkz3n8bK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "id": "8kMbYKkB0tw_"
      },
      "outputs": [],
      "source": [
        "np.save(stats_file, np.concatenate(stats_tuple_lists_array, axis=3))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M_6aHjB5evK_"
      },
      "source": [
        "# Template features"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-zjwhV5gevK_",
        "outputId": "40b5fd7c-7a07-4425-cdef-0942afa91f73"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\rDownloading data files:   0% 0/1 [00:00<?, ?it/s]\rDownloading data files: 100% 1/1 [00:00<00:00, 2059.06it/s]\n",
            "\rExtracting data files:   0% 0/1 [00:00<?, ?it/s]\rExtracting data files: 100% 1/1 [00:00<00:00, 83.18it/s]\n",
            "Generating dev split: 527 examples [00:00, 15419.27 examples/s]\n",
            "Template Feature Calc: 100% 1/1 [00:44<00:00, 44.80s/it]\n"
          ]
        }
      ],
      "source": [
        "model_dir_ = \"./bert-base-cased-en-cola_32_3e-05_lr_0.01_decay_balanced/\"\n",
        "for data_subset in [\"dev\"]:\n",
        "    d_dir = f\"./data/en-cola/{data_subset}.csv\"\n",
        "    !PYTHONPATH=%PYTHONPATH% python -m src.template_features --model_dir $model_dir_ --data_file $d_dir --num_of_workers 8"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!du -c -h $output_dir"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6-5NIgiUoR9j",
        "outputId": "ce93a9ba-dd06-49a1-cc99-9e0b005eacb4"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "10K\t./bert-base-cased-en-cola_32_3e-05_lr_0.01_decay_balanced/runs/Oct13_09-52-01_2d788f34c1a2/1697190734.2923024\n",
            "19K\t./bert-base-cased-en-cola_32_3e-05_lr_0.01_decay_balanced/runs/Oct13_09-52-01_2d788f34c1a2\n",
            "23K\t./bert-base-cased-en-cola_32_3e-05_lr_0.01_decay_balanced/runs\n",
            "593M\t./bert-base-cased-en-cola_32_3e-05_lr_0.01_decay_balanced/attentions\n",
            "12M\t./bert-base-cased-en-cola_32_3e-05_lr_0.01_decay_balanced/features\n",
            "1019M\t./bert-base-cased-en-cola_32_3e-05_lr_0.01_decay_balanced/\n",
            "1019M\ttotal\n"
          ]
        }
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuClass": "premium"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.9.7"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "d393dc43207b4e7cbdf4fe6cbc454983": {
          "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_ba16b2d9d4424f3088d921621fa8e006",
              "IPY_MODEL_403377afd6e343258fa2744653ed88a9",
              "IPY_MODEL_e432ce298758494aa5251feb5a3b4a44"
            ],
            "layout": "IPY_MODEL_1e94f2f586ea431caea6afeb21a4ff21"
          }
        },
        "ba16b2d9d4424f3088d921621fa8e006": {
          "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_15e456956579490e9b3bceda282c7aab",
            "placeholder": "​",
            "style": "IPY_MODEL_b7bfcdefb3374faeb9874dfe1a41ab22",
            "value": "Calculating topological features: 100%"
          }
        },
        "403377afd6e343258fa2744653ed88a9": {
          "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_d919d4f2a46d430f97a0c0b3c2294226",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_13a32f41fa294472bad73a99186b2d9b",
            "value": 1
          }
        },
        "e432ce298758494aa5251feb5a3b4a44": {
          "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_895206c698af486094e0caba08568f41",
            "placeholder": "​",
            "style": "IPY_MODEL_ca553fc7e7d041e9903cd8d8f20849ce",
            "value": " 1/1 [04:43&lt;00:00, 283.12s/it]"
          }
        },
        "1e94f2f586ea431caea6afeb21a4ff21": {
          "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
          }
        },
        "15e456956579490e9b3bceda282c7aab": {
          "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
          }
        },
        "b7bfcdefb3374faeb9874dfe1a41ab22": {
          "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": ""
          }
        },
        "d919d4f2a46d430f97a0c0b3c2294226": {
          "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
          }
        },
        "13a32f41fa294472bad73a99186b2d9b": {
          "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": ""
          }
        },
        "895206c698af486094e0caba08568f41": {
          "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
          }
        },
        "ca553fc7e7d041e9903cd8d8f20849ce": {
          "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": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}