FAISS.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "cKtMj52NvGlP"
      },
      "outputs": [],
      "source": [
        "import requests\n",
        "from io import StringIO\n",
        "import pandas as pd"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "res = requests.get('https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/sick2014/SICK_train.txt')\n",
        "# create dataframe\n",
        "data = pd.read_csv(StringIO(res.text), sep='\\t')\n",
        "data.head()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 293
        },
        "id": "T6J-FvRpvIL1",
        "outputId": "3f58e955-47a7-4613-b054-dcf910816920"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "   pair_ID                                         sentence_A  \\\n",
              "0        1  A group of kids is playing in a yard and an ol...   \n",
              "1        2  A group of children is playing in the house an...   \n",
              "2        3  The young boys are playing outdoors and the ma...   \n",
              "3        5  The kids are playing outdoors near a man with ...   \n",
              "4        9  The young boys are playing outdoors and the ma...   \n",
              "\n",
              "                                          sentence_B  relatedness_score  \\\n",
              "0  A group of boys in a yard is playing and a man...                4.5   \n",
              "1  A group of kids is playing in a yard and an ol...                3.2   \n",
              "2  The kids are playing outdoors near a man with ...                4.7   \n",
              "3  A group of kids is playing in a yard and an ol...                3.4   \n",
              "4  A group of kids is playing in a yard and an ol...                3.7   \n",
              "\n",
              "  entailment_judgment  \n",
              "0             NEUTRAL  \n",
              "1             NEUTRAL  \n",
              "2          ENTAILMENT  \n",
              "3             NEUTRAL  \n",
              "4             NEUTRAL  "
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-a0da876a-aa0e-432b-bacb-f15db09eb5b0\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>pair_ID</th>\n",
              "      <th>sentence_A</th>\n",
              "      <th>sentence_B</th>\n",
              "      <th>relatedness_score</th>\n",
              "      <th>entailment_judgment</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>1</td>\n",
              "      <td>A group of kids is playing in a yard and an ol...</td>\n",
              "      <td>A group of boys in a yard is playing and a man...</td>\n",
              "      <td>4.5</td>\n",
              "      <td>NEUTRAL</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>2</td>\n",
              "      <td>A group of children is playing in the house an...</td>\n",
              "      <td>A group of kids is playing in a yard and an ol...</td>\n",
              "      <td>3.2</td>\n",
              "      <td>NEUTRAL</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>3</td>\n",
              "      <td>The young boys are playing outdoors and the ma...</td>\n",
              "      <td>The kids are playing outdoors near a man with ...</td>\n",
              "      <td>4.7</td>\n",
              "      <td>ENTAILMENT</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>5</td>\n",
              "      <td>The kids are playing outdoors near a man with ...</td>\n",
              "      <td>A group of kids is playing in a yard and an ol...</td>\n",
              "      <td>3.4</td>\n",
              "      <td>NEUTRAL</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>9</td>\n",
              "      <td>The young boys are playing outdoors and the ma...</td>\n",
              "      <td>A group of kids is playing in a yard and an ol...</td>\n",
              "      <td>3.7</td>\n",
              "      <td>NEUTRAL</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-a0da876a-aa0e-432b-bacb-f15db09eb5b0')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-a0da876a-aa0e-432b-bacb-f15db09eb5b0 button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-a0da876a-aa0e-432b-bacb-f15db09eb5b0');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-3d5a9d99-c74e-484a-b131-3c4e486da977\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-3d5a9d99-c74e-484a-b131-3c4e486da977')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-3d5a9d99-c74e-484a-b131-3c4e486da977 button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "dataframe",
              "variable_name": "data",
              "summary": "{\n  \"name\": \"data\",\n  \"rows\": 4500,\n  \"fields\": [\n    {\n      \"column\": \"pair_ID\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 2866,\n        \"min\": 1,\n        \"max\": 10000,\n        \"num_unique_values\": 4500,\n        \"samples\": [\n          5197,\n          8173,\n          6572\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"sentence_A\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 3146,\n        \"samples\": [\n          \"A person with a red hooded jacket and a black backpack is walking near a rocky wall\",\n          \"A man is jumping off a low wall\",\n          \"The rhino is grazing on the grass\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"sentence_B\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 3132,\n        \"samples\": [\n          \"The crowd is watching two racing cars that are leaving the starting line\",\n          \"A man on a bike is jumping on a pyramid-shaped ramp\",\n          \"A dog that has a brown and white coat is trotting through shallow water\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"relatedness_score\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 1.0252348140884235,\n        \"min\": 1.0,\n        \"max\": 5.0,\n        \"num_unique_values\": 125,\n        \"samples\": [\n          3.635,\n          2.685,\n          4.6\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"entailment_judgment\",\n      \"properties\": {\n        \"dtype\": \"category\",\n        \"num_unique_values\": 3,\n        \"samples\": [\n          \"NEUTRAL\",\n          \"ENTAILMENT\",\n          \"CONTRADICTION\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}"
            }
          },
          "metadata": {},
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "sentences = data['sentence_A'].tolist()\n",
        "sentences[:5]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rqrM8qUjvO_0",
        "outputId": "7df9ac09-d14f-440b-8c63-be2793bd9b53"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['A group of kids is playing in a yard and an old man is standing in the background',\n",
              " 'A group of children is playing in the house and there is no man standing in the background',\n",
              " 'The young boys are playing outdoors and the man is smiling nearby',\n",
              " 'The kids are playing outdoors near a man with a smile',\n",
              " 'The young boys are playing outdoors and the man is smiling nearby']"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "sentences = data['sentence_A'].tolist()\n",
        "sentence_b = data['sentence_B'].tolist()\n",
        "sentences.extend(sentence_b)  # merge them\n",
        "len(set(sentences))  # together we have ~4.5K unique sentences"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fKMf5P76vRVB",
        "outputId": "ba78012f-0487-4c58-aa37-bc39198554ed"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "4802"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "urls = [\n",
        "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2012/MSRpar.train.tsv',\n",
        "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2012/MSRpar.test.tsv',\n",
        "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2012/OnWN.test.tsv',\n",
        "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2013/OnWN.test.tsv',\n",
        "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2014/OnWN.test.tsv',\n",
        "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2014/images.test.tsv',\n",
        "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2015/images.test.tsv'\n",
        "]"
      ],
      "metadata": {
        "id": "D8XOy4OXvTWl"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for url in urls:\n",
        "    res = requests.get(url)\n",
        "    # extract to dataframe\n",
        "    data = pd.read_csv(StringIO(res.text), sep='\\t', header=None, on_bad_lines='skip')\n",
        "    # add to columns 1 and 2 to sentences list\n",
        "    sentences.extend(data[1].tolist())\n",
        "    sentences.extend(data[2].tolist())"
      ],
      "metadata": {
        "id": "4wTZTmIwvVXF"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len(set(sentences))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kMX85YM0vjss",
        "outputId": "32878152-7dbd-425c-ec4f-7346d1d48964"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "14505"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "sentences = [word for word in list(set(sentences)) if type(word) is str]"
      ],
      "metadata": {
        "id": "qz0lzwfhvviV"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from sentence_transformers import SentenceTransformer\n",
        "# initialize sentence transformer model\n",
        "model = SentenceTransformer('bert-base-nli-mean-tokens')\n",
        "# create sentence embeddings\n",
        "sentence_embeddings = model.encode(sentences)\n",
        "sentence_embeddings.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tRqQYFqLvz2J",
        "outputId": "1bad3301-d91c-4085-ceea-e8eecc439b92"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/sentence_transformers/cross_encoder/CrossEncoder.py:13: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
            "  from tqdm.autonotebook import tqdm, trange\n",
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:89: UserWarning: \n",
            "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
            "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
            "You will be able to reuse this secret in all of your notebooks.\n",
            "Please note that authentication is recommended but still optional to access public models or datasets.\n",
            "  warnings.warn(\n",
            "/usr/local/lib/python3.10/dist-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(14504, 768)"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install faiss-gpu"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fU6dZF3qv1Et",
        "outputId": "611a343d-9dd1-4442-aefb-8a9f6098ad23"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: faiss-gpu in /usr/local/lib/python3.10/dist-packages (1.7.2)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import faiss\n"
      ],
      "metadata": {
        "id": "JplTmDBswC28"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "d = sentence_embeddings.shape[1]\n",
        "d"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-EnRA32swVYv",
        "outputId": "59e2adea-8f39-46c5-8a4c-8913d963b733"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "768"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "index = faiss.IndexFlatL2(d)"
      ],
      "metadata": {
        "id": "avAAA-87wXcc"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "index.is_trained"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E0DxObHFwZCY",
        "outputId": "7c16cac1-541d-4486-91fd-947be6f5e8ac"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "index.add(sentence_embeddings)"
      ],
      "metadata": {
        "id": "FNsi_lKnwamc"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "index.ntotal"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UKmOydbTwd90",
        "outputId": "b6c936c1-581f-489e-c37d-23a60684a2c8"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "14504"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "k = 4\n",
        "xq = model.encode([\"Someone sprints with a football\"])"
      ],
      "metadata": {
        "id": "rE8-iAbLwgSM"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "D, I = index.search(xq, k)  # search\n",
        "print(I)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E42e_km0wjsk",
        "outputId": "dbb4f8bf-6bd0-465b-dc5e-48a2297b33f5"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[1465 4653 1339 9118]]\n",
            "CPU times: user 16.4 ms, sys: 899 µs, total: 17.3 ms\n",
            "Wall time: 22.6 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "[f'{i}: {sentences[i]}' for i in I[0]]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W1Wde9P3wypF",
        "outputId": "375e83b8-2866-406d-cdf5-753ab1c22912"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['1465: A group of football players is running in the field',\n",
              " '4653: A group of people playing football is running in the field',\n",
              " '1339: Two groups of people are playing football',\n",
              " '9118: A person playing football is running past an official carrying a football']"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "\n",
        "vecs = np.zeros((k, d))\n",
        "# then iterate through each ID from I and add the reconstructed vector to our zero-array\n",
        "for i, val in enumerate(I[0].tolist()):\n",
        "    vecs[i, :] = index.reconstruct(val)"
      ],
      "metadata": {
        "id": "lc3yIrWCwmog"
      },
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vecs.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pYJpci_FwtWI",
        "outputId": "f8818351-a895-439e-cb86-a497ed7ca030"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(4, 768)"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "vecs[0][:100]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iQYnS3hPxnPC",
        "outputId": "99569e5b-9a57-4a7b-cd61-47baedc7201d"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([ 0.01627072,  0.2232592 , -0.15037425, -0.30747271, -0.27122465,\n",
              "       -0.10593155, -0.06460934,  0.04738171, -0.73349047, -0.37657681,\n",
              "       -0.76762789,  0.16902882,  0.53107649,  0.51176697,  1.14415824,\n",
              "       -0.08562881, -0.67240071, -0.96637076,  0.02545465, -0.21559809,\n",
              "       -1.25656545, -0.82982159, -0.09825006, -0.21850838,  0.50610238,\n",
              "        0.10527924,  0.50396848,  0.6524294 , -1.39458752,  0.65847486,\n",
              "       -0.21525319, -0.22487473,  0.818183  ,  0.08464295, -0.76141769,\n",
              "       -0.28928289, -0.09825794, -0.73046207,  0.07855801, -0.84354597,\n",
              "       -0.59242105,  0.7747131 , -1.20920527, -0.22757922, -1.30733585,\n",
              "       -0.23081493, -1.31322539,  0.01629098, -0.97285485,  0.19308187,\n",
              "        0.47424555,  1.18920887, -1.96741295, -0.70061141, -0.2963869 ,\n",
              "        0.60533738,  0.62407422, -0.70340371, -0.86754245,  0.17673171,\n",
              "       -0.19170482, -0.02951987,  0.22623563, -0.16695446, -0.80402559,\n",
              "       -0.45918921,  0.69675452, -0.24928184, -1.01478708, -0.92174506,\n",
              "       -0.33842692, -0.39296779, -0.83734852, -0.11479225,  0.4604969 ,\n",
              "       -1.45211208,  0.60310405,  0.38696292, -0.04061218,  0.00453171,\n",
              "        0.24117817,  0.0539629 ,  0.07506458,  1.05115831,  0.1238397 ,\n",
              "       -0.71281147,  0.11722863,  0.52238184, -0.04581185,  0.26827046,\n",
              "        0.85985392, -0.35669914, -0.64667076, -0.54357976, -0.04310478,\n",
              "        0.95139098, -0.15605754, -0.49625325, -0.11140213,  0.1561012 ])"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "nlist = 50  # how many cells\n",
        "quantizer = faiss.IndexFlatL2(d)\n",
        "index = faiss.IndexIVFFlat(quantizer, d, nlist)"
      ],
      "metadata": {
        "id": "Wvq5B8wwxokn"
      },
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "index.is_trained"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7JMDVI7Fxs6Q",
        "outputId": "5c582c16-53db-4db1-ee90-2997ea10a726"
      },
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "False"
            ]
          },
          "metadata": {},
          "execution_count": 38
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "index.train(sentence_embeddings)\n",
        "index.is_trained  # check if index is now trained"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DmpVaDanxybz",
        "outputId": "3816281a-de1a-4ce7-e0a7-b47321aab4dd"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "index.add(sentence_embeddings)\n",
        "index.ntotal  # number of embeddings indexed"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6ltiWK9Ixz5D",
        "outputId": "19aab722-08d8-4b0c-9332-968621d76693"
      },
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "14504"
            ]
          },
          "metadata": {},
          "execution_count": 40
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "D, I = index.search(xq, k)  # search\n",
        "print(I)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rnnf15Ejx1Nq",
        "outputId": "dee7abe1-a24f-4e17-abd8-ed8cd5021af7"
      },
      "execution_count": 41,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[1465 4653 1339 9118]]\n",
            "CPU times: user 729 µs, sys: 984 µs, total: 1.71 ms\n",
            "Wall time: 1.74 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "[f'{i}: {sentences[i]}' for i in I[0]]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qQGM_h763Q8D",
        "outputId": "1a07376e-57a8-4f01-fe17-500dd9d9167d"
      },
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['1465: A group of football players is running in the field',\n",
              " '4653: A group of people playing football is running in the field',\n",
              " '1339: Two groups of people are playing football',\n",
              " '9118: A person playing football is running past an official carrying a football']"
            ]
          },
          "metadata": {},
          "execution_count": 42
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "index.nprobe = 10"
      ],
      "metadata": {
        "id": "x-KKRUsnx6Uq"
      },
      "execution_count": 43,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "D, I = index.search(xq, k)  # search\n",
        "print(I)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lPRZ5Fwyx2cz",
        "outputId": "4aa5bad6-7edf-404b-dce1-e59a76f2c8c5"
      },
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[1465 4653 1339 9118]]\n",
            "CPU times: user 4.4 ms, sys: 29 µs, total: 4.43 ms\n",
            "Wall time: 3.67 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "[f'{i}: {sentences[i]}' for i in I[0]]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eFrN5Rqf3NSr",
        "outputId": "118065cd-23d2-4bab-c8a6-a019a74ea2d4"
      },
      "execution_count": 45,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['1465: A group of football players is running in the field',\n",
              " '4653: A group of people playing football is running in the field',\n",
              " '1339: Two groups of people are playing football',\n",
              " '9118: A person playing football is running past an official carrying a football']"
            ]
          },
          "metadata": {},
          "execution_count": 45
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "index.make_direct_map()"
      ],
      "metadata": {
        "id": "jakx7jnUx9dn"
      },
      "execution_count": 46,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "index.reconstruct(7460)[:100]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZqUrglOryC_v",
        "outputId": "0445ba4e-abd4-4b57-b116-14accb217f36"
      },
      "execution_count": 47,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([ 0.3074015 ,  0.29784325,  1.169849  ,  0.49171206,  0.15579304,\n",
              "        0.2688289 ,  0.7620417 ,  0.20051564,  0.27143443, -0.3085456 ,\n",
              "       -1.0383936 ,  0.5424543 , -0.07457361,  0.7191684 ,  0.98483235,\n",
              "       -0.3183309 , -0.80500513, -0.03861646,  1.7780961 ,  0.06267159,\n",
              "       -0.4335138 ,  0.536375  , -0.10101058, -0.7650174 ,  0.20436549,\n",
              "        0.68602496,  0.12260347, -0.33803013, -1.3555557 , -0.32722074,\n",
              "        0.16515845,  0.626782  ,  0.2441674 , -0.45406246, -0.53605086,\n",
              "        0.27161562,  0.7422892 , -0.51232   ,  0.3075459 ,  0.00270357,\n",
              "        0.42943513,  0.45058873,  0.06177604, -0.177378  , -0.6183112 ,\n",
              "       -0.44597647,  0.6259822 ,  0.38553518, -1.0830989 , -0.11827786,\n",
              "       -0.3288778 , -0.261656  ,  0.9885873 , -0.35565   , -0.6618367 ,\n",
              "       -0.369806  , -0.40991232, -0.8843489 , -0.4786378 ,  0.3552321 ,\n",
              "       -0.34558553, -0.6263227 ,  0.5656237 ,  0.9058484 , -0.58895785,\n",
              "       -0.02673524,  0.25150937, -0.2671459 , -1.2325639 , -0.54587233,\n",
              "        1.4144933 ,  0.11601061, -0.81625277,  0.27572548, -0.3395115 ,\n",
              "       -0.851712  , -0.08868876,  0.7501988 , -0.3911147 ,  0.84158885,\n",
              "       -0.23355263,  0.53382736,  0.5602867 ,  0.8338014 , -0.2642818 ,\n",
              "       -0.2811191 ,  1.3060477 ,  1.3242023 , -0.15578076,  0.21701209,\n",
              "        0.41253853,  0.78784215,  0.37363765, -1.5447763 ,  0.71700203,\n",
              "       -0.8158058 , -0.45044863,  0.721078  , -0.35364124, -0.25088114],\n",
              "      dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 47
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "m = 8  # number of centroid IDs in final compressed vectors\n",
        "bits = 8 # number of bits in each centroid\n",
        "\n",
        "quantizer = faiss.IndexFlatL2(d)  # we keep the same L2 distance flat index\n",
        "index = faiss.IndexIVFPQ(quantizer, d, nlist, m, bits)"
      ],
      "metadata": {
        "id": "LjjfBzcryD53"
      },
      "execution_count": 48,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "index.is_trained"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hZ-BM6MCyRMf",
        "outputId": "e553f002-5582-4396-dcbd-d02b9c06a13d"
      },
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "False"
            ]
          },
          "metadata": {},
          "execution_count": 49
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "index.train(sentence_embeddings)"
      ],
      "metadata": {
        "id": "OtUuuOgbySmn"
      },
      "execution_count": 50,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "index.add(sentence_embeddings)"
      ],
      "metadata": {
        "id": "hQgTlcwKyT6S"
      },
      "execution_count": 51,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "index.nprobe = 10"
      ],
      "metadata": {
        "id": "3YKR8Zl1yVQD"
      },
      "execution_count": 52,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "D, I = index.search(xq, k)\n",
        "print(I)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "w8eVaUuiyWYD",
        "outputId": "664699f4-1ed2-459f-d5eb-6bdd223fc26f"
      },
      "execution_count": 53,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[ 644  706 1444 1339]]\n",
            "CPU times: user 1.48 ms, sys: 29 µs, total: 1.51 ms\n",
            "Wall time: 1.52 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "[f'{i}: {sentences[i]}' for i in I[0]]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mUPIF2XByXTD",
        "outputId": "ede9b3e0-4328-428b-bbcb-987aa2a98e76"
      },
      "execution_count": 54,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['644: Football players are on the field.',\n",
              " '706: The crowd is watching the football at the game',\n",
              " '1444: position in football played by a team member',\n",
              " '1339: Two groups of people are playing football']"
            ]
          },
          "metadata": {},
          "execution_count": 54
        }
      ]
    }
  ]
}