pq_Faiss.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "79b74250-fa0d-4feb-b513-f07ca7ea0fab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "import urllib.request as request\n",
    "from contextlib import closing\n",
    "\n",
    "# first we download the Sift1M dataset\n",
    "with closing(request.urlopen('ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz')) as r:\n",
    "    with open('sift.tar.gz', 'wb') as f:\n",
    "        shutil.copyfileobj(r, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d43a1153-3fe0-49b9-99d1-9d4698977082",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tarfile\n",
    "\n",
    "# the download leaves us with a tar.gz file, we unzip it\n",
    "tar = tarfile.open('sift.tar.gz', \"r:gz\")\n",
    "tar.extractall()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4bfc84db-b952-45a2-97e0-a0a059c4e455",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# now define a function to read the fvecs file format of Sift1M dataset\n",
    "def read_fvecs(fp):\n",
    "    a = np.fromfile(fp, dtype='int32')\n",
    "    d = a[0]\n",
    "    return a.reshape(-1, d + 1)[:, 1:].copy().view('float32')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6889461f-8b6a-43f5-918d-d1e9eebd7814",
   "metadata": {},
   "outputs": [],
   "source": [
    "# data we will search through\n",
    "wb = read_fvecs('./sift/sift_base.fvecs')  # 1M samples\n",
    "# also get some query vectors to search with\n",
    "xq = read_fvecs('./sift/sift_query.fvecs')\n",
    "# take just one query (there are many in sift_learn.fvecs)\n",
    "xq = xq[0].reshape(1, xq.shape[1])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "903f3191-f7c8-4cd4-b816-faa2dc44fcad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 128)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xq.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8dd2b881-a6d4-46c4-a040-df7c33ef5354",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1000000, 128)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wb.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "b49e43c5-5e5b-4ae1-a35d-89504a931910",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  0.,  16.,  35.,   5.,  32.,  31.,  14.,  10.,  11.,  78.,  55.,\n",
       "         10.,  45.,  83.,  11.,   6.,  14.,  57., 102.,  75.,  20.,   8.,\n",
       "          3.,   5.,  67.,  17.,  19.,  26.,   5.,   0.,   1.,  22.,  60.,\n",
       "         26.,   7.,   1.,  18.,  22.,  84.,  53.,  85., 119., 119.,   4.,\n",
       "         24.,  18.,   7.,   7.,   1.,  81., 106., 102.,  72.,  30.,   6.,\n",
       "          0.,   9.,   1.,   9., 119.,  72.,   1.,   4.,  33., 119.,  29.,\n",
       "          6.,   1.,   0.,   1.,  14.,  52., 119.,  30.,   3.,   0.,   0.,\n",
       "         55.,  92., 111.,   2.,   5.,   4.,   9.,  22.,  89.,  96.,  14.,\n",
       "          1.,   0.,   1.,  82.,  59.,  16.,  20.,   5.,  25.,  14.,  11.,\n",
       "          4.,   0.,   0.,   1.,  26.,  47.,  23.,   4.,   0.,   0.,   4.,\n",
       "         38.,  83.,  30.,  14.,   9.,   4.,   9.,  17.,  23.,  41.,   0.,\n",
       "          0.,   2.,   8.,  19.,  25.,  23.,   1.]], dtype=float32)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q = wb[0].reshape(1, xq.shape[1])\n",
    "q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "def7bf85-1b09-4485-8966-152fe8443285",
   "metadata": {},
   "outputs": [],
   "source": [
    "import faiss\n",
    "\n",
    "D = xq.shape[1]\n",
    "m = 8\n",
    "assert D % m == 0\n",
    "nbits = 8  # number of bits per subquantizer, k* = 2**nbits\n",
    "index = faiss.IndexPQ(D, m, nbits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "03398bf8-ee86-4a10-b563-64e70d4d4e1a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index.is_trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "4e9261cc-9257-4fed-8315-62ad655ea59d",
   "metadata": {},
   "outputs": [],
   "source": [
    "index.train(wb) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "3fded7e3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index.is_trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "638b4d0a-1dfc-469a-bda6-3619b43f2772",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[3.4028235e+38, 3.4028235e+38, 3.4028235e+38]], dtype=float32)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "k = 3\n",
    "dist, I = index.search(q, k)\n",
    "dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "bfd8b4fe",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "49.7 µs ± 3.06 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "index.search(q, k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "879092e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n"
     ]
    }
   ],
   "source": [
    "l2_index = faiss.IndexFlatL2(D)\n",
    "l2_index.add(wb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "7278a45b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n"
     ]
    }
   ],
   "source": [
    "print(index.ntotal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b110f6f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 81.4 ms, sys: 10 µs, total: 81.4 ms\n",
      "Wall time: 80.2 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "l2_dist, l2_I = l2_index.search(xq, k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "58a17142",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum([1 for i in I[0] if i in l2_I])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7702ab78",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}