notebooks/ann_benchmark_recall.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a36fa029-39d5-422b-bd37-0c233cf87022",
   "metadata": {
    "papermill": {
     "duration": 0.011083,
     "end_time": "2022-11-02T14:10:23.656856",
     "exception": false,
     "start_time": "2022-11-02T14:10:23.645773",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "https://github.com/facebookresearch/faiss/blob/main/tutorial/python/2-IVFFlat.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6c3c499-5ed0-4cf7-a8f7-8a0686cf5256",
   "metadata": {
    "papermill": {
     "duration": 0.007661,
     "end_time": "2022-11-02T14:10:23.672386",
     "exception": false,
     "start_time": "2022-11-02T14:10:23.664725",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e238e21-fa45-478c-8695-15965df48263",
   "metadata": {
    "papermill": {
     "duration": 0.007533,
     "end_time": "2022-11-02T14:10:23.687399",
     "exception": false,
     "start_time": "2022-11-02T14:10:23.679866",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "16791642-b297-4997-a428-72ce05a79dd7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:23.703805Z",
     "iopub.status.busy": "2022-11-02T14:10:23.703478Z",
     "iopub.status.idle": "2022-11-02T14:10:23.711737Z",
     "shell.execute_reply": "2022-11-02T14:10:23.711086Z"
    },
    "papermill": {
     "duration": 0.01874,
     "end_time": "2022-11-02T14:10:23.713646",
     "exception": false,
     "start_time": "2022-11-02T14:10:23.694906",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import metrics_utils\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "208c2e9e-f948-479a-b51e-836722ef73b7",
   "metadata": {
    "papermill": {
     "duration": 0.007567,
     "end_time": "2022-11-02T14:10:23.728809",
     "exception": false,
     "start_time": "2022-11-02T14:10:23.721242",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3dfc722e-0e18-473c-afc2-db2948fa270e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:23.745239Z",
     "iopub.status.busy": "2022-11-02T14:10:23.745024Z",
     "iopub.status.idle": "2022-11-02T14:10:23.766352Z",
     "shell.execute_reply": "2022-11-02T14:10:23.765717Z"
    },
    "papermill": {
     "duration": 0.031501,
     "end_time": "2022-11-02T14:10:23.767992",
     "exception": false,
     "start_time": "2022-11-02T14:10:23.736491",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import faiss\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "045e013e-b55f-48ac-950f-35b8c295b901",
   "metadata": {
    "papermill": {
     "duration": 0.008046,
     "end_time": "2022-11-02T14:10:23.783674",
     "exception": false,
     "start_time": "2022-11-02T14:10:23.775628",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "61eccea2-da34-471f-993a-4a7ebfbc9cc9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:23.799903Z",
     "iopub.status.busy": "2022-11-02T14:10:23.799377Z",
     "iopub.status.idle": "2022-11-02T14:10:36.938629Z",
     "shell.execute_reply": "2022-11-02T14:10:36.937787Z"
    },
    "papermill": {
     "duration": 13.14954,
     "end_time": "2022-11-02T14:10:36.940625",
     "exception": false,
     "start_time": "2022-11-02T14:10:23.791085",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = pd.read_parquet(\"gs://np-public-training-temp/stackoverflow/final_subset/posts_with_embedding.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fc7f376-8542-4f46-b918-6fa16e0b8262",
   "metadata": {
    "papermill": {
     "duration": 0.007905,
     "end_time": "2022-11-02T14:10:36.956392",
     "exception": false,
     "start_time": "2022-11-02T14:10:36.948487",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "202ec369-cb53-45af-b54a-1f4dce4590a6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:36.973259Z",
     "iopub.status.busy": "2022-11-02T14:10:36.972797Z",
     "iopub.status.idle": "2022-11-02T14:10:36.980030Z",
     "shell.execute_reply": "2022-11-02T14:10:36.979336Z"
    },
    "papermill": {
     "duration": 0.017692,
     "end_time": "2022-11-02T14:10:36.981697",
     "exception": false,
     "start_time": "2022-11-02T14:10:36.964005",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "219841"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "24654d4a-ba1d-4d6f-8c38-575445a7e398",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:36.998589Z",
     "iopub.status.busy": "2022-11-02T14:10:36.998359Z",
     "iopub.status.idle": "2022-11-02T14:10:37.017916Z",
     "shell.execute_reply": "2022-11-02T14:10:37.017308Z"
    },
    "papermill": {
     "duration": 0.029574,
     "end_time": "2022-11-02T14:10:37.019337",
     "exception": false,
     "start_time": "2022-11-02T14:10:36.989763",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Id</th>\n",
       "      <th>AcceptedAnswerId</th>\n",
       "      <th>Title</th>\n",
       "      <th>QuestionBody</th>\n",
       "      <th>Tags</th>\n",
       "      <th>ViewCount</th>\n",
       "      <th>AnswerCount</th>\n",
       "      <th>CommentCount</th>\n",
       "      <th>Score</th>\n",
       "      <th>CreationDate</th>\n",
       "      <th>AnswerId</th>\n",
       "      <th>AcceptedAnswerBody</th>\n",
       "      <th>embedding</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>15020895</td>\n",
       "      <td>-1</td>\n",
       "      <td>Python int-byte efficient data structure</td>\n",
       "      <td>i am currently storing key-values of type int-...</td>\n",
       "      <td>python data-structures</td>\n",
       "      <td>155</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>2013-02-22 09:33:26.360</td>\n",
       "      <td>-1</td>\n",
       "      <td></td>\n",
       "      <td>[-0.012659057043492794, -0.00613340875133872, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>68487902</td>\n",
       "      <td>-1</td>\n",
       "      <td>Why does the Variance of Laplace very differen...</td>\n",
       "      <td>TL;DR: How can I use skimage.filters.laplace(i...</td>\n",
       "      <td>python opencv image-processing computer-vision...</td>\n",
       "      <td>391</td>\n",
       "      <td>0</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>2021-07-22 15:50:34.220</td>\n",
       "      <td>-1</td>\n",
       "      <td></td>\n",
       "      <td>[0.06565584242343903, -0.031634315848350525, 0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>61391327</td>\n",
       "      <td>-1</td>\n",
       "      <td>Why input never ends</td>\n",
       "      <td>I have python 3.7 installed and I have this co...</td>\n",
       "      <td>python python-3.x input</td>\n",
       "      <td>104</td>\n",
       "      <td>1</td>\n",
       "      <td>6</td>\n",
       "      <td>3</td>\n",
       "      <td>2020-04-23 15:43:03.497</td>\n",
       "      <td>-1</td>\n",
       "      <td></td>\n",
       "      <td>[0.01758185401558876, -0.07531370967626572, 0....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>28852710</td>\n",
       "      <td>-1</td>\n",
       "      <td>Crashes with piecewise linear objective for gu...</td>\n",
       "      <td>We have a complex optimization problem which i...</td>\n",
       "      <td>python crash gurobi piecewise</td>\n",
       "      <td>403</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>2015-03-04 10:58:16.370</td>\n",
       "      <td>-1</td>\n",
       "      <td></td>\n",
       "      <td>[0.010070848278701305, 0.06428138166666031, -0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>24043029</td>\n",
       "      <td>-1</td>\n",
       "      <td>Python TypeError: plotdatehist() got an unexpe...</td>\n",
       "      <td>apologies beforehand if this is a stupid quest...</td>\n",
       "      <td>python typeerror</td>\n",
       "      <td>419</td>\n",
       "      <td>0</td>\n",
       "      <td>7</td>\n",
       "      <td>0</td>\n",
       "      <td>2014-06-04 16:42:32.257</td>\n",
       "      <td>-1</td>\n",
       "      <td></td>\n",
       "      <td>[0.034622225910425186, 0.0662810429930687, 0.0...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         Id  AcceptedAnswerId  \\\n",
       "0  15020895                -1   \n",
       "1  68487902                -1   \n",
       "2  61391327                -1   \n",
       "3  28852710                -1   \n",
       "4  24043029                -1   \n",
       "\n",
       "                                               Title  \\\n",
       "0           Python int-byte efficient data structure   \n",
       "1  Why does the Variance of Laplace very differen...   \n",
       "2                               Why input never ends   \n",
       "3  Crashes with piecewise linear objective for gu...   \n",
       "4  Python TypeError: plotdatehist() got an unexpe...   \n",
       "\n",
       "                                        QuestionBody  \\\n",
       "0  i am currently storing key-values of type int-...   \n",
       "1  TL;DR: How can I use skimage.filters.laplace(i...   \n",
       "2  I have python 3.7 installed and I have this co...   \n",
       "3  We have a complex optimization problem which i...   \n",
       "4  apologies beforehand if this is a stupid quest...   \n",
       "\n",
       "                                                Tags  ViewCount  AnswerCount  \\\n",
       "0                             python data-structures        155            0   \n",
       "1  python opencv image-processing computer-vision...        391            0   \n",
       "2                            python python-3.x input        104            1   \n",
       "3                      python crash gurobi piecewise        403            1   \n",
       "4                                   python typeerror        419            0   \n",
       "\n",
       "   CommentCount  Score            CreationDate  AnswerId AcceptedAnswerBody  \\\n",
       "0             3      1 2013-02-22 09:33:26.360        -1                      \n",
       "1             5      1 2021-07-22 15:50:34.220        -1                      \n",
       "2             6      3 2020-04-23 15:43:03.497        -1                      \n",
       "3             1      3 2015-03-04 10:58:16.370        -1                      \n",
       "4             7      0 2014-06-04 16:42:32.257        -1                      \n",
       "\n",
       "                                           embedding  \n",
       "0  [-0.012659057043492794, -0.00613340875133872, ...  \n",
       "1  [0.06565584242343903, -0.031634315848350525, 0...  \n",
       "2  [0.01758185401558876, -0.07531370967626572, 0....  \n",
       "3  [0.010070848278701305, 0.06428138166666031, -0...  \n",
       "4  [0.034622225910425186, 0.0662810429930687, 0.0...  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a522ab60-1b85-4221-b88d-193d908a64f3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:37.036068Z",
     "iopub.status.busy": "2022-11-02T14:10:37.035818Z",
     "iopub.status.idle": "2022-11-02T14:10:37.146501Z",
     "shell.execute_reply": "2022-11-02T14:10:37.145782Z"
    },
    "papermill": {
     "duration": 0.121486,
     "end_time": "2022-11-02T14:10:37.148645",
     "exception": false,
     "start_time": "2022-11-02T14:10:37.027159",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_subset = df.sample(frac=1).head(500_000).reset_index(drop=True) \n",
    "df_subset['id_faiss'] = range(0, 0+len(df_subset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3fb3f0e2-20b2-4a5b-8b4a-f84e366d3b39",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:37.167408Z",
     "iopub.status.busy": "2022-11-02T14:10:37.166843Z",
     "iopub.status.idle": "2022-11-02T14:10:37.194086Z",
     "shell.execute_reply": "2022-11-02T14:10:37.193488Z"
    },
    "papermill": {
     "duration": 0.038299,
     "end_time": "2022-11-02T14:10:37.195630",
     "exception": false,
     "start_time": "2022-11-02T14:10:37.157331",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Id</th>\n",
       "      <th>AcceptedAnswerId</th>\n",
       "      <th>Title</th>\n",
       "      <th>QuestionBody</th>\n",
       "      <th>Tags</th>\n",
       "      <th>ViewCount</th>\n",
       "      <th>AnswerCount</th>\n",
       "      <th>CommentCount</th>\n",
       "      <th>Score</th>\n",
       "      <th>CreationDate</th>\n",
       "      <th>AnswerId</th>\n",
       "      <th>AcceptedAnswerBody</th>\n",
       "      <th>embedding</th>\n",
       "      <th>id_faiss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>71957513</td>\n",
       "      <td>-1</td>\n",
       "      <td>unit tests for a class containing methods doin...</td>\n",
       "      <td>I have a python class where I have a quicksigh...</td>\n",
       "      <td>python unit-testing mocking python-unittest py...</td>\n",
       "      <td>19</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2022-04-21 16:01:05.843</td>\n",
       "      <td>-1</td>\n",
       "      <td></td>\n",
       "      <td>[0.01994287222623825, 0.13373127579689026, 0.0...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>53981547</td>\n",
       "      <td>53981930</td>\n",
       "      <td>Calculate local clustering coefficient of a ve...</td>\n",
       "      <td>\\n\\nI found an example showing how to calculat...</td>\n",
       "      <td>r cluster-analysis igraph graph-theory</td>\n",
       "      <td>2255</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>2018-12-30 21:23:54.897</td>\n",
       "      <td>53981930</td>\n",
       "      <td>All of this can be done in igraph. It is nice ...</td>\n",
       "      <td>[-0.017001142725348473, 0.03255533054471016, -...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>60900596</td>\n",
       "      <td>-1</td>\n",
       "      <td>ModuleNotFoundError: No module named 'todo_lis...</td>\n",
       "      <td>I am following a tutorial online for a to-do a...</td>\n",
       "      <td>python django</td>\n",
       "      <td>777</td>\n",
       "      <td>1</td>\n",
       "      <td>6</td>\n",
       "      <td>0</td>\n",
       "      <td>2020-03-28 11:35:05.110</td>\n",
       "      <td>-1</td>\n",
       "      <td></td>\n",
       "      <td>[0.007443961221724749, -0.0237520020455122, 0....</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>34090439</td>\n",
       "      <td>34093079</td>\n",
       "      <td>index.html served by django cant locate socket...</td>\n",
       "      <td>I'm trying to add realtime feature to a page s...</td>\n",
       "      <td>javascript python django node.js sockets</td>\n",
       "      <td>207</td>\n",
       "      <td>1</td>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "      <td>2015-12-04 14:20:57.697</td>\n",
       "      <td>34093079</td>\n",
       "      <td>Since you're running Node.js on a separate URL...</td>\n",
       "      <td>[0.028874576091766357, -0.06416145712137222, 0...</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>70950415</td>\n",
       "      <td>-1</td>\n",
       "      <td>How to calculate gaussian bins manually?</td>\n",
       "      <td>A bin—sometimes called a class interval. I hav...</td>\n",
       "      <td>python gaussian</td>\n",
       "      <td>17</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>2022-02-02 04:11:16.967</td>\n",
       "      <td>-1</td>\n",
       "      <td></td>\n",
       "      <td>[-0.030148359015583992, 0.039703480899333954, ...</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>219836</th>\n",
       "      <td>16543573</td>\n",
       "      <td>16553503</td>\n",
       "      <td>How to reuse / clone an sqlalchemy query</td>\n",
       "      <td>It seems to me like going through the whole pr...</td>\n",
       "      <td>python caching sqlalchemy</td>\n",
       "      <td>4538</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>4</td>\n",
       "      <td>2013-05-14 12:46:41.287</td>\n",
       "      <td>16553503</td>\n",
       "      <td>\\n  It seems to me like going through the whol...</td>\n",
       "      <td>[0.042191971093416214, 0.08244752138853073, -0...</td>\n",
       "      <td>219836</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>219837</th>\n",
       "      <td>52081561</td>\n",
       "      <td>52082311</td>\n",
       "      <td>How Pandas read_html can get only selected col...</td>\n",
       "      <td>I'm trying to extract the particular columns f...</td>\n",
       "      <td>python pandas</td>\n",
       "      <td>2640</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>2018-08-29 15:52:14.677</td>\n",
       "      <td>52082311</td>\n",
       "      <td>to select a subset of columns you could use\\n\\...</td>\n",
       "      <td>[0.03314579650759697, 0.12591011822223663, -0....</td>\n",
       "      <td>219837</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>219838</th>\n",
       "      <td>60624604</td>\n",
       "      <td>-1</td>\n",
       "      <td>Python Logic to find the sum of a sub list sat...</td>\n",
       "      <td>I need to find the sum of the list elements be...</td>\n",
       "      <td>python list</td>\n",
       "      <td>144</td>\n",
       "      <td>6</td>\n",
       "      <td>11</td>\n",
       "      <td>1</td>\n",
       "      <td>2020-03-10 19:25:40.903</td>\n",
       "      <td>-1</td>\n",
       "      <td></td>\n",
       "      <td>[-0.0366428941488266, 0.06802389025688171, -0....</td>\n",
       "      <td>219838</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>219839</th>\n",
       "      <td>29354283</td>\n",
       "      <td>-1</td>\n",
       "      <td>Python/Pygame sprites using same class not upd...</td>\n",
       "      <td>[CODE]\\n\\nQuick rundown, shipX should move whe...</td>\n",
       "      <td>python class pygame sprite</td>\n",
       "      <td>240</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>2015-03-30 19:09:49.567</td>\n",
       "      <td>-1</td>\n",
       "      <td></td>\n",
       "      <td>[0.04933905228972435, -0.00944867730140686, 0....</td>\n",
       "      <td>219839</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>219840</th>\n",
       "      <td>59711060</td>\n",
       "      <td>-1</td>\n",
       "      <td>How to do a fuzzy match in PySpark UDF?</td>\n",
       "      <td>I am trying to run following code to generate ...</td>\n",
       "      <td>pyspark pyspark-sql fuzzy-logic fuzzy-comparis...</td>\n",
       "      <td>1802</td>\n",
       "      <td>0</td>\n",
       "      <td>6</td>\n",
       "      <td>1</td>\n",
       "      <td>2020-01-13 05:37:27.557</td>\n",
       "      <td>-1</td>\n",
       "      <td></td>\n",
       "      <td>[-0.010504454374313354, 0.07815464586019516, 0...</td>\n",
       "      <td>219840</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>219841 rows × 14 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "              Id  AcceptedAnswerId  \\\n",
       "0       71957513                -1   \n",
       "1       53981547          53981930   \n",
       "2       60900596                -1   \n",
       "3       34090439          34093079   \n",
       "4       70950415                -1   \n",
       "...          ...               ...   \n",
       "219836  16543573          16553503   \n",
       "219837  52081561          52082311   \n",
       "219838  60624604                -1   \n",
       "219839  29354283                -1   \n",
       "219840  59711060                -1   \n",
       "\n",
       "                                                    Title  \\\n",
       "0       unit tests for a class containing methods doin...   \n",
       "1       Calculate local clustering coefficient of a ve...   \n",
       "2       ModuleNotFoundError: No module named 'todo_lis...   \n",
       "3       index.html served by django cant locate socket...   \n",
       "4                How to calculate gaussian bins manually?   \n",
       "...                                                   ...   \n",
       "219836           How to reuse / clone an sqlalchemy query   \n",
       "219837  How Pandas read_html can get only selected col...   \n",
       "219838  Python Logic to find the sum of a sub list sat...   \n",
       "219839  Python/Pygame sprites using same class not upd...   \n",
       "219840            How to do a fuzzy match in PySpark UDF?   \n",
       "\n",
       "                                             QuestionBody  \\\n",
       "0       I have a python class where I have a quicksigh...   \n",
       "1       \\n\\nI found an example showing how to calculat...   \n",
       "2       I am following a tutorial online for a to-do a...   \n",
       "3       I'm trying to add realtime feature to a page s...   \n",
       "4       A bin—sometimes called a class interval. I hav...   \n",
       "...                                                   ...   \n",
       "219836  It seems to me like going through the whole pr...   \n",
       "219837  I'm trying to extract the particular columns f...   \n",
       "219838  I need to find the sum of the list elements be...   \n",
       "219839  [CODE]\\n\\nQuick rundown, shipX should move whe...   \n",
       "219840  I am trying to run following code to generate ...   \n",
       "\n",
       "                                                     Tags  ViewCount  \\\n",
       "0       python unit-testing mocking python-unittest py...         19   \n",
       "1                  r cluster-analysis igraph graph-theory       2255   \n",
       "2                                           python django        777   \n",
       "3                javascript python django node.js sockets        207   \n",
       "4                                         python gaussian         17   \n",
       "...                                                   ...        ...   \n",
       "219836                          python caching sqlalchemy       4538   \n",
       "219837                                      python pandas       2640   \n",
       "219838                                        python list        144   \n",
       "219839                         python class pygame sprite        240   \n",
       "219840  pyspark pyspark-sql fuzzy-logic fuzzy-comparis...       1802   \n",
       "\n",
       "        AnswerCount  CommentCount  Score            CreationDate  AnswerId  \\\n",
       "0                 0             0      0 2022-04-21 16:01:05.843        -1   \n",
       "1                 1             0      3 2018-12-30 21:23:54.897  53981930   \n",
       "2                 1             6      0 2020-03-28 11:35:05.110        -1   \n",
       "3                 1             4      1 2015-12-04 14:20:57.697  34093079   \n",
       "4                 0             1      0 2022-02-02 04:11:16.967        -1   \n",
       "...             ...           ...    ...                     ...       ...   \n",
       "219836            2             0      4 2013-05-14 12:46:41.287  16553503   \n",
       "219837            2             2      1 2018-08-29 15:52:14.677  52082311   \n",
       "219838            6            11      1 2020-03-10 19:25:40.903        -1   \n",
       "219839            1             0      1 2015-03-30 19:09:49.567        -1   \n",
       "219840            0             6      1 2020-01-13 05:37:27.557        -1   \n",
       "\n",
       "                                       AcceptedAnswerBody  \\\n",
       "0                                                           \n",
       "1       All of this can be done in igraph. It is nice ...   \n",
       "2                                                           \n",
       "3       Since you're running Node.js on a separate URL...   \n",
       "4                                                           \n",
       "...                                                   ...   \n",
       "219836  \\n  It seems to me like going through the whol...   \n",
       "219837  to select a subset of columns you could use\\n\\...   \n",
       "219838                                                      \n",
       "219839                                                      \n",
       "219840                                                      \n",
       "\n",
       "                                                embedding  id_faiss  \n",
       "0       [0.01994287222623825, 0.13373127579689026, 0.0...         0  \n",
       "1       [-0.017001142725348473, 0.03255533054471016, -...         1  \n",
       "2       [0.007443961221724749, -0.0237520020455122, 0....         2  \n",
       "3       [0.028874576091766357, -0.06416145712137222, 0...         3  \n",
       "4       [-0.030148359015583992, 0.039703480899333954, ...         4  \n",
       "...                                                   ...       ...  \n",
       "219836  [0.042191971093416214, 0.08244752138853073, -0...    219836  \n",
       "219837  [0.03314579650759697, 0.12591011822223663, -0....    219837  \n",
       "219838  [-0.0366428941488266, 0.06802389025688171, -0....    219838  \n",
       "219839  [0.04933905228972435, -0.00944867730140686, 0....    219839  \n",
       "219840  [-0.010504454374313354, 0.07815464586019516, 0...    219840  \n",
       "\n",
       "[219841 rows x 14 columns]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19fa87a5-6176-4da8-952c-a16082197c7d",
   "metadata": {
    "papermill": {
     "duration": 0.008617,
     "end_time": "2022-11-02T14:10:37.212752",
     "exception": false,
     "start_time": "2022-11-02T14:10:37.204135",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "76885f88-9d02-4264-9c78-1ef7c86a5f09",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:37.230752Z",
     "iopub.status.busy": "2022-11-02T14:10:37.230532Z",
     "iopub.status.idle": "2022-11-02T14:10:37.233567Z",
     "shell.execute_reply": "2022-11-02T14:10:37.232928Z"
    },
    "papermill": {
     "duration": 0.013906,
     "end_time": "2022-11-02T14:10:37.235128",
     "exception": false,
     "start_time": "2022-11-02T14:10:37.221222",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "k = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fab0ef10-cfa4-4102-9f8b-fe96df8516f6",
   "metadata": {
    "papermill": {
     "duration": 0.009624,
     "end_time": "2022-11-02T14:10:37.253457",
     "exception": false,
     "start_time": "2022-11-02T14:10:37.243833",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "ac1a1c92-67c0-443f-82df-202db85a1649",
   "metadata": {
    "papermill": {
     "duration": 0.008605,
     "end_time": "2022-11-02T14:10:37.271307",
     "exception": false,
     "start_time": "2022-11-02T14:10:37.262702",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "emmbedding corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "160bae62-31a7-4873-938f-824760232e95",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:37.291957Z",
     "iopub.status.busy": "2022-11-02T14:10:37.291700Z",
     "iopub.status.idle": "2022-11-02T14:10:38.121624Z",
     "shell.execute_reply": "2022-11-02T14:10:38.120819Z"
    },
    "papermill": {
     "duration": 0.841732,
     "end_time": "2022-11-02T14:10:38.123708",
     "exception": false,
     "start_time": "2022-11-02T14:10:37.281976",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "xb = df_subset['embedding'].tolist()\n",
    "xb = np.array(xb).astype('float32')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4ec4affe-3618-45cb-bf88-7e99efbb80aa",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:38.142887Z",
     "iopub.status.busy": "2022-11-02T14:10:38.142344Z",
     "iopub.status.idle": "2022-11-02T14:10:38.146849Z",
     "shell.execute_reply": "2022-11-02T14:10:38.146233Z"
    },
    "papermill": {
     "duration": 0.015823,
     "end_time": "2022-11-02T14:10:38.148241",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.132418",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(219841, 768)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xb.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bab05f6e-81f1-46fa-bfc2-a79167f35382",
   "metadata": {
    "papermill": {
     "duration": 0.008511,
     "end_time": "2022-11-02T14:10:38.165571",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.157060",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "d934cb7d-ca2c-4b1d-bf50-c12eda5123cb",
   "metadata": {
    "papermill": {
     "duration": 0.008609,
     "end_time": "2022-11-02T14:10:38.182662",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.174053",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8d558fbe-f0ba-4d98-aa9d-e0523ea20746",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:38.201055Z",
     "iopub.status.busy": "2022-11-02T14:10:38.200400Z",
     "iopub.status.idle": "2022-11-02T14:10:38.204719Z",
     "shell.execute_reply": "2022-11-02T14:10:38.204188Z"
    },
    "papermill": {
     "duration": 0.01527,
     "end_time": "2022-11-02T14:10:38.206268",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.190998",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "768"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d = xb.shape[-1]                       \n",
    "d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72efe73b-50ab-4c40-80de-3048bd61378e",
   "metadata": {
    "papermill": {
     "duration": 0.008459,
     "end_time": "2022-11-02T14:10:38.223164",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.214705",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e8c45c3-b376-4b9f-b9ea-4ddf720265c8",
   "metadata": {
    "papermill": {
     "duration": 0.008447,
     "end_time": "2022-11-02T14:10:38.240062",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.231615",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "5d918715-c288-4829-a60b-7617b18af542",
   "metadata": {
    "papermill": {
     "duration": 0.009296,
     "end_time": "2022-11-02T14:10:38.257737",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.248441",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Naive Index : flat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "db9880c4-9c96-4c16-832e-95ffed09e0c3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:38.276251Z",
     "iopub.status.busy": "2022-11-02T14:10:38.276003Z",
     "iopub.status.idle": "2022-11-02T14:10:38.279268Z",
     "shell.execute_reply": "2022-11-02T14:10:38.278680Z"
    },
    "papermill": {
     "duration": 0.014092,
     "end_time": "2022-11-02T14:10:38.280692",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.266600",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "index = faiss.IndexFlatL2(d)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "040cbda3-421c-4dc1-a1ef-cf1a3acc5dad",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:38.299075Z",
     "iopub.status.busy": "2022-11-02T14:10:38.298426Z",
     "iopub.status.idle": "2022-11-02T14:10:38.302461Z",
     "shell.execute_reply": "2022-11-02T14:10:38.301884Z"
    },
    "papermill": {
     "duration": 0.014768,
     "end_time": "2022-11-02T14:10:38.303890",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.289122",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "index.is_trained\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "541006b8-5d48-40b9-a449-3085a74392fb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:38.322501Z",
     "iopub.status.busy": "2022-11-02T14:10:38.322030Z",
     "iopub.status.idle": "2022-11-02T14:10:38.597411Z",
     "shell.execute_reply": "2022-11-02T14:10:38.596663Z"
    },
    "papermill": {
     "duration": 0.287026,
     "end_time": "2022-11-02T14:10:38.599612",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.312586",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "index.add(xb)                "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8df996c6-34c6-45ab-8aa2-426dd03e47f3",
   "metadata": {
    "papermill": {
     "duration": 0.008718,
     "end_time": "2022-11-02T14:10:38.617634",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.608916",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "vectors to search for"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "acbda0c3-d7e9-4595-abc1-4fc17270cafe",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:38.636754Z",
     "iopub.status.busy": "2022-11-02T14:10:38.636076Z",
     "iopub.status.idle": "2022-11-02T14:10:38.639451Z",
     "shell.execute_reply": "2022-11-02T14:10:38.638860Z"
    },
    "papermill": {
     "duration": 0.014492,
     "end_time": "2022-11-02T14:10:38.640898",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.626406",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "xq = xb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3a595f09-7b18-43e9-9668-619b39350d66",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:38.660127Z",
     "iopub.status.busy": "2022-11-02T14:10:38.659487Z",
     "iopub.status.idle": "2022-11-02T14:10:38.663553Z",
     "shell.execute_reply": "2022-11-02T14:10:38.662990Z"
    },
    "papermill": {
     "duration": 0.015469,
     "end_time": "2022-11-02T14:10:38.664997",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.649528",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "219841"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(xq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "cceeadd5-1f4b-475c-a26c-a27a3f89e9dc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:10:38.728422Z",
     "iopub.status.busy": "2022-11-02T14:10:38.727835Z",
     "iopub.status.idle": "2022-11-02T14:14:50.796108Z",
     "shell.execute_reply": "2022-11-02T14:14:50.795063Z"
    },
    "papermill": {
     "duration": 252.094156,
     "end_time": "2022-11-02T14:14:50.812076",
     "exception": false,
     "start_time": "2022-11-02T14:10:38.717920",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1h 3min, sys: 3min 28s, total: 1h 6min 28s\n",
      "Wall time: 4min 12s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "D, I = index.search(xq, k)     \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "dd78e807-4fe0-459b-978e-dd7b4b106e30",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:14:50.842964Z",
     "iopub.status.busy": "2022-11-02T14:14:50.842581Z",
     "iopub.status.idle": "2022-11-02T14:14:50.882521Z",
     "shell.execute_reply": "2022-11-02T14:14:50.881533Z"
    },
    "papermill": {
     "duration": 0.057739,
     "end_time": "2022-11-02T14:14:50.884610",
     "exception": false,
     "start_time": "2022-11-02T14:14:50.826871",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Title</th>\n",
       "      <th>id_faiss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>unit tests for a class containing methods doin...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Calculate local clustering coefficient of a ve...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>ModuleNotFoundError: No module named 'todo_lis...</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>index.html served by django cant locate socket...</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>How to calculate gaussian bins manually?</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>219836</th>\n",
       "      <td>How to reuse / clone an sqlalchemy query</td>\n",
       "      <td>219836</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>219837</th>\n",
       "      <td>How Pandas read_html can get only selected col...</td>\n",
       "      <td>219837</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>219838</th>\n",
       "      <td>Python Logic to find the sum of a sub list sat...</td>\n",
       "      <td>219838</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>219839</th>\n",
       "      <td>Python/Pygame sprites using same class not upd...</td>\n",
       "      <td>219839</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>219840</th>\n",
       "      <td>How to do a fuzzy match in PySpark UDF?</td>\n",
       "      <td>219840</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>219841 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                    Title  id_faiss\n",
       "0       unit tests for a class containing methods doin...         0\n",
       "1       Calculate local clustering coefficient of a ve...         1\n",
       "2       ModuleNotFoundError: No module named 'todo_lis...         2\n",
       "3       index.html served by django cant locate socket...         3\n",
       "4                How to calculate gaussian bins manually?         4\n",
       "...                                                   ...       ...\n",
       "219836           How to reuse / clone an sqlalchemy query    219836\n",
       "219837  How Pandas read_html can get only selected col...    219837\n",
       "219838  Python Logic to find the sum of a sub list sat...    219838\n",
       "219839  Python/Pygame sprites using same class not upd...    219839\n",
       "219840            How to do a fuzzy match in PySpark UDF?    219840\n",
       "\n",
       "[219841 rows x 2 columns]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_res = df_subset[['Title','id_faiss']].copy()\n",
    "df_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5aefe51a-16b5-43a0-a4ac-5199c3907943",
   "metadata": {
    "papermill": {
     "duration": 0.014264,
     "end_time": "2022-11-02T14:14:50.913529",
     "exception": false,
     "start_time": "2022-11-02T14:14:50.899265",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "2bfbc1a3-8423-4adc-a934-35c15e2d2690",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:14:50.942912Z",
     "iopub.status.busy": "2022-11-02T14:14:50.942574Z",
     "iopub.status.idle": "2022-11-02T14:14:50.948482Z",
     "shell.execute_reply": "2022-11-02T14:14:50.947623Z"
    },
    "papermill": {
     "duration": 0.022563,
     "end_time": "2022-11-02T14:14:50.950404",
     "exception": false,
     "start_time": "2022-11-02T14:14:50.927841",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def calculate_recall( index, target, k):\n",
    "    \n",
    "    res = index == target\n",
    "    \n",
    "    res = res[:k]\n",
    "    \n",
    "    return True in res\n",
    "\n",
    "\n",
    "def compute_recall( I, df, k):\n",
    "    \n",
    "    all_res = []\n",
    "    for idx, row in enumerate(I):\n",
    "        res = calculate_recall(row, idx, k)\n",
    "        all_res.append(res)\n",
    "        \n",
    "        \n",
    "    \n",
    "    \n",
    "    return all_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e458459b-ea8b-4d71-b13c-7bfeeddf03a8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:14:50.979521Z",
     "iopub.status.busy": "2022-11-02T14:14:50.979211Z",
     "iopub.status.idle": "2022-11-02T14:14:51.881330Z",
     "shell.execute_reply": "2022-11-02T14:14:51.880664Z"
    },
    "papermill": {
     "duration": 0.919186,
     "end_time": "2022-11-02T14:14:51.882976",
     "exception": false,
     "start_time": "2022-11-02T14:14:50.963790",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(219841, 219817)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#df_res['r@1'] = compute_recall(I,df_subset,k=1)\n",
    "\n",
    "res =   compute_recall(I,df_subset,k=1)\n",
    "\n",
    "len(res), sum(res)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "35c5c97a-8580-4fae-8c54-f7e8dec27afc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:14:51.903406Z",
     "iopub.status.busy": "2022-11-02T14:14:51.902863Z",
     "iopub.status.idle": "2022-11-02T14:14:51.906138Z",
     "shell.execute_reply": "2022-11-02T14:14:51.905542Z"
    },
    "papermill": {
     "duration": 0.015018,
     "end_time": "2022-11-02T14:14:51.907539",
     "exception": false,
     "start_time": "2022-11-02T14:14:51.892521",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "#calculate_recall (I[1] , 1 , k =1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "1bcc5c9d-0f78-4aed-a19c-d4c4d182b68f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:14:51.927278Z",
     "iopub.status.busy": "2022-11-02T14:14:51.927048Z",
     "iopub.status.idle": "2022-11-02T14:14:51.930151Z",
     "shell.execute_reply": "2022-11-02T14:14:51.929562Z"
    },
    "papermill": {
     "duration": 0.014913,
     "end_time": "2022-11-02T14:14:51.931751",
     "exception": false,
     "start_time": "2022-11-02T14:14:51.916838",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "#df_res['r@1'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16d3ac54-c44a-44ee-bb50-8902704dd98a",
   "metadata": {
    "papermill": {
     "duration": 0.009129,
     "end_time": "2022-11-02T14:14:51.950097",
     "exception": false,
     "start_time": "2022-11-02T14:14:51.940968",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71c413b6-d7b5-4359-9904-ad2a52d52b88",
   "metadata": {
    "papermill": {
     "duration": 0.009738,
     "end_time": "2022-11-02T14:14:51.987325",
     "exception": false,
     "start_time": "2022-11-02T14:14:51.977587",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "a4e56693-d85c-4331-b6bd-dd6166b22879",
   "metadata": {
    "papermill": {
     "duration": 0.009539,
     "end_time": "2022-11-02T14:14:52.006863",
     "exception": false,
     "start_time": "2022-11-02T14:14:51.997324",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## IVF Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cace3a29-339a-49a2-bd42-54de686d9fa8",
   "metadata": {
    "papermill": {
     "duration": 0.009792,
     "end_time": "2022-11-02T14:14:52.026144",
     "exception": false,
     "start_time": "2022-11-02T14:14:52.016352",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "027ed087-5bc7-48e0-aace-1cdde3178801",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:14:52.046968Z",
     "iopub.status.busy": "2022-11-02T14:14:52.046490Z",
     "iopub.status.idle": "2022-11-02T14:14:52.051918Z",
     "shell.execute_reply": "2022-11-02T14:14:52.051273Z"
    },
    "papermill": {
     "duration": 0.017556,
     "end_time": "2022-11-02T14:14:52.053461",
     "exception": false,
     "start_time": "2022-11-02T14:14:52.035905",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "nlist = 100 # nmber of clusters\n",
    "\n",
    "quantizer = faiss.IndexFlatL2(d)  # the other index\n",
    "index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "234fefeb-2e8c-4666-a3ea-70945b2aeeb3",
   "metadata": {
    "papermill": {
     "duration": 0.009475,
     "end_time": "2022-11-02T14:14:52.072208",
     "exception": false,
     "start_time": "2022-11-02T14:14:52.062733",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "73f63b7c-671e-4fe9-9863-49bdfe046b42",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:14:52.092826Z",
     "iopub.status.busy": "2022-11-02T14:14:52.092322Z",
     "iopub.status.idle": "2022-11-02T14:14:52.463505Z",
     "shell.execute_reply": "2022-11-02T14:14:52.462627Z"
    },
    "papermill": {
     "duration": 0.383711,
     "end_time": "2022-11-02T14:14:52.465687",
     "exception": false,
     "start_time": "2022-11-02T14:14:52.081976",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert not index.is_trained\n",
    "index.train(xb)\n",
    "assert index.is_trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "33f5ed61-65de-47ae-8fb5-e6f53b74c6ee",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:14:52.501567Z",
     "iopub.status.busy": "2022-11-02T14:14:52.501159Z",
     "iopub.status.idle": "2022-11-02T14:14:52.505293Z",
     "shell.execute_reply": "2022-11-02T14:14:52.504366Z"
    },
    "papermill": {
     "duration": 0.023162,
     "end_time": "2022-11-02T14:14:52.507362",
     "exception": false,
     "start_time": "2022-11-02T14:14:52.484200",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "xq = xb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "124caaf2-e6a9-4bd3-add9-b73e0ba46d4d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:14:52.537516Z",
     "iopub.status.busy": "2022-11-02T14:14:52.537205Z",
     "iopub.status.idle": "2022-11-02T14:14:52.913657Z",
     "shell.execute_reply": "2022-11-02T14:14:52.912833Z"
    },
    "papermill": {
     "duration": 0.393884,
     "end_time": "2022-11-02T14:14:52.915856",
     "exception": false,
     "start_time": "2022-11-02T14:14:52.521972",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "index.add(xb)         \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "1150e300-9360-4b2d-84bc-24a9430f6736",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:14:52.947802Z",
     "iopub.status.busy": "2022-11-02T14:14:52.947512Z",
     "iopub.status.idle": "2022-11-02T14:15:31.774241Z",
     "shell.execute_reply": "2022-11-02T14:15:31.773418Z"
    },
    "papermill": {
     "duration": 38.854458,
     "end_time": "2022-11-02T14:15:31.785731",
     "exception": false,
     "start_time": "2022-11-02T14:14:52.931273",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 10min 6s, sys: 290 ms, total: 10min 6s\n",
      "Wall time: 38.8 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "\n",
    "index.nprobe = 1              # default nprobe is 1\n",
    "\n",
    "D, I = index.search(xq, k)     # actual search\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "aa3662d4-2a1e-4a88-8fc7-bbe0173c3c89",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:15:31.806887Z",
     "iopub.status.busy": "2022-11-02T14:15:31.806263Z",
     "iopub.status.idle": "2022-11-02T14:15:32.636181Z",
     "shell.execute_reply": "2022-11-02T14:15:32.635503Z"
    },
    "papermill": {
     "duration": 0.842367,
     "end_time": "2022-11-02T14:15:32.637812",
     "exception": false,
     "start_time": "2022-11-02T14:15:31.795445",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True     219835\n",
       "False         6\n",
       "Name: r@1, dtype: int64"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_res['r@1'] = compute_recall(I,df_subset,k=1)\n",
    "df_res['r@1'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adfa6b87-34cc-44bc-bd43-67ff84de94bd",
   "metadata": {
    "papermill": {
     "duration": 0.009411,
     "end_time": "2022-11-02T14:15:32.657016",
     "exception": false,
     "start_time": "2022-11-02T14:15:32.647605",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "626639b7-7c37-4d4d-a1ff-d6d83b3ee238",
   "metadata": {
    "papermill": {
     "duration": 0.009273,
     "end_time": "2022-11-02T14:15:32.675429",
     "exception": false,
     "start_time": "2022-11-02T14:15:32.666156",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f71dc09-efc8-4121-8b85-50f1df45393e",
   "metadata": {
    "papermill": {
     "duration": 0.009307,
     "end_time": "2022-11-02T14:15:32.694080",
     "exception": false,
     "start_time": "2022-11-02T14:15:32.684773",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "b1a051f3-962d-4e50-8381-a241e2989470",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:15:32.714194Z",
     "iopub.status.busy": "2022-11-02T14:15:32.713900Z",
     "iopub.status.idle": "2022-11-02T14:18:35.932110Z",
     "shell.execute_reply": "2022-11-02T14:18:35.931369Z"
    },
    "papermill": {
     "duration": 183.240166,
     "end_time": "2022-11-02T14:18:35.943358",
     "exception": false,
     "start_time": "2022-11-02T14:15:32.703192",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 47min 58s, sys: 1.01 s, total: 47min 59s\n",
      "Wall time: 3min 3s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "index.nprobe = 5              # default nprobe is 1, max is nlist\n",
    "D, I = index.search(xq, k)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "8f1aa473-1a4d-48d1-9be8-10fec9daedae",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:18:35.964456Z",
     "iopub.status.busy": "2022-11-02T14:18:35.964199Z",
     "iopub.status.idle": "2022-11-02T14:18:36.853809Z",
     "shell.execute_reply": "2022-11-02T14:18:36.852909Z"
    },
    "papermill": {
     "duration": 0.902097,
     "end_time": "2022-11-02T14:18:36.855405",
     "exception": false,
     "start_time": "2022-11-02T14:18:35.953308",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True     219835\n",
       "False         6\n",
       "Name: r@1, dtype: int64"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_res['r@1'] = compute_recall(I,df_subset,k=1)\n",
    "df_res['r@1'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a568a0bc-4765-4e86-99e0-f1e52096edee",
   "metadata": {
    "papermill": {
     "duration": 0.010306,
     "end_time": "2022-11-02T14:18:36.875883",
     "exception": false,
     "start_time": "2022-11-02T14:18:36.865577",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "b0eccc84-9f04-4d94-8efb-f9ff7342f84b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:18:36.897791Z",
     "iopub.status.busy": "2022-11-02T14:18:36.896872Z",
     "iopub.status.idle": "2022-11-02T14:18:42.559795Z",
     "shell.execute_reply": "2022-11-02T14:18:42.558765Z"
    },
    "papermill": {
     "duration": 5.676196,
     "end_time": "2022-11-02T14:18:42.562042",
     "exception": false,
     "start_time": "2022-11-02T14:18:36.885846",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "faiss.write_index(index, \"/tmp/index_ivf.bin\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1f83905-cdf4-41bc-a922-00502aad1603",
   "metadata": {
    "papermill": {
     "duration": 0.067829,
     "end_time": "2022-11-02T14:18:42.640987",
     "exception": false,
     "start_time": "2022-11-02T14:18:42.573158",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### PQ Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "7efccc29-8ca1-462e-bd46-a779332c4329",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:18:42.662788Z",
     "iopub.status.busy": "2022-11-02T14:18:42.662126Z",
     "iopub.status.idle": "2022-11-02T14:18:42.668858Z",
     "shell.execute_reply": "2022-11-02T14:18:42.668212Z"
    },
    "papermill": {
     "duration": 0.01916,
     "end_time": "2022-11-02T14:18:42.670309",
     "exception": false,
     "start_time": "2022-11-02T14:18:42.651149",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "m = 8 # how many subchunks for the vector\n",
    "nbits = 8  # number of bits per subquantizer, k* = 2**nbits ; number of cluster_ids \n",
    "index = faiss.IndexPQ(d, m, nbits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "2f98fbf8-ab16-4171-a44d-bb45c3d6c7a3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:18:42.691696Z",
     "iopub.status.busy": "2022-11-02T14:18:42.691456Z",
     "iopub.status.idle": "2022-11-02T14:18:47.068984Z",
     "shell.execute_reply": "2022-11-02T14:18:47.068106Z"
    },
    "papermill": {
     "duration": 4.391081,
     "end_time": "2022-11-02T14:18:47.071388",
     "exception": false,
     "start_time": "2022-11-02T14:18:42.680307",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert not index.is_trained\n",
    "index.train(xb)\n",
    "assert index.is_trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "a3516271-9e4d-4a5d-b0ca-34eafba0f45d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:18:47.104597Z",
     "iopub.status.busy": "2022-11-02T14:18:47.104272Z",
     "iopub.status.idle": "2022-11-02T14:18:47.590395Z",
     "shell.execute_reply": "2022-11-02T14:18:47.589565Z"
    },
    "papermill": {
     "duration": 0.504464,
     "end_time": "2022-11-02T14:18:47.592627",
     "exception": false,
     "start_time": "2022-11-02T14:18:47.088163",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "D, I = index.search(xq, k)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "a10f8a23-bf41-481b-b858-3ef8a2e73d00",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:18:47.626493Z",
     "iopub.status.busy": "2022-11-02T14:18:47.626158Z",
     "iopub.status.idle": "2022-11-02T14:18:48.468960Z",
     "shell.execute_reply": "2022-11-02T14:18:48.468150Z"
    },
    "papermill": {
     "duration": 0.860948,
     "end_time": "2022-11-02T14:18:48.470752",
     "exception": false,
     "start_time": "2022-11-02T14:18:47.609804",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(219841, 0)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res =   compute_recall(I,df_subset,k=1)\n",
    "\n",
    "len(res), sum(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "9b72bb5f-9592-4806-bbe0-43aacd072311",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:18:48.492628Z",
     "iopub.status.busy": "2022-11-02T14:18:48.491949Z",
     "iopub.status.idle": "2022-11-02T14:18:48.505538Z",
     "shell.execute_reply": "2022-11-02T14:18:48.504864Z"
    },
    "papermill": {
     "duration": 0.026127,
     "end_time": "2022-11-02T14:18:48.507031",
     "exception": false,
     "start_time": "2022-11-02T14:18:48.480904",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " False,\n",
       " ...]"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "6ca5cc1b-ed42-4a47-9b9d-5ba520c703f2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:18:48.528618Z",
     "iopub.status.busy": "2022-11-02T14:18:48.528360Z",
     "iopub.status.idle": "2022-11-02T14:18:48.532139Z",
     "shell.execute_reply": "2022-11-02T14:18:48.531429Z"
    },
    "papermill": {
     "duration": 0.016359,
     "end_time": "2022-11-02T14:18:48.533609",
     "exception": false,
     "start_time": "2022-11-02T14:18:48.517250",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "faiss.write_index(index, \"/tmp/index_pq.bin\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80164fa3-90b5-4631-849a-b44cc2549389",
   "metadata": {
    "papermill": {
     "duration": 0.011235,
     "end_time": "2022-11-02T14:18:48.555149",
     "exception": false,
     "start_time": "2022-11-02T14:18:48.543914",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dd63490-48a0-4fe4-b149-366d6e5f576e",
   "metadata": {
    "papermill": {
     "duration": 0.010109,
     "end_time": "2022-11-02T14:18:48.575790",
     "exception": false,
     "start_time": "2022-11-02T14:18:48.565681",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "35ea8304-ff26-43cf-83d6-9e764453aea4",
   "metadata": {
    "papermill": {
     "duration": 0.010093,
     "end_time": "2022-11-02T14:18:48.596783",
     "exception": false,
     "start_time": "2022-11-02T14:18:48.586690",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Space consumed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "5fce6da3-7ac9-4371-a307-53672c76c05d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-02T14:18:48.619005Z",
     "iopub.status.busy": "2022-11-02T14:18:48.618246Z",
     "iopub.status.idle": "2022-11-02T14:18:48.783517Z",
     "shell.execute_reply": "2022-11-02T14:18:48.782538Z"
    },
    "papermill": {
     "duration": 0.178583,
     "end_time": "2022-11-02T14:18:48.785463",
     "exception": false,
     "start_time": "2022-11-02T14:18:48.606880",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-rw-r--r-- 1 jupyter jupyter 647M Nov  2 14:18 /tmp/index_ivf.bin\n",
      "-rw-r--r-- 1 jupyter jupyter 769K Nov  2 14:18 /tmp/index_pq.bin\n"
     ]
    }
   ],
   "source": [
    "!ls -lah /tmp/index_*"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "pytorch-gpu.1-12.m99",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/pytorch-gpu.1-12:m99"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 508.227095,
   "end_time": "2022-11-02T14:18:49.515814",
   "environment_variables": {},
   "exception": null,
   "input_path": "ann_benchmark_recall.ipynb",
   "output_path": "ann_benchmark_recall__out.ipynb",
   "parameters": {},
   "start_time": "2022-11-02T14:10:21.288719",
   "version": "2.3.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}