faiss_tutorial/intro.ipynb

{
 "metadata": {
  "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.8.8"
  },
  "orig_nbformat": 4,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.8 64-bit ('ml': conda)"
  },
  "interpreter": {
   "hash": "a683edd788238e5c64f9fa2e4bdd4387776bc5c6f4f0a84da0685f9a25e421d6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import faiss\n",
    "import requests\n",
    "from io import StringIO\n",
    "import pandas as pd"
   ]
  },
  {
   "source": [
    "# Dataset Initialization"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "'pair_ID\\tsentence_A\\tsentence_B\\trelatedness_score\\tentailment_judgment\\n1\\tA group of kids is playing in '"
      ]
     },
     "metadata": {},
     "execution_count": 2
    }
   ],
   "source": [
    "res = requests.get('https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/sick2014/SICK_train.txt')\n",
    "\n",
    "text = res.text\n",
    "text[:100]"
   ]
  },
  {
   "source": [
    "We need this in a dataframe, which we build from the `text` string like so:"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "   pair_ID                                         sentence_A  \\\n",
       "0        1  A group of kids is playing in a yard and an ol...   \n",
       "1        2  A group of children is playing in the house an...   \n",
       "2        3  The young boys are playing outdoors and the ma...   \n",
       "3        5  The kids are playing outdoors near a man with ...   \n",
       "4        9  The young boys are playing outdoors and the ma...   \n",
       "\n",
       "                                          sentence_B  relatedness_score  \\\n",
       "0  A group of boys in a yard is playing and a man...                4.5   \n",
       "1  A group of kids is playing in a yard and an ol...                3.2   \n",
       "2  The kids are playing outdoors near a man with ...                4.7   \n",
       "3  A group of kids is playing in a yard and an ol...                3.4   \n",
       "4  A group of kids is playing in a yard and an ol...                3.7   \n",
       "\n",
       "  entailment_judgment  \n",
       "0             NEUTRAL  \n",
       "1             NEUTRAL  \n",
       "2          ENTAILMENT  \n",
       "3             NEUTRAL  \n",
       "4             NEUTRAL  "
      ],
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>pair_ID</th>\n      <th>sentence_A</th>\n      <th>sentence_B</th>\n      <th>relatedness_score</th>\n      <th>entailment_judgment</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>1</td>\n      <td>A group of kids is playing in a yard and an ol...</td>\n      <td>A group of boys in a yard is playing and a man...</td>\n      <td>4.5</td>\n      <td>NEUTRAL</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>2</td>\n      <td>A group of children is playing in the house an...</td>\n      <td>A group of kids is playing in a yard and an ol...</td>\n      <td>3.2</td>\n      <td>NEUTRAL</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>3</td>\n      <td>The young boys are playing outdoors and the ma...</td>\n      <td>The kids are playing outdoors near a man with ...</td>\n      <td>4.7</td>\n      <td>ENTAILMENT</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>5</td>\n      <td>The kids are playing outdoors near a man with ...</td>\n      <td>A group of kids is playing in a yard and an ol...</td>\n      <td>3.4</td>\n      <td>NEUTRAL</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>9</td>\n      <td>The young boys are playing outdoors and the ma...</td>\n      <td>A group of kids is playing in a yard and an ol...</td>\n      <td>3.7</td>\n      <td>NEUTRAL</td>\n    </tr>\n  </tbody>\n</table>\n</div>"
     },
     "metadata": {},
     "execution_count": 3
    }
   ],
   "source": [
    "data = pd.read_csv(StringIO(text), sep='\\t')\n",
    "data.head()"
   ]
  },
  {
   "source": [
    "We will take all samples from `sentence_A` and build sentence embeddings for each - which we can then store in FAISS."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "['A group of kids is playing in a yard and an old man is standing in the background',\n",
       " 'A group of children is playing in the house and there is no man standing in the background',\n",
       " 'The young boys are playing outdoors and the man is smiling nearby',\n",
       " 'The kids are playing outdoors near a man with a smile',\n",
       " 'The young boys are playing outdoors and the man is smiling nearby']"
      ]
     },
     "metadata": {},
     "execution_count": 4
    }
   ],
   "source": [
    "sentences = data['sentence_A'].tolist()\n",
    "sentences[:5]"
   ]
  },
  {
   "source": [
    "And we will pull in the `sentence_B` column too, giving us ~4.5K *unique* sentences."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "4802"
      ]
     },
     "metadata": {},
     "execution_count": 5
    }
   ],
   "source": [
    "sentence_b = data['sentence_B'].tolist()\n",
    "sentences.extend(sentence_b)\n",
    "len(set(sentences))"
   ]
  },
  {
   "source": [
    "This isn't a particularly large number, so let's pull in a few more similar datasets."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "urls = [\n",
    "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2012/MSRpar.train.tsv',\n",
    "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2012/MSRpar.test.tsv',\n",
    "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2012/OnWN.test.tsv',\n",
    "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2013/OnWN.test.tsv',\n",
    "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2014/OnWN.test.tsv',\n",
    "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2014/images.test.tsv',\n",
    "    'https://raw.githubusercontent.com/brmson/dataset-sts/master/data/sts/semeval-sts/2015/images.test.tsv'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "b'Skipping line 191: expected 3 fields, saw 4\\nSkipping line 206: expected 3 fields, saw 4\\nSkipping line 295: expected 3 fields, saw 4\\nSkipping line 695: expected 3 fields, saw 4\\nSkipping line 699: expected 3 fields, saw 4\\n'\n",
      "b'Skipping line 104: expected 3 fields, saw 4\\nSkipping line 181: expected 3 fields, saw 4\\nSkipping line 317: expected 3 fields, saw 4\\nSkipping line 412: expected 3 fields, saw 5\\nSkipping line 508: expected 3 fields, saw 4\\n'\n"
     ]
    }
   ],
   "source": [
    "for url in urls:\n",
    "    res = requests.get(url)\n",
    "    # extract to dataframe\n",
    "    data = pd.read_csv(StringIO(res.text), sep='\\t', header=None, error_bad_lines=False)\n",
    "    # add to columns 1 and 2 to sentences list\n",
    "    sentences.extend(data[1].tolist())\n",
    "    sentences.extend(data[2].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "14505"
      ]
     },
     "metadata": {},
     "execution_count": 8
    }
   ],
   "source": [
    "len(set(sentences))"
   ]
  },
  {
   "source": [
    "Before converting to our sentence embeddings, we will save to text file as backup."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove duplicates and NaN\n",
    "sentences = [\n",
    "    sentence.replace('\\n', '') for sentence in list(set(sentences)) if type(sentence) is str\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('sentences.txt', 'w') as fp:\n",
    "    fp.write('\\n'.join(sentences))"
   ]
  },
  {
   "source": [
    "Now we have 14.5K *unique* sentences, a much better size. We'll go ahead and build the sentence embeddings (this can take some time, feel free to download the embeddings from [here]()). TODO add link"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "(14504, 768)"
      ]
     },
     "metadata": {},
     "execution_count": 11
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "model = SentenceTransformer('bert-base-nli-mean-tokens')\n",
    "\n",
    "sentence_embeddings = model.encode(sentences)\n",
    "sentence_embeddings.shape"
   ]
  },
  {
   "source": [
    "We can save/load from file in the case of needing to reload the notebook for any reason later."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "14504"
      ]
     },
     "metadata": {},
     "execution_count": 28
    }
   ],
   "source": [
    "sentence_embeddings.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'./sim_sentences/embeddings_X.npy', 'wb') as fp:\n",
    "    np.save(fp, sentence_embeddings[0:256])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "embeddings_0.npy | 0 -> 256\n",
      "embeddings_1.npy | 256 -> 512\n",
      "embeddings_2.npy | 512 -> 768\n",
      "embeddings_3.npy | 768 -> 1024\n",
      "embeddings_4.npy | 1024 -> 1280\n",
      "embeddings_5.npy | 1280 -> 1536\n",
      "embeddings_6.npy | 1536 -> 1792\n",
      "embeddings_7.npy | 1792 -> 2048\n",
      "embeddings_8.npy | 2048 -> 2304\n",
      "embeddings_9.npy | 2304 -> 2560\n",
      "embeddings_10.npy | 2560 -> 2816\n",
      "embeddings_11.npy | 2816 -> 3072\n",
      "embeddings_12.npy | 3072 -> 3328\n",
      "embeddings_13.npy | 3328 -> 3584\n",
      "embeddings_14.npy | 3584 -> 3840\n",
      "embeddings_15.npy | 3840 -> 4096\n",
      "embeddings_16.npy | 4096 -> 4352\n",
      "embeddings_17.npy | 4352 -> 4608\n",
      "embeddings_18.npy | 4608 -> 4864\n",
      "embeddings_19.npy | 4864 -> 5120\n",
      "embeddings_20.npy | 5120 -> 5376\n",
      "embeddings_21.npy | 5376 -> 5632\n",
      "embeddings_22.npy | 5632 -> 5888\n",
      "embeddings_23.npy | 5888 -> 6144\n",
      "embeddings_24.npy | 6144 -> 6400\n",
      "embeddings_25.npy | 6400 -> 6656\n",
      "embeddings_26.npy | 6656 -> 6912\n",
      "embeddings_27.npy | 6912 -> 7168\n",
      "embeddings_28.npy | 7168 -> 7424\n",
      "embeddings_29.npy | 7424 -> 7680\n",
      "embeddings_30.npy | 7680 -> 7936\n",
      "embeddings_31.npy | 7936 -> 8192\n",
      "embeddings_32.npy | 8192 -> 8448\n",
      "embeddings_33.npy | 8448 -> 8704\n",
      "embeddings_34.npy | 8704 -> 8960\n",
      "embeddings_35.npy | 8960 -> 9216\n",
      "embeddings_36.npy | 9216 -> 9472\n",
      "embeddings_37.npy | 9472 -> 9728\n",
      "embeddings_38.npy | 9728 -> 9984\n",
      "embeddings_39.npy | 9984 -> 10240\n",
      "embeddings_40.npy | 10240 -> 10496\n",
      "embeddings_41.npy | 10496 -> 10752\n",
      "embeddings_42.npy | 10752 -> 11008\n",
      "embeddings_43.npy | 11008 -> 11264\n",
      "embeddings_44.npy | 11264 -> 11520\n",
      "embeddings_45.npy | 11520 -> 11776\n",
      "embeddings_46.npy | 11776 -> 12032\n",
      "embeddings_47.npy | 12032 -> 12288\n",
      "embeddings_48.npy | 12288 -> 12544\n",
      "embeddings_49.npy | 12544 -> 12800\n",
      "embeddings_50.npy | 12800 -> 13056\n",
      "embeddings_51.npy | 13056 -> 13312\n",
      "embeddings_52.npy | 13312 -> 13568\n",
      "embeddings_53.npy | 13568 -> 13824\n",
      "embeddings_54.npy | 13824 -> 14080\n",
      "embeddings_55.npy | 14080 -> 14336\n",
      "embeddings_56.npy | 14336 -> 14505\n"
     ]
    }
   ],
   "source": [
    "# saving data\n",
    "split = 256\n",
    "file_count = 0\n",
    "for i in range(0, sentence_embeddings.shape[0], split):\n",
    "    end = i + split\n",
    "    if end > sentence_embeddings.shape[0] + 1:\n",
    "        end = sentence_embeddings.shape[0] + 1\n",
    "    file_count = '0' + str(file_count) if file_count < 0 else str(file_count)\n",
    "    with open(f'./sim_sentences/embeddings_{file_count}.npy', 'wb') as fp:\n",
    "        np.save(fp, sentence_embeddings[i:end, :])\n",
    "    print(f\"embeddings_{file_count}.npy | {i} -> {end}\")\n",
    "    file_count = int(file_count) + 1"
   ]
  },
  {
   "source": [
    "We setup our FAISS database dimensionality (number of dimensions per vector) based on these vectors."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "768"
      ]
     },
     "metadata": {},
     "execution_count": 16
    }
   ],
   "source": [
    "d = sentence_embeddings.shape[1]\n",
    "d"
   ]
  },
  {
   "source": [
    "# Flat L2 Index\n",
    "\n",
    "We initialize the flat L2 distance index `IndexFlatL2`, all we need is the specify the vector dimensionality - which in this case is `d == 768` (to align with the sentence-BERT model output embeddings of size `768`)."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = faiss.IndexFlatL2(d)"
   ]
  },
  {
   "source": [
    "Often, we will use indexes that require us to `train` them on our data before being used (if we are grouping or transforming the data in any way). `IndexFlatL2` however, is a simple operation and only requires that we calculate distances between vectors when we introduce our query vector `xq` during search. So, in this case, no training is required - which we can confirm by checking the `is_trained` attribute."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "metadata": {},
     "execution_count": 18
    }
   ],
   "source": [
    "index.is_trained"
   ]
  },
  {
   "source": [
    "Okay so once we're happy that our index is prepared, we then add new vectors using the `add` method."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "index.add(sentence_embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "14504"
      ]
     },
     "metadata": {},
     "execution_count": 20
    }
   ],
   "source": [
    "index.ntotal"
   ]
  },
  {
   "source": [
    "Then search given a query `xq` and number of nearest neigbors to return `k`."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 4\n",
    "xq = model.encode([\"Someone sprints with a football\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[[ 4586 10252 12465   190]]\nCPU times: user 16.2 ms, sys: 7.57 ms, total: 23.8 ms\nWall time: 9.05 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "D, I = index.search(xq, k)  # search\n",
    "print(I)  # k-nearest neigbors of the query vector | nprobe == 1: 6495 26392 61709 49932 | nprobe == 10: 36245  6495 57489  8705"
   ]
  },
  {
   "source": [
    "Here we're returning indices `7460`, `10940`, `3781`, and `5747`, which returns:"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "['4586: A group of football players is running in the field',\n",
       " '10252: A group of people playing football is running in the field',\n",
       " '12465: Two groups of people are playing football',\n",
       " '190: A person playing football is running past an official carrying a football']"
      ]
     },
     "metadata": {},
     "execution_count": 23
    }
   ],
   "source": [
    "[f'{i}: {sentences[i]}' for i in I[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "'A group of football players is running in the field'"
      ]
     },
     "metadata": {},
     "execution_count": 32
    }
   ],
   "source": [
    "sentences[7460]"
   ]
  },
  {
   "source": [
    "Clearly we have some good matches, everything returned includes people running with a football, or on the context of a football match. Now, if we'd rather extract the numerical vectors from FAISS, we can do that too."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "vecs = np.zeros((k, d))\n",
    "for i, val in enumerate(I[0].tolist()):\n",
    "    vecs[i, :] = index.reconstruct(val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "(4, 768)"
      ]
     },
     "metadata": {},
     "execution_count": 34
    }
   ],
   "source": [
    "vecs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "array([ 0.01627046,  0.22325929, -0.15037425, -0.30747262, -0.27122435,\n",
       "       -0.10593167, -0.0646093 ,  0.04738174, -0.73349041, -0.37657705,\n",
       "       -0.76762843,  0.16902871,  0.53107643,  0.51176691,  1.14415848,\n",
       "       -0.08562929, -0.67240077, -0.96637076,  0.02545463, -0.21559823,\n",
       "       -1.25656605, -0.82982153, -0.09825023, -0.21850856,  0.50610232,\n",
       "        0.10527933,  0.50396878,  0.65242976, -1.39458692,  0.6584745 ,\n",
       "       -0.21525341, -0.22487433,  0.81818348,  0.084643  , -0.76141697,\n",
       "       -0.28928292, -0.09825806, -0.73046201,  0.07855812, -0.84354657,\n",
       "       -0.59242058,  0.77471375, -1.20920527, -0.22757971, -1.30733597,\n",
       "       -0.23081468, -1.31322539,  0.01629073, -0.97285444,  0.19308169,\n",
       "        0.47424558,  1.18920875, -1.96741307, -0.70061135, -0.29638717,\n",
       "        0.60533702,  0.62407452, -0.70340395, -0.86754197,  0.17673187,\n",
       "       -0.1917053 , -0.02951987,  0.22623563, -0.16695444, -0.80402541,\n",
       "       -0.45918915,  0.69675523, -0.24928184, -1.01478696, -0.921745  ,\n",
       "       -0.33842632, -0.39296737, -0.83734828, -0.11479235,  0.46049669,\n",
       "       -1.4521122 ,  0.60310453,  0.38696304, -0.04061254,  0.00453161,\n",
       "        0.24117804,  0.05396278,  0.07506453,  1.05115867,  0.12383959,\n",
       "       -0.71281093,  0.11722917,  0.52238214, -0.04581215,  0.26827109,\n",
       "        0.8598538 , -0.3566995 , -0.64667088, -0.5435797 , -0.0431047 ,\n",
       "        0.95139188, -0.15605772, -0.49625337, -0.11140176,  0.15610115])"
      ]
     },
     "metadata": {},
     "execution_count": 35
    }
   ],
   "source": [
    "vecs[0][:100]"
   ]
  },
  {
   "source": [
    "---\n",
    "\n",
    "# Adding Partitioning to the Index\n",
    "\n",
    "FAISS allows us to add an additional step to optimize our search efficiency using a variety of different methods. A popular approach is to partition the index into *[Voronoi cells](https://en.wikipedia.org/wiki/Voronoi_diagram)*. Using this method we would take our query vector `xq`, identify the *cell* it belongs to, and then use our `IndexFlatL2` to search between the query vector `xq` and all indexed vectors belonging to that *cell*. We can also include vectors from other nearby cells too.\n",
    "\n",
    "We initialize our new partitioned index by first adding our previous `IndexFlatL2` operation as a quantization step (another step in the search process), and feeding this into the new `IndexIVFFlat` operation like so:"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlist = 50\n",
    "quantizer = faiss.IndexFlatL2(d)\n",
    "index = faiss.IndexIVFFlat(quantizer, d, nlist)"
   ]
  },
  {
   "source": [
    "Here we've added a new parameter `nlist`. We use `nlist` to define how many partitions we'd like our index to have. \n",
    "\n",
    "When we built the previous, `IndexFlatL2`-only index, we noted that no training was required as no grouping/transformation was required to build that index. Now that we've added partitioning using `IndexIVFFlat`, this is no longer the case. Let's take a look at the `is_trained` attribute."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "metadata": {},
     "execution_count": 37
    }
   ],
   "source": [
    "index.is_trained"
   ]
  },
  {
   "source": [
    "So, what we need to do now is `train` our index on our data, which we do *before* adding any data to the index (otherwise the index cannot know how to group each vector)."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "metadata": {},
     "execution_count": 38
    }
   ],
   "source": [
    "index.train(sentence_embeddings)\n",
    "index.is_trained"
   ]
  },
  {
   "source": [
    "Now our index is trained, we add our data just as we did before."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "14504"
      ]
     },
     "metadata": {},
     "execution_count": 39
    }
   ],
   "source": [
    "index.add(sentence_embeddings)\n",
    "index.ntotal"
   ]
  },
  {
   "source": [
    "Let's search again using the same indexed sentence embeddings and the same query `xq`."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[[ 7460 10940  3781  5747]]\nCPU times: user 3.83 ms, sys: 3.25 ms, total: 7.08 ms\nWall time: 2.15 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "D, I = index.search(xq, k)  # search\n",
    "print(I)"
   ]
  },
  {
   "source": [
    "We can increase the number of nearby cells to search too with `nprobe`."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "index.nprobe = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[[ 7460 10940  3781  5747]]\nCPU times: user 5.29 ms, sys: 2.7 ms, total: 7.99 ms\nWall time: 1.54 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "D, I = index.search(xq, k)  # search\n",
    "print(I)"
   ]
  },
  {
   "source": [
    "Increasing the number of `nprobe` will improve the accuracy of our search, but cost time. Our earlier `IndexFlatL2`-only search was *exhaustive* (it compared every single vector) and so it identified the closest matches with a perfect accuracy. The smaller our `nprobe` value, the smaller scope that we search. We received perfect results (that matched our previous `IndexFlatL2`-only results - `7460`, `10940`, `3781`, `5747`), however, if we found that we were not getting closely matching results, we could simply bump `nprobe` up further - improving accuracy, but increasing time-taken too.\n",
    "\n",
    "It's worth noting that the time taken can change with each run too, if we rerun the above block, we usually get a different time:"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[[ 7460 10940  3781  5747]]\nCPU times: user 4.88 ms, sys: 2.62 ms, total: 7.5 ms\nWall time: 1.32 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "D, I = index.search(xq, k)\n",
    "print(I)"
   ]
  },
  {
   "source": [
    "For IVF (and IMI) indexes, before attempting to use the `reconstruct` method, we need to call the `make_direct_map` method - otherwise we will return a `RunetimeError`."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "output_type": "error",
     "ename": "RuntimeError",
     "evalue": "Error in DirectMap::idx_t faiss::DirectMap::get(faiss::DirectMap::idx_t) const at /Users/distiller/miniconda/conda-bld/faiss-pkg_1623024446684/work/faiss/invlists/DirectMap.cpp:81: direct map not initialized",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-45-7708be30eb32>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mindex\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreconstruct\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m7460\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/opt/homebrew/anaconda3/envs/ml/lib/python3.8/site-packages/faiss/__init__.py\u001b[0m in \u001b[0;36mreplacement_reconstruct\u001b[0;34m(self, key, x)\u001b[0m\n\u001b[1;32m    404\u001b[0m             \u001b[0;32massert\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0md\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    405\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 406\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreconstruct_c\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mswig_ptr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    407\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    408\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/homebrew/anaconda3/envs/ml/lib/python3.8/site-packages/faiss/swigfaiss.py\u001b[0m in \u001b[0;36mreconstruct\u001b[0;34m(self, key, recons)\u001b[0m\n\u001b[1;32m   3105\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3106\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mreconstruct\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrecons\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3107\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0m_swigfaiss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mIndexIVF_reconstruct\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrecons\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3108\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3109\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mupdate_vectors\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnv\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Error in DirectMap::idx_t faiss::DirectMap::get(faiss::DirectMap::idx_t) const at /Users/distiller/miniconda/conda-bld/faiss-pkg_1623024446684/work/faiss/invlists/DirectMap.cpp:81: direct map not initialized"
     ]
    }
   ],
   "source": [
    "index.reconstruct(7460)"
   ]
  },
  {
   "source": [
    "With `make_direct_map`:"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "index.make_direct_map()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "array([ 0.01627046,  0.22325929, -0.15037425, -0.30747262, -0.27122435,\n",
       "       -0.10593167, -0.0646093 ,  0.04738174, -0.7334904 , -0.37657705,\n",
       "       -0.76762843,  0.16902871,  0.53107643,  0.5117669 ,  1.1441585 ,\n",
       "       -0.08562929, -0.6724008 , -0.96637076,  0.02545463, -0.21559823,\n",
       "       -1.256566  , -0.8298215 , -0.09825023, -0.21850856,  0.5061023 ,\n",
       "        0.10527933,  0.5039688 ,  0.65242976, -1.3945869 ,  0.6584745 ,\n",
       "       -0.21525341, -0.22487433,  0.8181835 ,  0.084643  , -0.761417  ,\n",
       "       -0.28928292, -0.09825806, -0.730462  ,  0.07855812, -0.84354657,\n",
       "       -0.5924206 ,  0.77471375, -1.2092053 , -0.22757971, -1.307336  ,\n",
       "       -0.23081468, -1.3132254 ,  0.01629073, -0.97285444,  0.19308169,\n",
       "        0.47424558,  1.1892087 , -1.9674131 , -0.70061135, -0.29638717,\n",
       "        0.605337  ,  0.6240745 , -0.70340395, -0.86754197,  0.17673187,\n",
       "       -0.1917053 , -0.02951987,  0.22623563, -0.16695444, -0.8040254 ,\n",
       "       -0.45918915,  0.69675523, -0.24928184, -1.014787  , -0.921745  ,\n",
       "       -0.33842632, -0.39296737, -0.8373483 , -0.11479235,  0.4604967 ,\n",
       "       -1.4521122 ,  0.60310453,  0.38696304, -0.04061254,  0.00453161,\n",
       "        0.24117804,  0.05396278,  0.07506453,  1.0511587 ,  0.12383959,\n",
       "       -0.71281093,  0.11722917,  0.52238214, -0.04581215,  0.2682711 ,\n",
       "        0.8598538 , -0.3566995 , -0.6466709 , -0.5435797 , -0.0431047 ,\n",
       "        0.9513919 , -0.15605772, -0.49625337, -0.11140176,  0.15610115],\n",
       "      dtype=float32)"
      ]
     },
     "metadata": {},
     "execution_count": 48
    }
   ],
   "source": [
    "index.reconstruct(7460)[:100]"
   ]
  },
  {
   "source": [
    "We've now significantly reduced the search time, what can we do next?\n",
    "\n",
    "# Quantization\n",
    "\n",
    "Well, when storing these vectors we're storing the full (eg `Flat`) vector. Now in very big datasets this can quickly become a problem. Typically, we look at big datasets, and when working with large dataset we will find that storing the full vectors consumes too much space.\n",
    "\n",
    "Fortunately, FAISS comes with the ability to *compress* our vectors using transformations based on *Product Quantization* (PQ). But, what is PQ? Well, we can view it as an additional approximation step similar to our use of **IVF**, which allowed us to *approximate* by reducing the scope of our search. PQ is slightly different however, and approximates the distance (or similarity) calculation instead.\n",
    "\n",
    "[PQ explanation TODO REMOVE](https://mccormickml.com/2017/10/13/product-quantizer-tutorial-part-1/)\n",
    "\n",
    "PQ achieves this approximated distance operation by compressing the vectors themselves. This consists of a few steps.\n",
    "\n",
    "1. We split the every original vector into several subvectors.\n",
    "\n",
    "2. For each set of subvectors, we perform a clustering operation, creating many centroids for each subvector set.\n",
    "\n",
    "3. In our vector of subvectors, we replace each subvector with the ID of it's nearest centroid."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = 8  # number of centroid IDs in final compressed vectors\n",
    "bits = 8 # number of bits in each centroid\n",
    "\n",
    "quantizer = faiss.IndexFlatL2(d)  # we keep the same L2 distance flat index\n",
    "index = faiss.IndexIVFPQ(quantizer, d, nlist, m, bits) "
   ]
  },
  {
   "source": [
    "Again, we'll need to `train` the index."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "metadata": {},
     "execution_count": 50
    }
   ],
   "source": [
    "index.is_trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "index.train(sentence_embeddings)"
   ]
  },
  {
   "source": [
    "And `add` our vectors."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "index.add(sentence_embeddings)"
   ]
  },
  {
   "source": [
    "Let's compare it to our previous index *without* PQ, and an `nprobe` value of `10`."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "index.nprobe = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[[ 5013 10940  7460  5370]]\nCPU times: user 3.04 ms, sys: 2.18 ms, total: 5.22 ms\nWall time: 1.33 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "D, I = index.search(xq, k)\n",
    "print(I)"
   ]
  },
  {
   "source": [
    "Through adding PQ we've reduced our search time from ~7.5ms to ~5ms, a small difference on a dataset of this size, but when scaled to larger datasets this can make a huge difference.\n",
    "\n",
    "Now, we should also notice the slightly different results being returned. Beforehand with our exhaustive L2 search we were returning `7460`, `10940`, `3781`, and `5747`. Now, we see a slightly different order to our results - and two different vectors, `5013` and `5370`.\n",
    "\n",
    "Each of our speed optimization operations, `IVF` and `PQ`, come at the cost of accuracy. Now, if we print out these results we will nonetheless find that each item is still a relevant match:"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "['5013: A group of football players running down the field.',\n",
       " '10940: A group of people playing football is running in the field',\n",
       " '7460: A group of football players is running in the field',\n",
       " '5370: A football player is running past an official carrying a football']"
      ]
     },
     "metadata": {},
     "execution_count": 61
    }
   ],
   "source": [
    "[f'{i}: {sentences[i]}' for i in I[0]]"
   ]
  },
  {
   "source": [
    "So although we might not get the *perfect* result, we still get close - and thanks to the approximation, we get a significant speed boost"
   ],
   "cell_type": "markdown",
   "metadata": {}
  }
 ]
}