notebooks/5 - Build Search Index.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prerequisites\n",
    "\n",
    "You should have completed steps 1-4 of this tutorial before beginning this exercise.  The files required for this notebook are generated by those previous steps.\n",
    "\n",
    "Creating the search engine for this example is extremely CPU and memory intensive.  We used an an AWS `x1.32xlarge` instance (128 cores) in order to achieve the maximum speed with building the search index. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import nmslib\n",
    "from lang_model_utils import load_lm_vocab, Query2Emb\n",
    "from general_utils import create_nmslib_search_index\n",
    "\n",
    "input_path = Path('./data/processed_data/')\n",
    "code2emb_path = Path('./data/code2emb/')\n",
    "output_path = Path('./data/search')\n",
    "output_path.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read in Metadata\n",
    "\n",
    "We will want to organize the data that we will want to display for the search results, which will be:\n",
    "\n",
    "1. The original code\n",
    "2. A link to the original code\n",
    "\n",
    "For convenience, we will collect this data into a pandas dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>url</th>\n",
       "      <th>code</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>https://github.com/fnl/libfnl/blob/master/src/...</td>\n",
       "      <td>def __init__(self, *leafs, **edges):\\n    self...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>https://github.com/fnl/libfnl/blob/master/src/...</td>\n",
       "      <td>def __eq__(self, other):\\n    if isinstance(ot...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>https://github.com/fnl/libfnl/blob/master/src/...</td>\n",
       "      <td>def __repr__(self):\\n    return 'Node&lt;leafs={}...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>https://github.com/fnl/libfnl/blob/master/src/...</td>\n",
       "      <td>@staticmethod\\ndef _isCapitalized(token):\\n   ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>https://github.com/fnl/libfnl/blob/master/src/...</td>\n",
       "      <td>@staticmethod\\ndef _isCapitalizeD(last, token)...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                 url  \\\n",
       "0  https://github.com/fnl/libfnl/blob/master/src/...   \n",
       "1  https://github.com/fnl/libfnl/blob/master/src/...   \n",
       "2  https://github.com/fnl/libfnl/blob/master/src/...   \n",
       "3  https://github.com/fnl/libfnl/blob/master/src/...   \n",
       "4  https://github.com/fnl/libfnl/blob/master/src/...   \n",
       "\n",
       "                                                code  \n",
       "0  def __init__(self, *leafs, **edges):\\n    self...  \n",
       "1  def __eq__(self, other):\\n    if isinstance(ot...  \n",
       "2  def __repr__(self):\\n    return 'Node<leafs={}...  \n",
       "3  @staticmethod\\ndef _isCapitalized(token):\\n   ...  \n",
       "4  @staticmethod\\ndef _isCapitalizeD(last, token)...  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# read file of urls\n",
    "url_df = pd.read_csv(input_path/'without_docstrings.lineage', header=None, names=['url'])\n",
    "\n",
    "\n",
    "# read original code\n",
    "code_df = pd.read_json(input_path/'without_docstrings_original_function.json.gz')\n",
    "code_df.columns = ['code']\n",
    "\n",
    "# make sure these files have same number of rows\n",
    "assert code_df.shape[0] == url_df.shape[0]\n",
    "\n",
    "# collect these two together into a dataframe\n",
    "ref_df = pd.concat([url_df, code_df], axis = 1).reset_index(drop=True)\n",
    "ref_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For reference the above files are also available for download incase you skipped step 1:\n",
    "\n",
    "`without_docstrings.lineage`: https://storage.googleapis.com/kubeflow-examples/code_search/data/without_docstrings.lineage\n",
    "\n",
    "`without_docstrings_original_function.json.gz`: https://storage.googleapis.com/kubeflow-examples/code_search/data/without_docstrings_original_function.json.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Search Index For Vectorized Code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First read in the vectorized code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "nodoc_vecs = np.load(code2emb_path/'nodoc_vecs.npy')\n",
    "assert nodoc_vecs.shape[0] == ref_df.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now build the search index. **Warning:** this step takes ~ 18 minutes on an `x1.32xlarge` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1d 7h 37min 53s, sys: 3min, total: 1d 7h 40min 53s\n",
      "Wall time: 17min 33s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "search_index = create_nmslib_search_index(nodoc_vecs)\n",
    "search_index.saveIndex('./data/search/search_index.nmslib')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This cached version of this index can be downloaded here:  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create A Minimal Search Engine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can find the cached version of the required files on google cloud:\n",
    "\n",
    "`lang_model_cpu_v2.torch`: https://storage.googleapis.com/kubeflow-examples/code_search/data/lang_model/lang_model_cpu_v2.torch\n",
    "\n",
    "`vocab_v2.cls`: https://storage.googleapis.com/kubeflow-examples/code_search/data/lang_model/vocab_v2.cls\n",
    "\n",
    "`search_index.nmslib`: https://storage.googleapis.com/kubeflow-examples/code_search/data/search/search_index.nmslib\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Loaded vocab of size 23,687\n",
      "WARNING:root:Processing 1 rows\n"
     ]
    }
   ],
   "source": [
    "lang_model = torch.load('./data/lang_model/lang_model_cpu_v2.torch', \n",
    "                        map_location=lambda storage, loc: storage)\n",
    "\n",
    "vocab = load_lm_vocab('./data/lang_model/vocab_v2.cls')\n",
    "q2emb = Query2Emb(lang_model = lang_model.cpu(),\n",
    "                  vocab = vocab)\n",
    "\n",
    "search_index = nmslib.init(method='hnsw', space='cosinesimil')\n",
    "search_index.loadIndex('./data/search/search_index.nmslib')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Query2Emb` is a helper class that will vectorize sentences using the language model trained in Part 3.  \n",
    "\n",
    "In this case, we call the method `emb_mean` because we are taking the mean over the time steps of the hidden states in order to construct a sentence embedding for the query supplied by the user.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Processing 1 rows\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(1, 500)"
      ]
     },
     "execution_count": 177,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test = q2emb.emb_mean('Hello World!  This is a test.')\n",
    "test.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an object to make the process of showing search results easier\n",
    "\n",
    "The below object organizes all the pieces together for searching the index and displaying the results with a method call.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "metadata": {},
   "outputs": [],
   "source": [
    "class search_engine:\n",
    "    \"\"\"Organizes all the necessary elements we need to make a search engine.\"\"\"\n",
    "    def __init__(self, \n",
    "                 nmslib_index, \n",
    "                 ref_df, \n",
    "                 query2emb_func):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ==========\n",
    "        nmslib_index : nmslib object\n",
    "            This is pre-computed search index.\n",
    "        ref_df : pandas.DataFrame\n",
    "            This dataframe contains meta-data for search results, \n",
    "            must contain the columns 'code' and 'url'.\n",
    "        query2emb_func : callable\n",
    "            This is a function that takes as input a string and returns a vector\n",
    "            that is in the same vector space as what is loaded into the search index.\n",
    "\n",
    "        \"\"\"\n",
    "        assert 'url' in ref_df.columns\n",
    "        assert 'code' in ref_df.columns\n",
    "        \n",
    "        self.search_index = nmslib_index\n",
    "        self.ref_df = ref_df\n",
    "        self.query2emb_func = query2emb_func\n",
    "    \n",
    "    def search(self, str_search, k=2):\n",
    "        \"\"\"\n",
    "        Prints the code that are the nearest neighbors (by cosine distance)\n",
    "        to the search query.\n",
    "        \n",
    "        Parameters\n",
    "        ==========\n",
    "        str_search : str\n",
    "            a search query.  Ex: \"read data into pandas dataframe\"\n",
    "        k : int\n",
    "            the number of nearest neighbors to return.  Defaults to 2.\n",
    "        \n",
    "        \"\"\"\n",
    "        query = self.query2emb_func(str_search)\n",
    "        idxs, dists = self.search_index.knnQuery(query, k=k)\n",
    "        \n",
    "        for idx, dist in zip(idxs, dists):\n",
    "            code = self.ref_df.iloc[idx].code\n",
    "            url = self.ref_df.iloc[idx].url\n",
    "            print(f'cosine dist:{dist:.4f}  url: {url}\\n---------------\\n')\n",
    "            print(code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "metadata": {},
   "outputs": [],
   "source": [
    "se = search_engine(nmslib_index=search_index,\n",
    "                   ref_df=ref_df,\n",
    "                   query2emb_func=q2emb.emb_mean)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run Some Queries Against The Index!!\n",
    "\n",
    "Now that we have instantiated the search engine, we can use the `search` method to display the results.\n",
    "\n",
    "**Warning:** some of the displayed links may not work since this is historical data retrieved from a [historical open dataset Google has hosted on BigQuery](https://cloud.google.com/bigquery/public-data/github)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Processing 1 rows\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cosine dist:0.1045  url: https://github.com/timvandermeij/sentiment-analysis/blob/master/plot.py#L23\n",
      "---------------\n",
      "\n",
      "def read_data(self, cols):\n",
      "    return pd.read_csv(self.data_file, delimiter='\\t', names=cols)\n",
      "\n",
      "cosine dist:0.1072  url: https://github.com/littlezz/ESL-Model/blob/master/tests/utils.py#L3\n",
      "---------------\n",
      "\n",
      "def read_data(filename):\n",
      "    df = read_csv(filename, sep='\\t')\n",
      "    return df\n",
      "\n"
     ]
    }
   ],
   "source": [
    "se.search('read data into pandas dataframe')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use Custom Ipython Magic Function To Create A Fake Search Box\n",
    "\n",
    "You don't know how to build a website?  No problem!  You can still impress your friends by using a [custom magic function](https://ipython.org/ipython-doc/3/config/custommagics.html) to allow you to do a live demonstration in a Jupyter notebook.  This is what I did when I first created this prototype!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.magic import (register_line_magic, register_cell_magic,\n",
    "                                register_line_cell_magic)\n",
    "@register_cell_magic\n",
    "def search(line, cell):\n",
    "    return se.search(cell)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Live Semantic Search of Code (Searching Holdout Set Only)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%search\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}