faiss_tests.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!conda install -y faiss-gpu -c pytorch\n",
    "from scipy.spatial.distance import cosine\n",
    "import numpy as np\n",
    "import os\n",
    "from sentence_transformers import SentenceTransformer\n",
    "import psycopg2\n",
    "from psycopg2 import extras\n",
    "import time\n",
    "\n",
    "d=768\n",
    "HOST = '54.196.150.193'\n",
    "USER = 'postgres'\n",
    "PASS = os.environ.get('PGPASS')\n",
    "PGSSLROOTCERT = os.environ.get('PGSSLROOTCERT')\n",
    "if PASS == None or PGSSLROOTCERT == None:\n",
    "    print(\"Please set PG_PASS and PGSSLROOTCERT env variable\")\n",
    "    raise SystemExit()\n",
    "DBNAME = 'fever'\n",
    "POSTGRES_DSN = f'''dbname='fever' user='{USER}' host='{HOST}' password='{PASS}' '''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResultHeap:\n",
    "    \"\"\" Combine query results from a sliced dataset \"\"\"\n",
    "\n",
    "    def __init__(self, nq, k):\n",
    "        \" nq: number of query vectors, k: number of results per query \"\n",
    "        self.I = np.zeros((nq, k), dtype='int64')\n",
    "        self.D = np.zeros((nq, k), dtype='float32')\n",
    "        self.nq, self.k = nq, k\n",
    "        heaps = faiss.float_maxheap_array_t()\n",
    "        heaps.k = k\n",
    "        heaps.nh = nq\n",
    "        heaps.val = faiss.swig_ptr(self.D)\n",
    "        heaps.ids = faiss.swig_ptr(self.I)\n",
    "        heaps.heapify()\n",
    "        self.heaps = heaps\n",
    "\n",
    "    def add_batch_result(self, D, I, i0):\n",
    "        assert D.shape == (self.nq, self.k)\n",
    "        assert I.shape == (self.nq, self.k)\n",
    "        I += i0\n",
    "        self.heaps.addn_with_ids(\n",
    "            self.k, faiss.swig_ptr(D),\n",
    "            faiss.swig_ptr(I), self.k)\n",
    "\n",
    "    def finalize(self):\n",
    "        self.heaps.reorder()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'faiss' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-3-043f13a2a21a>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mdef\u001b[0m \u001b[0msearch_knn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mxq\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mxb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdistance_type\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfaiss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mMETRIC_L2\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      2\u001b[0m     \u001b[0;34m\"\"\" wrapper around the faiss knn functions without index \"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m     \u001b[0mnq\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mxq\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0mnb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mxb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0;32massert\u001b[0m \u001b[0md\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0md2\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'faiss' is not defined"
     ]
    }
   ],
   "source": [
    "def search_knn(xq, xb, k, distance_type=faiss.METRIC_L2): \n",
    "    \"\"\" wrapper around the faiss knn functions without index \"\"\"\n",
    "    nq, d = xq.shape\n",
    "    nb, d2 = xb.shape\n",
    "    assert d == d2\n",
    "    \n",
    "    I = np.empty((nq, k), dtype='int64')\n",
    "    D = np.empty((nq, k), dtype='float32')\n",
    "    \n",
    "    if distance_type == faiss.METRIC_L2: \n",
    "        heaps = faiss.float_maxheap_array_t()\n",
    "        heaps.k = k\n",
    "        heaps.nh = nq\n",
    "        heaps.val = faiss.swig_ptr(D)\n",
    "        heaps.ids = faiss.swig_ptr(I)\n",
    "        faiss.knn_L2sqr(\n",
    "            faiss.swig_ptr(xq), faiss.swig_ptr(xb), \n",
    "            d, nq, nb, heaps\n",
    "        )\n",
    "    elif distance_type == faiss.METRIC_INNER_PRODUCT: \n",
    "        heaps = faiss.float_minheap_array_t()\n",
    "        heaps.k = k\n",
    "        heaps.nh = nq\n",
    "        heaps.val = faiss.swig_ptr(D)\n",
    "        heaps.ids = faiss.swig_ptr(I)\n",
    "        faiss.knn_inner_product(\n",
    "            faiss.swig_ptr(xq), faiss.swig_ptr(xb), \n",
    "            d, nq, nb, heaps\n",
    "        )\n",
    "    return D, I "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_raw_array_pytorch(res, xb, xq, k, D=None, I=None,\n",
    "                             metric=faiss.METRIC_L2):\n",
    "    assert xb.device == xq.device\n",
    "\n",
    "    nq, d = xq.size()\n",
    "    if xq.is_contiguous():\n",
    "        xq_row_major = True\n",
    "    elif xq.t().is_contiguous():\n",
    "        xq = xq.t()    # I initially wrote xq:t(), Lua is still haunting me :-)\n",
    "        xq_row_major = False\n",
    "    else:\n",
    "        raise TypeError('matrix should be row or column-major')\n",
    "\n",
    "    xq_ptr = swig_ptr_from_FloatTensor(xq)\n",
    "\n",
    "    nb, d2 = xb.size()\n",
    "    assert d2 == d\n",
    "    if xb.is_contiguous():\n",
    "        xb_row_major = True\n",
    "    elif xb.t().is_contiguous():\n",
    "        xb = xb.t()\n",
    "        xb_row_major = False\n",
    "    else:\n",
    "        raise TypeError('matrix should be row or column-major')\n",
    "    xb_ptr = swig_ptr_from_FloatTensor(xb)\n",
    "\n",
    "    if D is None:\n",
    "        D = torch.empty(nq, k, device=xb.device, dtype=torch.float32)\n",
    "    else:\n",
    "        assert D.shape == (nq, k)\n",
    "        assert D.device == xb.device\n",
    "\n",
    "    if I is None:\n",
    "        I = torch.empty(nq, k, device=xb.device, dtype=torch.int64)\n",
    "    else:\n",
    "        assert I.shape == (nq, k)\n",
    "        assert I.device == xb.device\n",
    "\n",
    "    D_ptr = swig_ptr_from_FloatTensor(D)\n",
    "    I_ptr = swig_ptr_from_LongTensor(I)\n",
    "\n",
    "    faiss.bruteForceKnn(res, metric,\n",
    "                        xb_ptr, xb_row_major, nb,\n",
    "                        xq_ptr, xq_row_major, nq,\n",
    "                        d, k, D_ptr, I_ptr)\n",
    "\n",
    "    return D, I"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "getting from postgres\n"
     ]
    }
   ],
   "source": [
    "model = SentenceTransformer('./fever-model/')\n",
    "conn = psycopg2.connect(POSTGRES_DSN)\n",
    "cur = conn.cursor()\n",
    "print('getting from postgres')\n",
    "cur.execute('''\n",
    "SELECT id, claim FROM sets.claims \n",
    "WHERE is_test_set AND verifiable\n",
    "''')\n",
    "res = cur.fetchall()\n",
    "claim_ids, claims = zip(*res)\n",
    "query_vectors = model.encode(claims, batch_size=32, show_progress_bar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def swig_ptr_from_FloatTensor(x):\n",
    "    assert x.is_contiguous()\n",
    "    assert x.dtype == torch.float32\n",
    "    return faiss.cast_integer_to_float_ptr(\n",
    "        x.storage().data_ptr() + x.storage_offset() * 4)\n",
    "\n",
    "def swig_ptr_from_LongTensor(x):\n",
    "    assert x.is_contiguous()\n",
    "    assert x.dtype == torch.int64, 'dtype=%s' % x.dtype\n",
    "    return faiss.cast_integer_to_long_ptr(\n",
    "        x.storage().data_ptr() + x.storage_offset() * 8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "id: 5396096 rate: 3356661.9277707604 per second\n",
      "id: 5396096 rate: 578878.6751916535 per second\n",
      "id: 5396096 rate: 334943.72122431686 per second\n",
      "id: 5396096 rate: 224512.5100982765 per second\n",
      "id: 5396096 rate: 174543.4830695143 per second\n",
      "id: 5396096 rate: 139417.42379495435 per second\n",
      "id: 5396096 rate: 119263.85486368318 per second\n",
      "id: 5396096 rate: 100591.02669696555 per second\n",
      "id: 5396096 rate: 88611.10875752025 per second\n",
      "id: 5396096 rate: 79675.53612302964 per second\n",
      "id: 5396096 rate: 71978.7402600271 per second\n",
      "id: 5396096 rate: 66065.22132283906 per second\n",
      "id: 5396096 rate: 60448.92514482267 per second\n",
      "id: 5396096 rate: 55759.14861954831 per second\n",
      "id: 5396096 rate: 52239.31387624846 per second\n",
      "id: 5396096 rate: 49053.41797575435 per second\n",
      "id: 5396096 rate: 45982.28167145725 per second\n",
      "id: 5396096 rate: 43360.889504654304 per second\n",
      "id: 5396096 rate: 40572.20296014832 per second\n",
      "id: 5396096 rate: 38418.69396491631 per second\n",
      "id: 5396096 rate: 36450.79419099415 per second\n",
      "id: 5396096 rate: 34686.228169105634 per second\n",
      "id: 5396096 rate: 33216.82027529402 per second\n",
      "id: 5396096 rate: 31613.170493764697 per second\n",
      "id: 5396096 rate: 30176.72589041001 per second\n",
      "id: 5396096 rate: 29046.922853206685 per second\n",
      "id: 5396096 rate: 27946.32138277378 per second\n",
      "id: 5396096 rate: 26947.174222598656 per second\n",
      "id: 5396096 rate: 26066.985810831517 per second\n",
      "id: 5396096 rate: 25234.091403114322 per second\n",
      "id: 5396096 rate: 24413.24944265776 per second\n",
      "id: 5396096 rate: 23613.30656754474 per second\n",
      "id: 5396096 rate: 22867.489145139843 per second\n",
      "id: 5396096 rate: 22195.518353402287 per second\n",
      "id: 5396096 rate: 21554.106898909227 per second\n",
      "id: 5396096 rate: 20911.532542143337 per second\n",
      "id: 5396096 rate: 20220.400982782838 per second\n",
      "id: 5396096 rate: 19707.71210763228 per second\n",
      "id: 5396096 rate: 19244.746115338487 per second\n",
      "id: 5396096 rate: 18771.082540291285 per second\n",
      "id: 5396096 rate: 18301.26647964942 per second\n",
      "id: 5396096 rate: 17811.826673008512 per second\n",
      "id: 5396096 rate: 17309.104025399527 per second\n",
      "id: 5396096 rate: 16938.30758910044 per second\n",
      "id: 5396096 rate: 16574.954458508273 per second\n",
      "id: 5396096 rate: 16172.172976897757 per second\n",
      "id: 5396096 rate: 15858.52428588172 per second\n",
      "id: 5396096 rate: 15518.8712898936 per second\n",
      "id: 5396096 rate: 15134.589790993123 per second\n",
      "id: 5396096 rate: 14829.895170597232 per second\n",
      "id: 5396096 rate: 14561.431348525493 per second\n",
      "id: 5396096 rate: 14271.24971074287 per second\n",
      "id: 5396096 rate: 13990.64671783938 per second\n",
      "id: 5396096 rate: 13664.236710129706 per second\n"
     ]
    }
   ],
   "source": [
    "k = 100\n",
    "import faiss\n",
    "import torch\n",
    "\n",
    "res = faiss.StandardGpuResources()\n",
    "res.setDefaultNullStreamAllDevices()\n",
    "query_vectors = np.array(query_vectors)\n",
    "results = ResultHeap(len(query_vectors), k)\n",
    "start = time.time()\n",
    "xq_t = torch.from_numpy(query_vectors).cuda()\n",
    "\n",
    "import glob\n",
    "for file in glob.glob('./fever-embs/*.npy'):\n",
    "    arr = np.load(file)\n",
    "    ids = arr[:, 0].astype(np.int64)\n",
    "    vectors = np.ascontiguousarray(arr[:, 1:].astype(np.float32))\n",
    "    xb_t = torch.from_numpy(vectors).cuda()\n",
    "    D, I = search_raw_array_pytorch(res, xb_t, xq_t, k)\n",
    "    I = I.cpu().numpy()\n",
    "    I = np.ascontiguousarray(ids[I])\n",
    "    results.add_batch_result(D.cpu().numpy(), I , 0)\n",
    "    print(f'id: {i0} rate: {i0/(time.time() - start)} per second')\n",
    "results.finalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "from psycopg2.extensions import register_adapter, AsIs\n",
    "psycopg2.extensions.register_adapter(np.int64, psycopg2._psycopg.AsIs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "template = \"INSERT INTO test.vector_benchmark (claim_id, article_ids) VALUES %s\"\n",
    "buffer = []\n",
    "conn = psycopg2.connect(POSTGRES_DSN)\n",
    "cur = conn.cursor()\n",
    "for i, (claim_id, res) in enumerate(zip(claim_ids, results.I)):\n",
    "    buffer.append((claim_id, list(res)))\n",
    "    if (i+1) % 5000 == 0:\n",
    "        extras.execute_values(cur, template, buffer)\n",
    "        conn.commit()\n",
    "        buffer = []\n",
    "extras.execute_values(cur,template,buffer)\n",
    "buffer = []\n",
    "conn.commit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = faiss.index_factory(d, \"IDMap,Flat\") # no compression/ loss\n",
    "# index = faiss.index_factory(d, \"PCAR64,IVF65536_HNSW32,SQ8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "xb = np.ascontiguousarray(xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-13-74d861dfd6c2>\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[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mxb\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      2\u001b[0m \u001b[0mindex\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mxb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0marr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mquery\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mxb\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m5\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/anaconda3/lib/python3.7/site-packages/faiss/__init__.py\u001b[0m in \u001b[0;36mreplacement_train\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    140\u001b[0m         \u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md\u001b[0m \u001b[0;34m=\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[0m\n\u001b[1;32m    141\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0md\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[0m\n\u001b[0;32m--> 142\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_c\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mn\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    143\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    144\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mreplacement_search\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mk\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/opt/anaconda3/lib/python3.7/site-packages/faiss/swigfaiss_avx2.py\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(self, n, x)\u001b[0m\n\u001b[1;32m   1804\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1805\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\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[0;32m-> 1806\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0m_swigfaiss_avx2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mIndexPreTransform_train\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\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[0m\n\u001b[0m\u001b[1;32m   1807\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1808\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0madd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\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[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "index.train(xb)\n",
    "index.add(xb)\n",
    "arr = None\n",
    "query = xb[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[    0  6311 72709 18162]\n",
      " [    1 59766 58111 52040]\n",
      " [    2 31479 58959  9319]\n",
      " [    3 35738 64225 74826]\n",
      " [    4 96100 29453 16855]]\n",
      "[[ 63.39306   64.44453   65.00524   65.72736 ]\n",
      " [ 43.355114  43.770733  43.770733  43.770733]\n",
      " [146.49068  172.00456  172.46603  175.20573 ]\n",
      " [173.89133  185.32256  193.85425  193.94987 ]\n",
      " [ 79.9476    83.01735   83.22531   83.37677 ]]\n"
     ]
    }
   ],
   "source": [
    "k = 4                          # we want to see 4 nearest neighbors\n",
    "D, I = index.search(query, k) # sanity check\n",
    "print(I)\n",
    "print(D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}