notebooks/04_ann.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d6298118-b5f8-4250-bd82-e2a3787914ca",
   "metadata": {},
   "source": [
    "# Benchmarking Aproximate Nearest Neighbors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2af1fec-6519-40d0-8826-f201d0acba0b",
   "metadata": {},
   "source": [
    "# About"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a45ed78-348b-4193-8870-cbcbaff240e1",
   "metadata": {},
   "source": [
    "In order for embedding retrieval to work at scale, need to use a vector database.\n",
    "We also need to use Approximate Nearest Search instead of brute force.\n",
    "\n",
    "\n",
    "In this notebook, we will use [FAISS]() a library from facebook.\n",
    "\n",
    "We will compare a brute force and the speedup gained from `IVF`.\n",
    "\n",
    "For a more detailed comparision, take a look <a href=\"http://ann-benchmarks.com/\">here</a> to find other solutions and benchmark data.\n",
    "\n",
    "\n",
    "We will look at `performance` and `recall@1`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e45e627d-9c8c-4f19-ab91-0e64ed8677d7",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "62298ef5-fa20-4164-8ec7-e7f43bf85c20",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:36.834089Z",
     "iopub.status.busy": "2023-04-26T14:52:36.833681Z",
     "iopub.status.idle": "2023-04-26T14:52:37.774946Z",
     "shell.execute_reply": "2023-04-26T14:52:37.774030Z",
     "shell.execute_reply.started": "2023-04-26T14:52:36.834036Z"
    }
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import faiss\n",
    "import datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b9a73d2-aa07-46ef-b94d-8780ca9ecb68",
   "metadata": {},
   "source": [
    "## Load the embeddings of the image corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "74c3f8aa-cb27-4ae8-a9a8-0060001d357c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:37.778006Z",
     "iopub.status.busy": "2023-04-26T14:52:37.776915Z",
     "iopub.status.idle": "2023-04-26T14:52:44.572677Z",
     "shell.execute_reply": "2023-04-26T14:52:44.571814Z",
     "shell.execute_reply.started": "2023-04-26T14:52:37.777973Z"
    }
   },
   "outputs": [],
   "source": [
    "dset = datasets.load_from_disk(\"../data/processed_embeddings\")\n",
    "## these embeddings will be used to create the search space.\n",
    "corpus = dset['embeddings']\n",
    "\n",
    "\n",
    "corpus = np.array(corpus).astype('float32')\n",
    "corpus = np.unique(corpus, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "94eeb2f5-88ff-428f-8e9c-5014234427b8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:44.574012Z",
     "iopub.status.busy": "2023-04-26T14:52:44.573742Z",
     "iopub.status.idle": "2023-04-26T14:52:44.580175Z",
     "shell.execute_reply": "2023-04-26T14:52:44.579367Z",
     "shell.execute_reply.started": "2023-04-26T14:52:44.573987Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(24954, 512)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "corpus.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "939108fa-daf9-472c-b026-19c6d9708a77",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:44.581364Z",
     "iopub.status.busy": "2023-04-26T14:52:44.581116Z",
     "iopub.status.idle": "2023-04-26T14:52:44.590669Z",
     "shell.execute_reply": "2023-04-26T14:52:44.589884Z",
     "shell.execute_reply.started": "2023-04-26T14:52:44.581340Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-0.08344752,  0.01604629,  0.03037108, ...,  0.03962855,\n",
       "        -0.02023211, -0.01102281],\n",
       "       [-0.07890625,  0.02533851,  0.00522987, ...,  0.02622218,\n",
       "        -0.05418065, -0.00765004],\n",
       "       [-0.0781679 ,  0.03937826, -0.01087696, ...,  0.04282334,\n",
       "        -0.02091636, -0.01027698],\n",
       "       ...,\n",
       "       [ 0.0878398 ,  0.01232621,  0.00077178, ..., -0.00705758,\n",
       "         0.01574707, -0.01541145],\n",
       "       [ 0.0882502 ,  0.03615745, -0.00961868, ...,  0.01392467,\n",
       "         0.00077467, -0.02139922],\n",
       "       [ 0.09195283,  0.04004925, -0.00255262, ...,  0.0036222 ,\n",
       "        -0.0181689 , -0.04212729]], dtype=float32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63581c47-4d2a-4106-883f-3f42c9070e99",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "69bf8642-772f-4996-b3c7-910a1f38b1b2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:44.592214Z",
     "iopub.status.busy": "2023-04-26T14:52:44.591565Z",
     "iopub.status.idle": "2023-04-26T14:52:44.600591Z",
     "shell.execute_reply": "2023-04-26T14:52:44.599816Z",
     "shell.execute_reply.started": "2023-04-26T14:52:44.592163Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "512"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dimension = corpus.shape[-1]\n",
    "dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5fbf481-46ac-45d3-b277-e9e19185c214",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "06a59231-abcf-45f9-a69a-69149e73c2f8",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Flat Index / Brute Force\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2b9e1b9-c75b-4700-a3c9-cb2b8ecab451",
   "metadata": {},
   "source": [
    "FAISS supports a bruteforce index.    \n",
    "This index is good if you want perfect recall.   \n",
    "It requires all the data to be fit in memory.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "926f4977-2ee0-48bd-8f50-f13e6ed82897",
   "metadata": {},
   "source": [
    "## Create the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a447b05c-f30d-4d43-a6ed-fd29a99477a2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:44.601898Z",
     "iopub.status.busy": "2023-04-26T14:52:44.601570Z",
     "iopub.status.idle": "2023-04-26T14:52:44.609907Z",
     "shell.execute_reply": "2023-04-26T14:52:44.609154Z",
     "shell.execute_reply.started": "2023-04-26T14:52:44.601873Z"
    }
   },
   "outputs": [],
   "source": [
    "x_corpus = corpus\n",
    "x_corpus.shape\n",
    "dimension = x_corpus.shape[-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb28c78c-3156-4da9-b5c1-3b192cb4c70f",
   "metadata": {},
   "source": [
    "initialize the flat index for data dimension.    \n",
    "In current example it is 512\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5bf3d65c-c435-407a-a932-8cdd9655ff5a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:44.612718Z",
     "iopub.status.busy": "2023-04-26T14:52:44.612384Z",
     "iopub.status.idle": "2023-04-26T14:52:44.664297Z",
     "shell.execute_reply": "2023-04-26T14:52:44.663442Z",
     "shell.execute_reply.started": "2023-04-26T14:52:44.612692Z"
    }
   },
   "outputs": [],
   "source": [
    "index = faiss.IndexFlatL2(dimension)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c121ab02-4fe3-4012-af73-038ae78f872e",
   "metadata": {},
   "source": [
    "since it is a brute force index, there is no \"training\" or parameters to learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3c2d2c60-0afb-44cf-9cf6-5032811725a7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:44.666136Z",
     "iopub.status.busy": "2023-04-26T14:52:44.665598Z",
     "iopub.status.idle": "2023-04-26T14:52:44.670880Z",
     "shell.execute_reply": "2023-04-26T14:52:44.670173Z",
     "shell.execute_reply.started": "2023-04-26T14:52:44.666092Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index.is_trained\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a547440c-fd96-4ce5-9996-0210f00617a7",
   "metadata": {},
   "source": [
    "add data to the index. This is a CPU based index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7aa739ee-42e1-4ac5-b1d8-9876ec777129",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:44.672070Z",
     "iopub.status.busy": "2023-04-26T14:52:44.671820Z",
     "iopub.status.idle": "2023-04-26T14:52:44.747616Z",
     "shell.execute_reply": "2023-04-26T14:52:44.746751Z",
     "shell.execute_reply.started": "2023-04-26T14:52:44.672047Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "index.add(x_corpus)               "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ab9e1b45-1cf2-4e88-960c-6999ad312e22",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:44.749140Z",
     "iopub.status.busy": "2023-04-26T14:52:44.748763Z",
     "iopub.status.idle": "2023-04-26T14:52:44.754419Z",
     "shell.execute_reply": "2023-04-26T14:52:44.753707Z",
     "shell.execute_reply.started": "2023-04-26T14:52:44.749112Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "24954"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(x_corpus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfabea5d-26bd-45d8-8b7a-97f179bc4013",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "348f697d-db43-4093-aa4e-4858d4c058f0",
   "metadata": {},
   "source": [
    "number of vectors / results to retrieve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6200be29-b926-4847-9e42-daf90d99319d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:44.755735Z",
     "iopub.status.busy": "2023-04-26T14:52:44.755402Z",
     "iopub.status.idle": "2023-04-26T14:52:44.764114Z",
     "shell.execute_reply": "2023-04-26T14:52:44.763389Z",
     "shell.execute_reply.started": "2023-04-26T14:52:44.755710Z"
    }
   },
   "outputs": [],
   "source": [
    "k =1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d57e8258-b7fa-49cb-a9f7-52a62a0dda17",
   "metadata": {},
   "source": [
    "#### Index Search\n",
    "search method returns query indices (I) similar to search query vector and their euclidean distances (D) from the search query vector."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87ea66c3-ebda-4600-b5d8-34ad796cc7c2",
   "metadata": {},
   "source": [
    "search for single vector and get top 1 result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "306e3a89-331c-4253-b7cf-ce7a0d951b42",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:44.765288Z",
     "iopub.status.busy": "2023-04-26T14:52:44.765050Z",
     "iopub.status.idle": "2023-04-26T14:52:48.341486Z",
     "shell.execute_reply": "2023-04-26T14:52:48.340493Z",
     "shell.execute_reply.started": "2023-04-26T14:52:44.765265Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.38 ms ± 39.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "D, I = index.search(x_corpus[:1], k=1)   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2382cf6-2718-4240-b805-c96195af51f0",
   "metadata": {},
   "source": [
    "search for all vectors in corpus and get top 1 result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "badcb20e-f872-4763-baf5-1876d5dd617d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:48.342963Z",
     "iopub.status.busy": "2023-04-26T14:52:48.342590Z",
     "iopub.status.idle": "2023-04-26T14:52:58.653601Z",
     "shell.execute_reply": "2023-04-26T14:52:58.652643Z",
     "shell.execute_reply.started": "2023-04-26T14:52:48.342935Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 30.3 s, sys: 8.92 ms, total: 30.3 s\n",
      "Wall time: 10.3 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "D, I = index.search(x_corpus, k=1)     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acf3d26f-119c-4f5c-8c24-fd28e3fccdfc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "1eac1d91-007d-4143-b0a3-ae19c763cc6c",
   "metadata": {},
   "source": [
    "distance of vector in corpus to query vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0b555673-cb87-4b11-aa81-ed2ff69d513e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:58.660044Z",
     "iopub.status.busy": "2023-04-26T14:52:58.657749Z",
     "iopub.status.idle": "2023-04-26T14:52:58.667064Z",
     "shell.execute_reply": "2023-04-26T14:52:58.666198Z",
     "shell.execute_reply.started": "2023-04-26T14:52:58.660006Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.0000000e+00],\n",
       "       [0.0000000e+00],\n",
       "       [3.5762787e-07],\n",
       "       ...,\n",
       "       [0.0000000e+00],\n",
       "       [1.3113022e-06],\n",
       "       [7.1525574e-07]], dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "D"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4d271c3-850a-4375-87a2-100ffab7a416",
   "metadata": {},
   "source": [
    "top vertex id \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "65310b96-aab0-4a34-a9eb-48d0dbefc0ec",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:58.668403Z",
     "iopub.status.busy": "2023-04-26T14:52:58.668046Z",
     "iopub.status.idle": "2023-04-26T14:52:58.690783Z",
     "shell.execute_reply": "2023-04-26T14:52:58.689783Z",
     "shell.execute_reply.started": "2023-04-26T14:52:58.668375Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[    0],\n",
       "       [    1],\n",
       "       [    2],\n",
       "       ...,\n",
       "       [24951],\n",
       "       [24952],\n",
       "       [24953]])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "I"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90283075-8783-4607-bf28-e5ce9f55c08c",
   "metadata": {},
   "source": [
    "because we are using the entire corpus and the ids are sequential, the ideal recall would be sequential too"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b3fb40a2-dd43-4676-a766-3d198943f957",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:58.692358Z",
     "iopub.status.busy": "2023-04-26T14:52:58.691962Z",
     "iopub.status.idle": "2023-04-26T14:52:58.703731Z",
     "shell.execute_reply": "2023-04-26T14:52:58.702726Z",
     "shell.execute_reply.started": "2023-04-26T14:52:58.692330Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ True,  True,  True, ...,  True,  True,  True])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = I[:,0] == np.array( list(range(len(x_corpus))))\n",
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "27641884-4563-4c1a-9d37-1b8e61ee5322",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:58.705103Z",
     "iopub.status.busy": "2023-04-26T14:52:58.704836Z",
     "iopub.status.idle": "2023-04-26T14:52:58.713367Z",
     "shell.execute_reply": "2023-04-26T14:52:58.712403Z",
     "shell.execute_reply.started": "2023-04-26T14:52:58.705078Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([], dtype=int64),)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.where(res == False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ec2181ae-7a70-4a3a-bc34-bd9899168fe5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:58.714883Z",
     "iopub.status.busy": "2023-04-26T14:52:58.714532Z",
     "iopub.status.idle": "2023-04-26T14:52:58.726900Z",
     "shell.execute_reply": "2023-04-26T14:52:58.725900Z",
     "shell.execute_reply.started": "2023-04-26T14:52:58.714856Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'recall@1': 24954, 'num_vectors': 24954, 'mismatch': 0}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "{\n",
    " \"recall@1\":  res.sum()\n",
    " , \"num_vectors\":  len(res)\n",
    " , \"mismatch\":    len(res) - res.sum()\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f52a636e-041b-43e6-9b94-5c31447f31cc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T00:02:55.003678Z",
     "iopub.status.busy": "2023-04-26T00:02:55.002818Z",
     "iopub.status.idle": "2023-04-26T00:02:55.010824Z",
     "shell.execute_reply": "2023-04-26T00:02:55.010099Z",
     "shell.execute_reply.started": "2023-04-26T00:02:55.003640Z"
    }
   },
   "source": [
    "For this corpus, we are able to find the query vector as position 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e88e34b6-f9e0-4835-90cc-ba0e6b2c0414",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02bac14d-9696-4a49-be12-9541beeb45a2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "7e79d4fc-8191-4151-b330-01b41a2b05d4",
   "metadata": {},
   "source": [
    "# FAISS IVF"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c6f3e4c-2cad-461c-a8b7-5ffba0a5b354",
   "metadata": {},
   "source": [
    "<img src=\"https://d33wubrfki0l68.cloudfront.net/44acb1425f25e30ca058daec92bdb209c6c47ad2/e92fc/images/faiss5.png\" width=\"500\"/>\n",
    "\n",
    "<p> Image from Pinecone Faiss Tutorial </p>\n",
    "https://www.pinecone.io/learn/faiss-tutorial/\n",
    "\n",
    "\n",
    "**Parameters**:\n",
    "- nlist : number of clusters\n",
    "- nprobe: number of clusters to search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "342966dc-d361-4fec-8ebe-c3c67864736e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:58.728467Z",
     "iopub.status.busy": "2023-04-26T14:52:58.728210Z",
     "iopub.status.idle": "2023-04-26T14:52:58.785614Z",
     "shell.execute_reply": "2023-04-26T14:52:58.784521Z",
     "shell.execute_reply.started": "2023-04-26T14:52:58.728443Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "nlist = 20 # number of clusters\n",
    "quantizer = faiss.IndexFlatL2(dimension)  # the other index\n",
    "index = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_L2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "fec38a91-8c3c-40fe-935d-a160474a7e4f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:58.787144Z",
     "iopub.status.busy": "2023-04-26T14:52:58.786846Z",
     "iopub.status.idle": "2023-04-26T14:52:58.883952Z",
     "shell.execute_reply": "2023-04-26T14:52:58.882788Z",
     "shell.execute_reply.started": "2023-04-26T14:52:58.787117Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert not index.is_trained\n",
    "index.train(x_corpus)\n",
    "assert index.is_trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "6b577199-2e7b-439f-93a4-2653c7545eef",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:58.885895Z",
     "iopub.status.busy": "2023-04-26T14:52:58.885561Z",
     "iopub.status.idle": "2023-04-26T14:52:58.946474Z",
     "shell.execute_reply": "2023-04-26T14:52:58.945447Z",
     "shell.execute_reply.started": "2023-04-26T14:52:58.885865Z"
    }
   },
   "outputs": [],
   "source": [
    "index.add(x_corpus)         "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55f7c7b3-7298-4e53-b1e0-2c3a0f568579",
   "metadata": {},
   "source": [
    "we need to train the index first with a sample of vectors before indexing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a50b8c1-ed8b-45ac-9f19-4325f42b265a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "f8d0824c-0712-46fb-a5d4-16b32ee695c0",
   "metadata": {},
   "source": [
    "search for single vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "451988b2-0e03-4e12-92e2-8df408094526",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:52:58.956643Z",
     "iopub.status.busy": "2023-04-26T14:52:58.953741Z",
     "iopub.status.idle": "2023-04-26T14:53:08.215103Z",
     "shell.execute_reply": "2023-04-26T14:53:08.214024Z",
     "shell.execute_reply.started": "2023-04-26T14:52:58.956606Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "114 µs ± 729 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "index.nprobe = 1              # default nprobe is 1\n",
    "\n",
    "D, I = index.search(x_corpus[:1], k)     # actual search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28e0a32e-5380-48da-9022-0be90db13e75",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T00:07:36.782792Z",
     "iopub.status.busy": "2023-04-26T00:07:36.782538Z",
     "iopub.status.idle": "2023-04-26T00:07:36.788109Z",
     "shell.execute_reply": "2023-04-26T00:07:36.786807Z",
     "shell.execute_reply.started": "2023-04-26T00:07:36.782768Z"
    }
   },
   "source": [
    "in the above, we are only querying 1/20 of the search space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7547048-000e-4f32-b1c9-fccfd83b3918",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "27ef12d6-d585-4be3-9f5f-8b35ee192c38",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:53:08.216344Z",
     "iopub.status.busy": "2023-04-26T14:53:08.216081Z",
     "iopub.status.idle": "2023-04-26T14:53:20.757164Z",
     "shell.execute_reply": "2023-04-26T14:53:20.755944Z",
     "shell.execute_reply.started": "2023-04-26T14:53:08.216319Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.55 ms ± 13.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "\n",
    "index.nprobe = 10              # default nprobe is 1\n",
    "\n",
    "D, I = index.search(x_corpus[:1], k)     # actual search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e98f114-4d94-4b59-84f2-73232fc834da",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2023-04-26T00:07:36.790059Z",
     "iopub.status.idle": "2023-04-26T00:07:36.790366Z",
     "shell.execute_reply": "2023-04-26T00:07:36.790209Z",
     "shell.execute_reply.started": "2023-04-26T00:07:36.790195Z"
    }
   },
   "source": [
    "in the above, we are only querying half of the search space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1045fcc-358f-4216-a00e-e0db6e1811a4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50e83d62-b203-4048-95df-74e61fc2aa0c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "2a140a52-213a-4da1-8077-8be794a36f30",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:53:20.758823Z",
     "iopub.status.busy": "2023-04-26T14:53:20.758438Z",
     "iopub.status.idle": "2023-04-26T14:53:24.717603Z",
     "shell.execute_reply": "2023-04-26T14:53:24.716457Z",
     "shell.execute_reply.started": "2023-04-26T14:53:20.758795Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.88 ms ± 57.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "\n",
    "index.nprobe = 20              # default nprobe is 1\n",
    "\n",
    "D, I = index.search(x_corpus[:1], k)     # actual search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "971f14de-ec0c-4ac7-b63a-135ec7444834",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2023-04-26T00:07:36.792692Z",
     "iopub.status.idle": "2023-04-26T00:07:36.792990Z",
     "shell.execute_reply": "2023-04-26T00:07:36.792856Z",
     "shell.execute_reply.started": "2023-04-26T00:07:36.792842Z"
    }
   },
   "source": [
    "in the above, we are querying the entire search space. This is the same as using Brute Force."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec6f1eb2-a9cd-47d9-86bc-5c58ed5dbb1d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb11b2cd-eef5-4b15-a719-e9c5cbe0fb4f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "d8524b69-4efd-4a6d-9f10-499f299ff762",
   "metadata": {},
   "source": [
    "search for entire corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "60795b88-8957-4c49-a729-22f935c4dc3f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:53:24.719250Z",
     "iopub.status.busy": "2023-04-26T14:53:24.718697Z",
     "iopub.status.idle": "2023-04-26T14:53:25.774560Z",
     "shell.execute_reply": "2023-04-26T14:53:25.773678Z",
     "shell.execute_reply.started": "2023-04-26T14:53:24.719219Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8.18 s, sys: 23.5 ms, total: 8.21 s\n",
      "Wall time: 1.05 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "\n",
    "index.nprobe = 1              \n",
    "\n",
    "D, I = index.search(x_corpus, k)     # actual search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "ddbabfeb-2b27-412c-b77b-1d795b00aa29",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:53:25.781104Z",
     "iopub.status.busy": "2023-04-26T14:53:25.778777Z",
     "iopub.status.idle": "2023-04-26T14:53:25.791108Z",
     "shell.execute_reply": "2023-04-26T14:53:25.790296Z",
     "shell.execute_reply.started": "2023-04-26T14:53:25.781066Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'recall@1': 24954, 'num_vectors': 24954, 'mismatch': 0}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z = I[:,0] == np.array( list(range(len(x_corpus))))\n",
    "{\n",
    " \"recall@1\":  z.sum()\n",
    " , \"num_vectors\":  len(z)\n",
    " , \"mismatch\":    len(z) - z.sum()\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c9da4ac-9d1f-480a-ac9f-ea97e1d16f9d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "2b5515dc-09b7-459c-b93c-80424faac839",
   "metadata": {},
   "source": [
    "increase the number of cells that are probed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "687d36fe-de5d-4ee4-9941-7ca604b446fb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "31c0b86b-2d5e-480d-be94-c136e7fe07d1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:53:25.792414Z",
     "iopub.status.busy": "2023-04-26T14:53:25.792073Z",
     "iopub.status.idle": "2023-04-26T14:54:03.768762Z",
     "shell.execute_reply": "2023-04-26T14:54:03.767827Z",
     "shell.execute_reply.started": "2023-04-26T14:53:25.792388Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.73 s ± 230 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "index.nprobe = 5              # default nprobe is 1\n",
    "\n",
    "D, I = index.search(x_corpus, k)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "fbffb654-a3a6-4d42-956f-602dab97124e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-26T14:54:03.777354Z",
     "iopub.status.busy": "2023-04-26T14:54:03.769965Z",
     "iopub.status.idle": "2023-04-26T14:54:03.787506Z",
     "shell.execute_reply": "2023-04-26T14:54:03.786386Z",
     "shell.execute_reply.started": "2023-04-26T14:54:03.777305Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'recall@1': 24954, 'num_vectors': 24954, 'mismatch': 0}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z = I[:,0] == np.array( list(range(len(x_corpus))))\n",
    "{\n",
    " \"recall@1\":  z.sum()\n",
    " , \"num_vectors\":  len(z)\n",
    " , \"mismatch\":    len(z) - z.sum()\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd7e1354-2002-45dc-8732-9506ef6200cf",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "pytorch-gpu.1-13.m107",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/pytorch-gpu.1-13:m107"
  },
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}