Demo-ColBERT.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Demo-ColBERT.ipynb",
      "provenance": [],
      "mount_file_id": "1W8L2AcLErd872wXDonK1Dtcrf3t-ZJAB",
      "authorship_tag": "ABX9TyOVIPkJvhQk93foKribftE8",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/finardi/tutos/blob/master/Demo-ColBERT.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6rEPRrmvJsC-"
      },
      "source": [
        "%%capture\n",
        "!pip install -q transformers"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oxHw8OZbQnhj"
      },
      "source": [
        "import gc\n",
        "import os\n",
        "import torch\n",
        "import pickle\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from functools import partial\n",
        "\n",
        "from transformers import logging\n",
        "from transformers import BertPreTrainedModel, BertModel, BertTokenizer\n",
        "\n",
        "logging.set_verbosity_error()\n",
        "\n",
        "# better pandas viz\n",
        "pd.set_option('display.max_columns', 100)  \n",
        "pd.set_option('display.expand_frame_repr', 100)\n",
        "pd.set_option('max_colwidth', 700)\n",
        "pd.set_option('display.max_rows', 5000)\n",
        "  \n",
        "# save/load pickles\n",
        "def pickle_file(path, data=None):\n",
        "    if data is None:\n",
        "        with open(path, 'rb') as f:\n",
        "            return pickle.load(f)\n",
        "    if data is not None:\n",
        "        with open(path, 'wb') as handle:\n",
        "            pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
        " \n",
        "# path base\n",
        "path_base = '/content/drive/MyDrive/BACEN/FAQ/'\n",
        "\n",
        "# load dataframes\n",
        "df_train = pd.read_parquet(path_base+'data/df_FAQ_TRAIN.parquet.gzip')\n",
        "\n",
        "print(f'unique docs:      {df_train[\"Doc\"].nunique()}')\n",
        "print(f'unique questions: {df_train[\"Query\"].nunique()}')\n",
        "df_triplet = pd.read_parquet(path_base+'data/df_FAQ_triplet_IDS_TRAIN.parquet.gzip')\n",
        "\n",
        "# load data dicts\n",
        "query_to_qid = pickle_file(path_base+'data/query_to_qid_TRAIN' )\n",
        "qid_to_query = pickle_file(path_base+'data/qid_to_query_TRAIN')\n",
        "doc_to_pid   = pickle_file(path_base+'data/doc_to_pid_TRAIN' )\n",
        "pid_to_doc   = pickle_file(path_base+'data/pid_to_doc_TRAIN')\n",
        "\n",
        "df_train.head()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vGpD2jlkWU3F"
      },
      "source": [
        "path_model = 'bert-base-multilingual-uncased'\n",
        "tok = BertTokenizer.from_pretrained(path_model)"
      ],
      "execution_count": 98,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wjOhd80YRJQw",
        "outputId": "9eb66af0-71af-489f-f732-67e6b169a338"
      },
      "source": [
        "sample = df_train.iloc[:6]\n",
        "queries = sample.Query.to_list()[:3]\n",
        "pos_docs = sample.Doc.to_list()[:3]\n",
        "neg_docs = sample.Doc.to_list()[3:]\n",
        "\n",
        "for i, (query, pos_doc, neg_doc) in enumerate(zip(queries, pos_docs, neg_docs)):\n",
        "    print(i, query)\n",
        "    print('\\t -> POSITIVE::: ', pos_doc)\n",
        "    print('\\t -> NEGATIVE::: ', neg_doc)"
      ],
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0 Quais as condições básicas para concessão dos créditos de investimento?\n",
            "\t -> POSITIVE:::   Os créditos de investimento devem ser concedidos mediante apresentação de projeto técnico, o qual poderá ser substituído, a critério da instituição financeira, por proposta simplificada de crédito, desde que as inversões programadas envolvam técnicas simples e bem assimiladas pelos agricultores da região ou se trate de crédito destinado à ampliação dos investimentos já financiados. Os créditos de investimento se destinam a promover o aumento da produção e da produtividade e a redução dos custos de produção, visando a elevação da renda da família produtora rural. Os créditos de investimento estão restritos ao financiamento de itens diretamente relacionados com a implantação, ampliação ou modernização da estrutura das atividades de produção, de armazenagem, de transporte ou de serviços agropecuários ou não agropecuários, no estabelecimento rural ou em áreas comunitárias rurais próximas, sendo passível de financiamento, ainda, a aquisição de equipamentos e de programas de informática voltados para melhoria da gestão dos empreendimentos rurais, de acordo com projetos técnicos específicos.\n",
            "\t -> NEGATIVE:::  A conta de pagamento é uma conta utilizada pelo cliente para a realização de pagamentos e transferência de recursos como, por exemplo, pagamento de contas e boletos, emissão de TED e DOC, além de transferência de recursos entre clientes de uma mesma instituição de pagamento.\n",
            "1 O que é \"Crédito Comercial\"?\n",
            "\t -> POSITIVE:::  Créditos comerciais compreendemfinanciamentos concedidos diretamente entre exportador e importador para aquisição de bens ou serviços em transações de comércio exterior. Podem assumir duas formas: Importador residente no Brasil efetua o pagamento ao exportador não residente, que assume o compromisso de, no futuro, entregar o bem ou serviço (adiantamento de compras). Implica saída de recursos financeiros do País e é um ativo externo recebível em bens ou serviços; Exportador residente no Brasil envia o bem ou presta o serviço ao importador não residente, que assume o compromisso de, no futuro, efetuar o pagamento devido (exportações a receber). Não implica saída de recursos financeiros do País e é um ativo externo exigível em moeda.    Atenção: Ainda que o financiamento esteja associado ao comércio de bens e serviços, se houver instituição financeira como credora, trata-se de empréstimo e não de crédito comercial.\n",
            "\t -> NEGATIVE:::   Não, se ele for associado de uma cooperativa que congregue somente funcionários de uma empresa ou grupo de empresas. Caso seja associado de uma cooperativa cujo vínculo não seja o empregador ou de uma cooperativa de livre admissão, não há necessidade de se desligar da cooperativa. Nos termos da regulamentação vigente, a cooperativa singular de crédito que não seja de livre admissão de associados pode fazer constar de seus estatutos previsão de associação de aposentados que, quando em atividade, atendiam os critérios estatutários de associação. A Lei 5.764, de 1971, em seu artigo 35, exige a exclusão de associados que deixem de atender aos requisitos estatutários de ingresso ou permanência na cooperativa. Assim, a administração da cooperativa está obrigada a providenciar a sua exclusão, nos termos legais. Adicionalmente, de acordo com o inciso III do artigo 21 da referida Lei, deve constar no estatuto social da cooperativa a forma de devolução do capital ao associado que se desliga.\n",
            "2 Onde acesso o protocolo de entrega?\n",
            "\t -> POSITIVE:::   Na tela de “Lista de declarações”, o número de protocolo pode ser observado ao lado de cada declaração entregue. Na versão impressa da declaração também consta o número do protocolo, data e hora de entrega.\n",
            "\t -> NEGATIVE:::  Não. As administradoras podem definir o prazo de duração do grupo, informação que é obrigatória constar no contrato. O prazo de duração do grupo somente pode ser alterado por meio de assembleia geral extraordinária na ocorrência de fatos que onerem em demasia os consorciados ou de outros eventos que dificultem a satisfação de suas obrigações. Os contratos de todos os consorciados, que façam parte do mesmo grupo de consórcio, devem ter prazo de conclusão na mesma data de encerramento do grupo, admitida a antecipação da liquidação do saldo devedor por iniciativa do cotista. Em situações específicas, no caso de consorciado admitido em grupo em andamento (em decorrência de venda de cota nova ou de cota de reposição), o prazo de duração do contrato para o novo cotista será o prazo remanescente para o término do grupo de consórcio, devendo a administradora considerar este prazo para calcular o valor da parcela.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HfxLlW3ESSBa"
      },
      "source": [
        "# Query Tokenization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hO9nRbzbSd0a",
        "outputId": "8a16b543-65b1-4bbd-9c2b-a6339053b6cb"
      },
      "source": [
        "query_maxlen = 18\n",
        "bsize = 3\n",
        "\n",
        "# =============================\n",
        "# ✨ build the q_obj tokenizer\n",
        "# =============================\n",
        "q_obj = tok(\n",
        "    queries, \n",
        "    padding='max_length', \n",
        "    truncation=True,\n",
        "    return_tensors='pt', \n",
        "    max_length=query_maxlen,\n",
        "    )\n",
        "\n",
        "# use only ids e mask keys in the q_obj dict\n",
        "q_ids, q_mask = q_obj['input_ids'], q_obj['attention_mask']\n",
        "q_ids, q_mask"
      ],
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(tensor([[  101, 23840, 10146, 56749, 61997, 10107, 10239, 10173, 15101, 22877,\n",
              "          10426, 87174, 10102, 10104, 54093, 11605,   136,   102],\n",
              "         [  101,   157, 10126,   147,   107, 62372, 20347,   107,   136,   102,\n",
              "              0,     0,     0,     0,     0,     0,     0,     0],\n",
              "         [  101, 12243, 45555,   157, 33212, 10132, 10102, 45276,   136,   102,\n",
              "              0,     0,     0,     0,     0,     0,     0,     0]]),\n",
              " tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
              "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],\n",
              "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]]))"
            ]
          },
          "metadata": {},
          "execution_count": 29
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iTzmInXQTqnJ",
        "outputId": "5839a0e5-0491-4072-f442-a7dd16c43043"
      },
      "source": [
        "query_batches = []\n",
        "# range (0, N, N)\n",
        "for offset in range(0, q_ids.size(0), bsize):\n",
        "    query_batches.append((q_ids[offset:offset+bsize], q_mask[offset:offset+bsize]))\n",
        "query_batches    "
      ],
      "execution_count": 70,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[(tensor([[  101, 23840, 10146, 56749, 61997, 10107, 10239, 10173, 15101, 22877,\n",
              "           10426, 87174, 10102, 10104, 54093, 11605,   136,   102],\n",
              "          [  101,   157, 10126,   147,   107, 62372, 20347,   107,   136,   102,\n",
              "               0,     0,     0,     0,     0,     0,     0,     0],\n",
              "          [  101, 12243, 45555,   157, 33212, 10132, 10102, 45276,   136,   102,\n",
              "               0,     0,     0,     0,     0,     0,     0,     0]]),\n",
              "  tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
              "          [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],\n",
              "          [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]]))]"
            ]
          },
          "metadata": {},
          "execution_count": 70
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gd-UrzmTWYdd"
      },
      "source": [
        "# Doc Tokenization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MtkN7fQGWe8u",
        "outputId": "a8ab9d17-e274-4098-b68a-7cb02e465635"
      },
      "source": [
        "# ===========================================================\n",
        "# ✨ build one batch with positive and negative docs \n",
        "# ===========================================================\n",
        "doc_bsize = pos_docs + neg_docs\n",
        "len(doc_bsize)"
      ],
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "6"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7epS4pHHWsQM",
        "outputId": "61f7133e-aabf-43e3-e5d8-34c9fad0e360"
      },
      "source": [
        "doc_maxlen = 64\n",
        "\n",
        "d_obj = tok(\n",
        "    doc_bsize, \n",
        "    padding='max_length', \n",
        "    truncation=True,\n",
        "    return_tensors='pt', \n",
        "    max_length=doc_maxlen,\n",
        "    )\n",
        "\n",
        "# utiliza somente ids e mask keys\n",
        "d_ids, d_mask = d_obj['input_ids'], d_obj['attention_mask']\n",
        "d_ids, d_mask"
      ],
      "execution_count": 55,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(tensor([[  101, 10253, 87174, 10102, 10104, 54093, 11605, 75143, 10542, 10173,\n",
              "          84677, 10107, 18255, 85267, 10102, 28903, 22760,   117,   157, 13249,\n",
              "          13129, 10112, 10542, 77984,   117,   143, 24590, 57460, 10141, 78689,\n",
              "          24363, 13198,   117, 10190, 36054, 33873, 64422, 41707, 10102, 62372,\n",
              "            117, 11328, 10126, 10146, 10104, 78426, 10165, 13668, 11313, 10109,\n",
              "          35710, 20145, 33658, 28579,   147, 19060, 17911, 14633, 11313, 18247,\n",
              "          13353, 18754, 75246,   102],\n",
              "         [  101, 87174, 81134, 85065, 28635, 11525, 10150, 54682, 13823, 25942,\n",
              "          10173, 84677, 10107, 79673, 10405, 40230, 19841,   147, 50373, 19841,\n",
              "          10239, 23955, 45426, 12574, 10102, 56651, 10391, 43189, 10252, 20041,\n",
              "          35564, 10102, 25451, 21954,   119, 21256, 94744, 15408, 22670,   131,\n",
              "          50373, 19841, 74850, 10181, 12369, 33575, 30114, 10112,   157, 80914,\n",
              "          10620, 40230, 19841, 11373, 74850,   117, 10126, 42954,   157, 90177,\n",
              "          11513, 10102,   117,   102],\n",
              "         [  101, 10135, 27994, 10102,   100, 14247, 10102, 58262, 12965,   100,\n",
              "            117,   157, 11855, 10102, 33212, 10132, 14396, 10542, 63834, 10351,\n",
              "          10620, 15841, 10102, 11822, 58262, 11115, 10405, 17734,   119, 10135,\n",
              "          24588, 73296, 51215, 10141, 58262, 11115, 12014, 30736,   157, 11855,\n",
              "          10154, 33212, 10132,   117, 10248,   147, 20018, 10102, 45276,   119,\n",
              "            102,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
              "              0,     0,     0,     0],\n",
              "         [  101,   143, 21471, 10102, 80914,   147, 10477, 21471, 36692, 12108,\n",
              "          62876, 10239,   143, 95541, 10102, 80914, 10107,   147, 63736, 10102,\n",
              "          23402, 10245,   117, 10190, 21107,   117, 80914, 10102, 21471, 10107,\n",
              "            147, 74788, 13320,   117, 25990, 87852, 10102, 21674,   147, 22747,\n",
              "            117, 17077, 10102, 63736, 10102, 23402, 10405, 67210, 10102, 10477,\n",
              "          18777, 78689, 10102, 80914,   119,   102,     0,     0,     0,     0,\n",
              "              0,     0,     0,     0],\n",
              "         [  101, 11373,   117, 10128, 12002, 10139, 13967, 52078, 10351, 10102,\n",
              "          10477, 89578, 10126, 11638, 49389, 12274, 39703, 58478, 10102, 10477,\n",
              "          14443, 10391, 12098, 10102, 24924,   119, 12887, 26625, 13967, 52078,\n",
              "          10351, 10102, 10477, 89578, 58459, 30626, 28485, 11373, 26625,   157,\n",
              "          10252, 31546, 32554, 10391, 10102, 10477, 89578, 10102, 16557, 94110,\n",
              "          22877,   117, 11373, 10240, 83606, 10102, 10128, 10143, 15519, 10131,\n",
              "          10141, 89578,   119,   102],\n",
              "         [  101, 11373,   119, 10146, 82051, 10416, 21256, 42274,   157, 43898,\n",
              "          11560, 10102, 30739, 11115, 10154, 12098,   117, 64466, 10126,   147,\n",
              "          15547, 48855, 27872, 30736, 10131, 10181, 26025,   119,   157, 43898,\n",
              "          11560, 10102, 30739, 11115, 10154, 12098, 39703, 14396, 10542, 17757,\n",
              "          11511, 10190, 24250, 10102, 62681, 24632, 17708, 21366, 64947, 10135,\n",
              "          49927, 13823, 10102, 37808, 10107, 10126, 10399, 19280, 10252, 33538,\n",
              "          10278, 10253, 10173,   102]]),\n",
              " tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
              "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
              "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "          1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
              "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "          1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],\n",
              "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
              "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]))"
            ]
          },
          "metadata": {},
          "execution_count": 55
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2S_G4zQbXrvj",
        "outputId": "6d78f5ee-e44c-46d6-8c86-842c801f5cd6"
      },
      "source": [
        "# ==========================\n",
        "# ✨ sort the doc's lengths \n",
        "# ==========================\n",
        "indices = d_mask.sum(-1).sort().indices\n",
        "print(indices)\n",
        "d_ids = d_ids[indices]\n",
        "d_mask = d_mask[indices]\n",
        "d_mask.sum(-1)"
      ],
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([2, 3, 0, 1, 4, 5])\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([51, 56, 64, 64, 64, 64])"
            ]
          },
          "metadata": {},
          "execution_count": 56
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MkgT-VYcXmsR",
        "outputId": "641e888e-9c57-40ed-cb75-120fb424dba6"
      },
      "source": [
        "# ==========================================\n",
        "# ✨ build the doc's batch: list(ids, mask)\n",
        "# ==========================================\n",
        "d_batches = []\n",
        "# range (0, N, N)\n",
        "for offset in range(0, d_ids.size(0), bsize):\n",
        "    d_batches.append((d_ids[offset:offset+bsize], d_mask[offset:offset+bsize]))\n",
        "d_batches    "
      ],
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[(tensor([[  101, 10135, 27994, 10102,   100, 14247, 10102, 58262, 12965,   100,\n",
              "             117,   157, 11855, 10102, 33212, 10132, 14396, 10542, 63834, 10351,\n",
              "           10620, 15841, 10102, 11822, 58262, 11115, 10405, 17734,   119, 10135,\n",
              "           24588, 73296, 51215, 10141, 58262, 11115, 12014, 30736,   157, 11855,\n",
              "           10154, 33212, 10132,   117, 10248,   147, 20018, 10102, 45276,   119,\n",
              "             102,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
              "               0,     0,     0,     0],\n",
              "          [  101,   143, 21471, 10102, 80914,   147, 10477, 21471, 36692, 12108,\n",
              "           62876, 10239,   143, 95541, 10102, 80914, 10107,   147, 63736, 10102,\n",
              "           23402, 10245,   117, 10190, 21107,   117, 80914, 10102, 21471, 10107,\n",
              "             147, 74788, 13320,   117, 25990, 87852, 10102, 21674,   147, 22747,\n",
              "             117, 17077, 10102, 63736, 10102, 23402, 10405, 67210, 10102, 10477,\n",
              "           18777, 78689, 10102, 80914,   119,   102,     0,     0,     0,     0,\n",
              "               0,     0,     0,     0],\n",
              "          [  101, 10253, 87174, 10102, 10104, 54093, 11605, 75143, 10542, 10173,\n",
              "           84677, 10107, 18255, 85267, 10102, 28903, 22760,   117,   157, 13249,\n",
              "           13129, 10112, 10542, 77984,   117,   143, 24590, 57460, 10141, 78689,\n",
              "           24363, 13198,   117, 10190, 36054, 33873, 64422, 41707, 10102, 62372,\n",
              "             117, 11328, 10126, 10146, 10104, 78426, 10165, 13668, 11313, 10109,\n",
              "           35710, 20145, 33658, 28579,   147, 19060, 17911, 14633, 11313, 18247,\n",
              "           13353, 18754, 75246,   102]]),\n",
              "  tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "           1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
              "          [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "           1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],\n",
              "          [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])),\n",
              " (tensor([[  101, 87174, 81134, 85065, 28635, 11525, 10150, 54682, 13823, 25942,\n",
              "           10173, 84677, 10107, 79673, 10405, 40230, 19841,   147, 50373, 19841,\n",
              "           10239, 23955, 45426, 12574, 10102, 56651, 10391, 43189, 10252, 20041,\n",
              "           35564, 10102, 25451, 21954,   119, 21256, 94744, 15408, 22670,   131,\n",
              "           50373, 19841, 74850, 10181, 12369, 33575, 30114, 10112,   157, 80914,\n",
              "           10620, 40230, 19841, 11373, 74850,   117, 10126, 42954,   157, 90177,\n",
              "           11513, 10102,   117,   102],\n",
              "          [  101, 11373,   117, 10128, 12002, 10139, 13967, 52078, 10351, 10102,\n",
              "           10477, 89578, 10126, 11638, 49389, 12274, 39703, 58478, 10102, 10477,\n",
              "           14443, 10391, 12098, 10102, 24924,   119, 12887, 26625, 13967, 52078,\n",
              "           10351, 10102, 10477, 89578, 58459, 30626, 28485, 11373, 26625,   157,\n",
              "           10252, 31546, 32554, 10391, 10102, 10477, 89578, 10102, 16557, 94110,\n",
              "           22877,   117, 11373, 10240, 83606, 10102, 10128, 10143, 15519, 10131,\n",
              "           10141, 89578,   119,   102],\n",
              "          [  101, 11373,   119, 10146, 82051, 10416, 21256, 42274,   157, 43898,\n",
              "           11560, 10102, 30739, 11115, 10154, 12098,   117, 64466, 10126,   147,\n",
              "           15547, 48855, 27872, 30736, 10131, 10181, 26025,   119,   157, 43898,\n",
              "           11560, 10102, 30739, 11115, 10154, 12098, 39703, 14396, 10542, 17757,\n",
              "           11511, 10190, 24250, 10102, 62681, 24632, 17708, 21366, 64947, 10135,\n",
              "           49927, 13823, 10102, 37808, 10107, 10126, 10399, 19280, 10252, 33538,\n",
              "           10278, 10253, 10173,   102]]),\n",
              "  tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
              "          [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
              "          [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]))]"
            ]
          },
          "metadata": {},
          "execution_count": 57
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ttHrpFx0aEgd"
      },
      "source": [
        "# Tensorize Triplets"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WRrxbldjaKBM",
        "outputId": "eac85054-f19f-4fc7-9e85-d9e2d63d12fb"
      },
      "source": [
        "# ========================================\n",
        "# ✨ assign ids e mask tokens and reshape\n",
        "# ========================================\n",
        "N = bsize\n",
        "Q_ids, Q_mask = q_ids, q_mask\n",
        "D_ids, D_mask = d_ids, d_mask\n",
        "D_ids, D_mask = D_ids.view(2, N, -1), D_mask.view(2, N, -1)\n",
        "\n",
        "print(f'Q_ids:  {q_ids.shape}')\n",
        "print(f'Q_mask: {q_mask.shape}\\n')\n",
        "\n",
        "print(f'D_ids:                {d_ids.shape}')\n",
        "print(f'D_ids.view(2, N, -1): {d_ids.view(2, N, -1).shape}\\n')\n",
        "print(f'D_mask:                {d_mask.shape}')\n",
        "print(f'D_mask.view(2, N, -1): {d_mask.view(2, N, -1).shape}')"
      ],
      "execution_count": 59,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Q_ids:  torch.Size([3, 18])\n",
            "Q_mask: torch.Size([3, 18])\n",
            "\n",
            "D_ids:                torch.Size([6, 64])\n",
            "D_ids.view(2, N, -1): torch.Size([2, 3, 64])\n",
            "\n",
            "D_mask:                torch.Size([6, 64])\n",
            "D_mask.view(2, N, -1): torch.Size([2, 3, 64])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XUg_gsjRaVTA",
        "outputId": "fc8a7327-c603-4722-c954-b866706702b1"
      },
      "source": [
        "# ==================================================================================================\n",
        "# ✨ get the max value between the len of i-th positive and the len of the i-th negative for i in N\n",
        "# ==================================================================================================\n",
        "maxlens = D_mask.sum(-1).max(0).values\n",
        "print(maxlens)\n",
        "indices = maxlens.sort().indices\n",
        "indices"
      ],
      "execution_count": 60,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([64, 64, 64])\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([0, 1, 2])"
            ]
          },
          "metadata": {},
          "execution_count": 60
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G6yPFTesKYSG"
      },
      "source": [
        "# =============================\n",
        "# ✨ sort Q_* e D_* by maxlens\n",
        "# =============================\n",
        "Q_ids, Q_mask = Q_ids[indices], Q_mask[indices]\n",
        "D_ids, D_mask = D_ids[:, indices], D_mask[:, indices]"
      ],
      "execution_count": 66,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HuOhT_2hLwvT",
        "outputId": "bc69883b-0670-4196-a7ee-f998a79ea0f8"
      },
      "source": [
        "# ==================================================================\n",
        "# ✨ split the positive e negative ids and mask from D_ids e D_mask\n",
        "# ==================================================================\n",
        "(positive_ids, negative_ids), (positive_mask, negative_mask) = D_ids, D_mask\n",
        "positive_ids.shape, negative_ids.shape, positive_mask.shape, negative_mask.shape"
      ],
      "execution_count": 67,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(torch.Size([3, 64]),\n",
              " torch.Size([3, 64]),\n",
              " torch.Size([3, 64]),\n",
              " torch.Size([3, 64]))"
            ]
          },
          "metadata": {},
          "execution_count": 67
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PExtEqAvLwsq"
      },
      "source": [
        "# ====================================================\n",
        "# ✨ build batches to positive_docs and negative_docs\n",
        "# ====================================================\n",
        "(positive_ids, negative_ids), (positive_mask, negative_mask) = D_ids, D_mask\n",
        "\n",
        "positive_batches = []\n",
        "for offset in range(0, positive_ids.size(0), bsize):\n",
        "    positive_batches.append((positive_ids[offset:offset+bsize], positive_mask[offset:offset+bsize]))    \n",
        "\n",
        "negative_batches = []\n",
        "for offset in range(0, negative_ids.size(0), bsize):\n",
        "    negative_batches.append((negative_ids[offset:offset+bsize], negative_mask[offset:offset+bsize]))    "
      ],
      "execution_count": 72,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SL1meeaSKMHc",
        "outputId": "7c62dd67-93ce-4189-db31-7f218af236f8"
      },
      "source": [
        "# ===========================================================\n",
        "# ✨ group the batches: (query_, positive_, negative)batches\n",
        "# ===========================================================\n",
        "batches = []\n",
        "for (q_ids, q_mask), (p_ids, p_mask), (n_ids, n_mask) in zip(query_batches, positive_batches, negative_batches):\n",
        "    Q = (torch.cat((q_ids, q_ids)), torch.cat((q_mask, q_mask))) # <- duplicate Q (one to pos docs and another for neg docs)\n",
        "    D = (torch.cat((p_ids, n_ids)), torch.cat((p_mask, n_mask)))\n",
        "    batches.append((Q, D))\n",
        "batches"
      ],
      "execution_count": 73,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[((tensor([[  101, 23840, 10146, 56749, 61997, 10107, 10239, 10173, 15101, 22877,\n",
              "            10426, 87174, 10102, 10104, 54093, 11605,   136,   102],\n",
              "           [  101,   157, 10126,   147,   107, 62372, 20347,   107,   136,   102,\n",
              "                0,     0,     0,     0,     0,     0,     0,     0],\n",
              "           [  101, 12243, 45555,   157, 33212, 10132, 10102, 45276,   136,   102,\n",
              "                0,     0,     0,     0,     0,     0,     0,     0],\n",
              "           [  101, 23840, 10146, 56749, 61997, 10107, 10239, 10173, 15101, 22877,\n",
              "            10426, 87174, 10102, 10104, 54093, 11605,   136,   102],\n",
              "           [  101,   157, 10126,   147,   107, 62372, 20347,   107,   136,   102,\n",
              "                0,     0,     0,     0,     0,     0,     0,     0],\n",
              "           [  101, 12243, 45555,   157, 33212, 10132, 10102, 45276,   136,   102,\n",
              "                0,     0,     0,     0,     0,     0,     0,     0]]),\n",
              "   tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
              "           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],\n",
              "           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],\n",
              "           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
              "           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],\n",
              "           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]])),\n",
              "  (tensor([[  101, 10135, 27994, 10102,   100, 14247, 10102, 58262, 12965,   100,\n",
              "              117,   157, 11855, 10102, 33212, 10132, 14396, 10542, 63834, 10351,\n",
              "            10620, 15841, 10102, 11822, 58262, 11115, 10405, 17734,   119, 10135,\n",
              "            24588, 73296, 51215, 10141, 58262, 11115, 12014, 30736,   157, 11855,\n",
              "            10154, 33212, 10132,   117, 10248,   147, 20018, 10102, 45276,   119,\n",
              "              102,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
              "                0,     0,     0,     0],\n",
              "           [  101,   143, 21471, 10102, 80914,   147, 10477, 21471, 36692, 12108,\n",
              "            62876, 10239,   143, 95541, 10102, 80914, 10107,   147, 63736, 10102,\n",
              "            23402, 10245,   117, 10190, 21107,   117, 80914, 10102, 21471, 10107,\n",
              "              147, 74788, 13320,   117, 25990, 87852, 10102, 21674,   147, 22747,\n",
              "              117, 17077, 10102, 63736, 10102, 23402, 10405, 67210, 10102, 10477,\n",
              "            18777, 78689, 10102, 80914,   119,   102,     0,     0,     0,     0,\n",
              "                0,     0,     0,     0],\n",
              "           [  101, 10253, 87174, 10102, 10104, 54093, 11605, 75143, 10542, 10173,\n",
              "            84677, 10107, 18255, 85267, 10102, 28903, 22760,   117,   157, 13249,\n",
              "            13129, 10112, 10542, 77984,   117,   143, 24590, 57460, 10141, 78689,\n",
              "            24363, 13198,   117, 10190, 36054, 33873, 64422, 41707, 10102, 62372,\n",
              "              117, 11328, 10126, 10146, 10104, 78426, 10165, 13668, 11313, 10109,\n",
              "            35710, 20145, 33658, 28579,   147, 19060, 17911, 14633, 11313, 18247,\n",
              "            13353, 18754, 75246,   102],\n",
              "           [  101, 87174, 81134, 85065, 28635, 11525, 10150, 54682, 13823, 25942,\n",
              "            10173, 84677, 10107, 79673, 10405, 40230, 19841,   147, 50373, 19841,\n",
              "            10239, 23955, 45426, 12574, 10102, 56651, 10391, 43189, 10252, 20041,\n",
              "            35564, 10102, 25451, 21954,   119, 21256, 94744, 15408, 22670,   131,\n",
              "            50373, 19841, 74850, 10181, 12369, 33575, 30114, 10112,   157, 80914,\n",
              "            10620, 40230, 19841, 11373, 74850,   117, 10126, 42954,   157, 90177,\n",
              "            11513, 10102,   117,   102],\n",
              "           [  101, 11373,   117, 10128, 12002, 10139, 13967, 52078, 10351, 10102,\n",
              "            10477, 89578, 10126, 11638, 49389, 12274, 39703, 58478, 10102, 10477,\n",
              "            14443, 10391, 12098, 10102, 24924,   119, 12887, 26625, 13967, 52078,\n",
              "            10351, 10102, 10477, 89578, 58459, 30626, 28485, 11373, 26625,   157,\n",
              "            10252, 31546, 32554, 10391, 10102, 10477, 89578, 10102, 16557, 94110,\n",
              "            22877,   117, 11373, 10240, 83606, 10102, 10128, 10143, 15519, 10131,\n",
              "            10141, 89578,   119,   102],\n",
              "           [  101, 11373,   119, 10146, 82051, 10416, 21256, 42274,   157, 43898,\n",
              "            11560, 10102, 30739, 11115, 10154, 12098,   117, 64466, 10126,   147,\n",
              "            15547, 48855, 27872, 30736, 10131, 10181, 26025,   119,   157, 43898,\n",
              "            11560, 10102, 30739, 11115, 10154, 12098, 39703, 14396, 10542, 17757,\n",
              "            11511, 10190, 24250, 10102, 62681, 24632, 17708, 21366, 64947, 10135,\n",
              "            49927, 13823, 10102, 37808, 10107, 10126, 10399, 19280, 10252, 33538,\n",
              "            10278, 10253, 10173,   102]]),\n",
              "   tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "            1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
              "           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "            1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],\n",
              "           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
              "           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
              "           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
              "           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
              "            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])))]"
            ]
          },
          "metadata": {},
          "execution_count": 73
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_UBloctzgqaj"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tFJKmRl3g6hd",
        "outputId": "20835eda-69b9-4901-d4ce-ab1251684c3b"
      },
      "source": [
        "query_ids = batches[0][0][0]\n",
        "query_mask = batches[0][0][1]\n",
        "print('q_input_ids', query_ids.shape)\n",
        "print('q_attn_mask', query_mask.shape)\n",
        "\n",
        "doc_ids = batches[0][1][0]\n",
        "doc_mask = batches[0][1][1]\n",
        "print('doc_input_ids', doc_ids.shape)\n",
        "print('doc_attn_mask', doc_mask.shape)"
      ],
      "execution_count": 96,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "q_input_ids torch.Size([6, 18])\n",
            "q_attn_mask torch.Size([6, 18])\n",
            "doc_input_ids torch.Size([6, 64])\n",
            "doc_attn_mask torch.Size([6, 64])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z4NpcFXfg2G4"
      },
      "source": [
        "# =============================\n",
        "# ✨ init BERT Model\n",
        "# =============================\n",
        "model = BertModel.from_pretrained(path_model, return_dict=True)\n",
        "\n",
        "# =================================\n",
        "# ✨ build a dense layer\n",
        "# =================================\n",
        "linear = torch.nn.Linear(model.config.hidden_size, 128, bias=False)"
      ],
      "execution_count": 104,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EH1gKpWRunwA"
      },
      "source": [
        "# Prepare Q "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "P8JLVf4pp4cj",
        "outputId": "372afa65-292e-4d85-ec7d-810fac2d7a82"
      },
      "source": [
        "# ====================================================\n",
        "# ✨ input the input_ids and attention_mask into BERT\n",
        "# ====================================================\n",
        "outs = model(input_ids=query_ids, attention_mask=query_mask)\n",
        "Q = outs['last_hidden_state']\n",
        "Q.shape"
      ],
      "execution_count": 105,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([6, 18, 768])"
            ]
          },
          "metadata": {},
          "execution_count": 105
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_XAlx3Q0p4Z8",
        "outputId": "736fe9d4-d4fd-462e-9b7e-7f38f8b2fcfd"
      },
      "source": [
        "# ================================\n",
        "# ✨ perform Q in the dense layer\n",
        "# ================================\n",
        "Q = linear(Q)\n",
        "Q.shape"
      ],
      "execution_count": 106,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([6, 18, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 106
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TLb94IuRp4XY",
        "outputId": "9e081884-4fa0-4d2f-be07-cf92f8ae5d5c"
      },
      "source": [
        "# ===================\n",
        "# ✨ normalize in L2\n",
        "# ===================\n",
        "Q = torch.nn.functional.normalize(Q, p=2, dim=2)\n",
        "Q.shape"
      ],
      "execution_count": 107,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([6, 18, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 107
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lBv52Tqnp4VT"
      },
      "source": [
        "# Prepare D"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Sv_4Leyvp4Sx",
        "outputId": "833e7358-e926-40c6-a2e0-e72b7e80fcfb"
      },
      "source": [
        "# ================================================\n",
        "# ✨ input input_ids and attention_mask into BERT\n",
        "# ================================================\n",
        "outs = model(input_ids=doc_ids, attention_mask=doc_mask)\n",
        "D = outs['last_hidden_state']\n",
        "D.shape"
      ],
      "execution_count": 109,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([6, 64, 768])"
            ]
          },
          "metadata": {},
          "execution_count": 109
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AXmCWW-tp4Qg",
        "outputId": "71c96ab7-6591-4fb9-f806-a52f5978ab96"
      },
      "source": [
        "# ================================\n",
        "# ✨ perform D in the dense layer\n",
        "# ================================\n",
        "D = linear(D)\n",
        "D.shape"
      ],
      "execution_count": 110,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([6, 64, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 110
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q0_4QhtIp4Oa",
        "outputId": "20c7f431-14e9-43fb-feab-94af8bfda186"
      },
      "source": [
        "# ======================\n",
        "# ✨ filter D with mask\n",
        "# ======================\n",
        "mask = torch.tensor([[(x != 0) for x in d] for d in doc_ids.cpu().tolist()])\n",
        "print(mask.shape)\n",
        "D = D * mask.unsqueeze(2).float()\n",
        "print(D.shape)\n",
        "\n",
        "# ===================\n",
        "# ✨ normalize in L2\n",
        "# ===================\n",
        "D = torch.nn.functional.normalize(D, p=2, dim=2)\n",
        "D.shape"
      ],
      "execution_count": 111,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([6, 64])\n",
            "torch.Size([6, 64, 128])\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([6, 64, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 111
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "peAovi3gp4KB"
      },
      "source": [
        "# Get the score between Q and D"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Avs1uP1gp4H5",
        "outputId": "0f7aaba3-30e6-4fba-f5a5-5a2cd3e46f2b"
      },
      "source": [
        "print(f'Q shape: {Q.size()} -- D shape: {D.size()}\\n')\n",
        "\n",
        "scores = torch.einsum('nqe, nde -> nqd', Q, D) \n",
        "print(f'Score shape:                       {scores.size()}')\n",
        "\n",
        "scores = scores.max(2)\n",
        "print(f'Scores.max(2).values shape:        {scores.values.size()}')\n",
        "\n",
        "scores = scores.values.sum(1)\n",
        "print(f'Scores.max(2).values.sum(1) shape: {scores.size()}')\n",
        "\n",
        "print(f'Scores: {[float(str(s.item())[:7]) for s in scores]:}')"
      ],
      "execution_count": 128,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Q shape: torch.Size([6, 18, 128]) -- D shape: torch.Size([6, 64, 128])\n",
            "\n",
            "Score shape:                       torch.Size([6, 18, 64])\n",
            "Scores.max(2).values shape:        torch.Size([6, 18])\n",
            "Scores.max(2).values.sum(1) shape: torch.Size([6])\n",
            "Scores: [10.7851, 9.7143, 9.93258, 11.5629, 10.0149, 10.114]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bXv8OowN1h1P"
      },
      "source": [
        "# Otimization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "l-PUt6c0pxlz",
        "outputId": "6f43d950-e820-4b95-d77f-154f76695271"
      },
      "source": [
        "# ======================================================\n",
        "# ✨ build pseudo-labels for the classe 0 (torch zeros)\n",
        "# ======================================================\n",
        "# labels has the size of the batch bsize\n",
        "labels = torch.zeros(bsize, dtype=torch.long)\n",
        "print('labels shape', labels.shape, '\\n')"
      ],
      "execution_count": 129,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "labels shape torch.Size([3]) \n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Xq23gr9p2PvV",
        "outputId": "217627eb-5c7d-49e7-bde2-5fc90bc806cd"
      },
      "source": [
        "# =======================\n",
        "# ✨ reshape the  scores\n",
        "# =======================\n",
        "scores = scores.view(2, -1).permute(1, 0)\n",
        "print(scores.shape)\n",
        "scores"
      ],
      "execution_count": 130,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([3, 2])\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[10.7852, 11.5629],\n",
              "        [ 9.7143, 10.0150],\n",
              "        [ 9.9326, 10.1141]], grad_fn=<PermuteBackward>)"
            ]
          },
          "metadata": {},
          "execution_count": 130
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eBx_92kIoj9Z",
        "outputId": "2ae3127e-8957-413e-f7d5-d2424fde1483"
      },
      "source": [
        "# ============\n",
        "# ✨ get loss\n",
        "# ============\n",
        "# init the CE Loss\n",
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "loss = criterion(scores, labels[:scores.size(0)])\n",
        "loss.item()"
      ],
      "execution_count": 131,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.9328517913818359"
            ]
          },
          "metadata": {},
          "execution_count": 131
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XC9sJrysnO9h"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}