notebooks/3 - Train Language Model Using FastAI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Environment is available as a publically available docker container: `hamelsmu/ml-gpu`\n",
    "\n",
    "### Pre-Requisite: Make Sure you have the right files prepared from Step 1\n",
    "\n",
    "You should have these files in the root of the `./data/processed_data/` directory:\n",
    "\n",
    "1. `{train/valid/test.function}` - these are python function definitions tokenized (by space), 1 line per function.\n",
    "2. `{train/valid/test.docstring}` - these are docstrings that correspond to each of the python function definitions, and have a 1:1 correspondence with the lines in *.function files.\n",
    "3. `{train/valid/test.lineage}` - every line in this file contains a link back to the original location (github repo link) where the code was retrieved.  There is a 1:1 correspondence with the lines in this file and the other two files. This is useful for debugging.\n",
    "\n",
    "\n",
    "### Set the value of `use_cache` appropriately.  \n",
    "\n",
    "If `use_cache = True`, data will be downloaded where possible instead of re-computing.  However, it is highly recommended that you set `use_cache = False` for this tutorial as it will be less confusing, and you will learn more by runing these steps yourself.  This notebook was run on AWS on a `p3.8xlarge` in approximately 8 hours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Optional: you can set what GPU you want to use in a notebook like this.  \n",
    "# # Useful if you want to run concurrent experiments at the same time on different GPUs.\n",
    "# import os\n",
    "# os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_cache = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Download pre-processed data if you want to run from tutorial from this step.##\n",
    "from general_utils import get_step2_prerequisite_files\n",
    "\n",
    "if use_cache:\n",
    "    get_step2_prerequisite_files(output_directory = './data/processed_data')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build Language Model From Docstrings\n",
    "\n",
    "The goal is to build a language model using the docstrings, and use that language model to generate an embedding for each docstring.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n",
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import torch,cv2\n",
    "from lang_model_utils import lm_vocab, load_lm_vocab, train_lang_model\n",
    "from general_utils import save_file_pickle, load_file_pickle\n",
    "import logging\n",
    "from pathlib import Path\n",
    "from fastai.text import *\n",
    "\n",
    "source_path = Path('./data/processed_data/')\n",
    "\n",
    "with open(source_path/'train.docstring', 'r') as f:\n",
    "    trn_raw = f.readlines()\n",
    "\n",
    "with open(source_path/'valid.docstring', 'r') as f:\n",
    "    val_raw = f.readlines()\n",
    "    \n",
    "with open(source_path/'test.docstring', 'r') as f:\n",
    "    test_raw = f.readlines()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Preview what the raw data looks like: here are 10 docstrings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['\"generates batches of samples : param data : array - like , shape = ( n_samples , n_features ) : param labels : array - like , shape = ( n_samples , ) : return :\"\\n',\n",
       " '\"converts labels as single integer to row vectors . for instance , given a three class problem , labels would be mapped as label_1 : [ 1 0 0 ] , label_2 : [ 0 1 0 ] , label_3 : [ 0 , 0 , 1 ] where labels can be either int or string . : param labels : array - like , shape = ( n_samples , ) : return :\"\\n',\n",
       " '\"sigmoid function . : param x : array - like , shape = ( n_features , ) : return :\"\\n',\n",
       " '\"compute sigmoid first derivative . : param x : array - like , shape = ( n_features , ) : return :\"\\n',\n",
       " '\"rectified linear function . : param x : array - like , shape = ( n_features , ) : return :\"\\n',\n",
       " '\"rectified linear first derivative . : param x : array - like , shape = ( n_features , ) : return :\"\\n',\n",
       " '\"hyperbolic tangent function . : param x : array - like , shape = ( n_features , ) : return :\"\\n',\n",
       " '\"hyperbolic tangent first derivative . : param x : array - like , shape = ( n_features , ) : return :\"\\n',\n",
       " 'blueflask command line .\\n',\n",
       " 'args : name ( str ) : app name . directory ( str ) : path of the directory where the project is going to be created . defaults to the current directory .\\n']"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trn_raw[:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-process data for language model\n",
    "\n",
    "We will use the class  `build_lm_vocab` to prepare our data for the language model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Processing 1,021,954 rows\n",
      "WARNING:root:Vocab Size 23,687\n",
      "WARNING:root:Transforming 1,021,954 rows.\n",
      "WARNING:root:Removed 152,325 duplicate rows.\n"
     ]
    }
   ],
   "source": [
    "vocab = lm_vocab(max_vocab=50000,\n",
    "                 min_freq=10)\n",
    "\n",
    "# fit the transform on the training data, then transform\n",
    "trn_flat_idx = vocab.fit_transform_flattened(trn_raw)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at the transformed data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([   2,   61,   61,   90, 1567,   70,   61,   61,    3,  862])"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trn_flat_idx[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['_xbos_', '*', '*', '[', 'abstract', ']', '*', '*', 'the', 'real']"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[vocab.itos[x] for x in trn_flat_idx[:10]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Transforming 206,035 rows.\n",
      "WARNING:root:Removed 22,179 duplicate rows.\n"
     ]
    }
   ],
   "source": [
    "# apply transform to validation data\n",
    "val_flat_idx = vocab.transform_flattened(val_raw)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save files for later use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Saved vocab to data/lang_model/vocab_v2.cls\n"
     ]
    }
   ],
   "source": [
    "if not use_cache:\n",
    "    vocab.save('./data/lang_model/vocab_v2.cls')\n",
    "    save_file_pickle('./data/lang_model/trn_flat_idx_list.pkl_v2', trn_flat_idx)\n",
    "    save_file_pickle('./data/lang_model/val_flat_idx_list.pkl_v2', val_flat_idx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Fast.AI Language Model\n",
    "\n",
    "This model will read in files that were created and train a [fast.ai](https://github.com/fastai/fastai/tree/master/fastai) language model.  This model learns to predict the next word in the sentence using fast.ai's implementation of [AWD LSTM](https://github.com/salesforce/awd-lstm-lm).  \n",
    "\n",
    "The goal of training this model is to build a general purpose feature extractor for text that can be used in downstream models.  In this case, we will utilize this model to produce embeddings for function docstrings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Loaded vocab of size 49,811\n"
     ]
    }
   ],
   "source": [
    "vocab = load_lm_vocab('./data/lang_model/vocab.cls')\n",
    "trn_flat_idx = load_file_pickle('./data/lang_model/trn_flat_idx_list.pkl')\n",
    "val_flat_idx = load_file_pickle('./data/lang_model/val_flat_idx_list.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7da5d6e59cf64c4facb2c5d7cc518441",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Epoch', max=7), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch      trn_loss   val_loss                                \n",
      "    0      4.294783   4.28535   \n",
      "    1      4.127489   4.140274                                \n",
      "    2      4.069605   4.078538                                \n",
      "    3      4.024407   4.046905                                \n",
      "    4      4.004515   4.024419                                \n",
      "    5      3.982315   4.01069                                 \n",
      "                                                              \r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:State dict for the best model saved here:\n",
      "data/lang_model_weights_v2/models/langmodel_best.h5\n",
      "WARNING:root:Not re-training language model because use_cache=True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    6      3.971787   3.998788  \n",
      "\n"
     ]
    }
   ],
   "source": [
    "if not use_cache:\n",
    "    fastai_learner, lang_model = train_lang_model(model_path = './data/lang_model_weights_v2',\n",
    "                                                  trn_indexed = trn_flat_idx,\n",
    "                                                  val_indexed = val_flat_idx,\n",
    "                                                  vocab_size = vocab.vocab_size,\n",
    "                                                  lr=3e-3,\n",
    "                                                  em_sz= 500,\n",
    "                                                  nh= 500,\n",
    "                                                  bptt=20,\n",
    "                                                  cycle_len=1,\n",
    "                                                  n_cycle=3,\n",
    "                                                  cycle_mult=2,\n",
    "                                                  bs = 200,\n",
    "                                                  wd = 1e-6)\n",
    "    \n",
    "elif use_cache:    \n",
    "    logging.warning('Not re-training language model because use_cache=True')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "346e8a550de242ce83509287b3fa3473",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Epoch', max=6), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch      trn_loss   val_loss                                \n",
      "    0      3.954703   3.989164  \n",
      "    1      3.907728   3.975681                                \n",
      "    2      3.936994   3.976287                                \n",
      "    3      3.871557   3.96412                                 \n",
      "    4      3.927649   3.969976                                \n",
      "    5      3.873011   3.956639                                \n",
      "\n"
     ]
    }
   ],
   "source": [
    "if not use_cache:\n",
    "    fastai_learner.fit(1e-3, 3, wds=1e-6, cycle_len=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f3f2ad969a2742a1ba089d8a88fc560d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Epoch', max=9), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch      trn_loss   val_loss                                \n",
      "    0      3.925804   3.971093  \n",
      "    1      3.857519   3.951696                                \n",
      "    2      3.840948   3.946251                                \n",
      "    3      3.907309   3.970567                                \n",
      "    4      3.879899   3.956719                                \n",
      "    5      3.840587   3.947983                                \n",
      "    6      3.823401   3.935096                                \n",
      "    7      3.838912   3.929217                                \n",
      "    8      3.778818   3.930717                                \n",
      "\n"
     ]
    }
   ],
   "source": [
    "if not use_cache:\n",
    "    fastai_learner.fit(1e-3, 2, wds=1e-6, cycle_len=3, cycle_mult=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8c4372e58847485e80d5404eb02a2129",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Epoch', max=33), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch      trn_loss   val_loss                                \n",
      "    0      3.86375    3.953147  \n",
      "    1      3.851326   3.930299                                \n",
      "    2      3.773453   3.927069                                \n",
      "    3      3.879102   3.957266                                \n",
      "    4      3.858202   3.954743                                \n",
      "    5      3.852824   3.951508                                \n",
      "    6      3.837561   3.9509                                  \n",
      "    7      3.818845   3.947756                                \n",
      "    8      3.809637   3.944036                                \n",
      "    9      3.835555   3.942263                                \n",
      "    10     3.824583   3.935868                                \n",
      "    11     3.827287   3.932043                                \n",
      "    12     3.817058   3.927741                                \n",
      "    13     3.778389   3.927357                                \n",
      "    14     3.779933   3.925774                                \n",
      "    15     3.780848   3.918761                                \n",
      "    16     3.746735   3.920191                                \n",
      "    17     3.743517   3.915674                                \n",
      "    18     3.752455   3.911835                                \n",
      "    19     3.758213   3.908067                                \n",
      "    20     3.768209   3.904584                                \n",
      "    21     3.711149   3.904635                                \n",
      "    22     3.770484   3.898746                                \n",
      "    23     3.767993   3.897296                                \n",
      "    24     3.707685   3.898568                                \n",
      "    25     3.694116   3.898346                                \n",
      "    26     3.749094   3.89368                                 \n",
      "    27     3.727432   3.894122                                \n",
      "    28     3.682065   3.89575                                 \n",
      "    29     3.712119   3.894845                                \n",
      "    30     3.721573   3.894399                                \n",
      "    31     3.668023   3.89601                                 \n",
      "    32     3.710865   3.896029                                \n",
      "\n"
     ]
    }
   ],
   "source": [
    "if not use_cache:\n",
    "    fastai_learner.fit(1e-3, 2, wds=1e-6, cycle_len=3, cycle_mult=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save language model and learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not use_cache:\n",
    "    fastai_learner.save('lang_model_learner_v2.fai')\n",
    "    lang_model_new = fastai_learner.model.eval()\n",
    "    torch.save(lang_model_new, './data/lang_model/lang_model_gpu_v2.torch')\n",
    "    torch.save(lang_model_new.cpu(), './data/lang_model/lang_model_cpu_v2.torch')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Model and Encode All Docstrings\n",
    "\n",
    "Now that we have trained the language model, the next step is to use the language model to encode all of the docstrings into a vector. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "** Note that checkpointed versions of the language model artifacts are available for download: **\n",
    "\n",
    "1. `lang_model_cpu_v2.torch` : https://storage.googleapis.com/kubeflow-examples/code_search/data/lang_model/lang_model_cpu_v2.torch \n",
    "2. `lang_model_gpu_v2.torch` : https://storage.googleapis.com/kubeflow-examples/code_search/data/lang_model/lang_model_gpu_v2.torch\n",
    "3. `vocab_v2.cls` : https://storage.googleapis.com/kubeflow-examples/code_search/data/lang_model/vocab_v2.cls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Loaded vocab of size 23,687\n",
      "WARNING:root:Processing 1,227,989 rows\n"
     ]
    }
   ],
   "source": [
    "from lang_model_utils import load_lm_vocab\n",
    "vocab = load_lm_vocab('./data/lang_model/vocab_v2.cls')\n",
    "idx_docs = vocab.transform(trn_raw + val_raw, max_seq_len=30, padding=False)\n",
    "lang_model = torch.load('./data/lang_model/lang_model_gpu_v2.torch', \n",
    "                        map_location=lambda storage, loc: storage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequentialRNN(\n",
       "  (0): RNN_Encoder(\n",
       "    (encoder): Embedding(23687, 500, padding_idx=1)\n",
       "    (encoder_with_dropout): EmbeddingDropout(\n",
       "      (embed): Embedding(23687, 500, padding_idx=1)\n",
       "    )\n",
       "    (rnns): ModuleList(\n",
       "      (0): WeightDrop(\n",
       "        (module): LSTM(500, 500)\n",
       "      )\n",
       "      (1): WeightDrop(\n",
       "        (module): LSTM(500, 500)\n",
       "      )\n",
       "      (2): WeightDrop(\n",
       "        (module): LSTM(500, 500)\n",
       "      )\n",
       "    )\n",
       "    (dropouti): LockedDropout(\n",
       "    )\n",
       "    (dropouths): ModuleList(\n",
       "      (0): LockedDropout(\n",
       "      )\n",
       "      (1): LockedDropout(\n",
       "      )\n",
       "      (2): LockedDropout(\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (1): LinearDecoder(\n",
       "    (decoder): Linear(in_features=500, out_features=23687, bias=False)\n",
       "    (dropout): LockedDropout(\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lang_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** the below code extracts embeddings for docstrings one docstring at a time, which is very inefficient.  Ideally, you want to extract embeddings in batch but account for the fact that you will have padding, etc. when extracting the hidden states.  For this tutorial, we only provide this minimal example, however you are welcome to improve upon this and sumbit a PR!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "def list2arr(l):\n",
    "    \"Convert list into pytorch Variable.\"\n",
    "    return V(np.expand_dims(np.array(l), -1)).cpu()\n",
    "\n",
    "def make_prediction_from_list(model, l):\n",
    "    \"\"\"\n",
    "    Encode a list of integers that represent a sequence of tokens.  The\n",
    "    purpose is to encode a sentence or phrase.\n",
    "\n",
    "    Parameters\n",
    "    -----------\n",
    "    model : fastai language model\n",
    "    l : list\n",
    "        list of integers, representing a sequence of tokens that you want to encode\n",
    "\n",
    "    \"\"\"\n",
    "    arr = list2arr(l)# turn list into pytorch Variable with bs=1\n",
    "    model.reset()  # language model is stateful, so you must reset upon each prediction\n",
    "    hidden_states = model(arr)[-1][-1] # RNN Hidden Layer output is last output, and only need the last layer\n",
    "\n",
    "    #return avg-pooling, max-pooling, and last hidden state\n",
    "    return hidden_states.mean(0), hidden_states.max(0)[0], hidden_states[-1]\n",
    "\n",
    "\n",
    "def get_embeddings(lm_model, list_list_int):\n",
    "    \"\"\"\n",
    "    Vectorize a list of sequences List[List[int]] using a fast.ai language model.\n",
    "\n",
    "    Paramters\n",
    "    ---------\n",
    "    lm_model : fastai language model\n",
    "    list_list_int : List[List[int]]\n",
    "        A list of sequences to encode\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    tuple: (avg, mean, last)\n",
    "        A tuple that returns the average-pooling, max-pooling over time steps as well as the last time step.\n",
    "    \"\"\"\n",
    "    n_rows = len(list_list_int)\n",
    "    n_dim = lm_model[0].nhid\n",
    "    avgarr = np.empty((n_rows, n_dim))\n",
    "    maxarr = np.empty((n_rows, n_dim))\n",
    "    lastarr = np.empty((n_rows, n_dim))\n",
    "\n",
    "    for i in tqdm_notebook(range(len(list_list_int))):\n",
    "        avg_, max_, last_ = make_prediction_from_list(lm_model, list_list_int[i])\n",
    "        avgarr[i,:] = avg_.data.numpy()\n",
    "        maxarr[i,:] = max_.data.numpy()\n",
    "        lastarr[i,:] = last_.data.numpy()\n",
    "\n",
    "    return avgarr, maxarr, lastarr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db326f7eccd14f7d819181fe371bb74b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=1227989), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "CPU times: user 3d 3h 8min 12s, sys: 4min 3s, total: 3d 3h 12min 15s\n",
      "Wall time: 4h 48min 48s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "avg_hs, max_hs, last_hs = get_embeddings(lang_model, idx_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Do the same thing for the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Processing 177,220 rows\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "50c64aed447f4592b82069463318633c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=177220), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "idx_docs_test = vocab.transform(test_raw, max_seq_len=30, padding=False)\n",
    "avg_hs_test, max_hs_test, last_hs_test = get_embeddings(lang_model, idx_docs_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save Language Model Embeddings For Docstrings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "savepath = Path('./data/lang_model_emb/')\n",
    "np.save(savepath/'avg_emb_dim500_v2.npy', avg_hs)\n",
    "np.save(savepath/'max_emb_dim500_v2.npy', max_hs)\n",
    "np.save(savepath/'last_emb_dim500_v2.npy', last_hs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the test set embeddings also\n",
    "np.save(savepath/'avg_emb_dim500_test_v2.npy', avg_hs_test)\n",
    "np.save(savepath/'max_emb_dim500_test_v2.npy', max_hs_test)\n",
    "np.save(savepath/'last_emb_dim500_test_v2.npy', last_hs_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "** Note that the embeddings saved to disk above have also been cached and are are available for download: **\n",
    "\n",
    "Train + Validation docstrings vectorized:\n",
    "\n",
    "1. `avg_emb_dim500_v2.npy` : https://storage.googleapis.com/kubeflow-examples/code_search/data/lang_model_emb/avg_emb_dim500_v2.npy\n",
    "2. `max_emb_dim500_v2.npy` : https://storage.googleapis.com/kubeflow-examples/code_search/data/lang_model_emb/last_emb_dim500_v2.npy\n",
    "3. `last_emb_dim500_v2.npy` : https://storage.googleapis.com/kubeflow-examples/code_search/data/lang_model_emb/max_emb_dim500_v2.npy\n",
    "\n",
    "Test set docstrings vectorized:\n",
    "\n",
    "1. `avg_emb_dim500_test_v2.npy`: https://storage.googleapis.com/kubeflow-examples/code_search/data/lang_model_emb/avg_emb_dim500_test_v2.npy\n",
    "\n",
    "2. `max_emb_dim500_test_v2.npy`: https://storage.googleapis.com/kubeflow-examples/code_search/data/lang_model_emb/last_emb_dim500_test_v2.npy\n",
    "\n",
    "3. `last_emb_dim500_test_v2.npy`: https://storage.googleapis.com/kubeflow-examples/code_search/data/lang_model_emb/max_emb_dim500_test_v2.npy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate Sentence Embeddings\n",
    "\n",
    "One popular way of evaluating sentence embeddings is to measure the efficacy of these embeddings in downstream tasks like sentiment analysis, textual similarity etc.  Usually you can use general-purpose benchmarks such as the examples outlined [here](https://github.com/facebookresearch/SentEval) to measure the quality of your embeddings.  However, since this is a very domain specific dataset - those general purpose benchmarks may not be appropriate.  Unfortunately, we have not designed downstream tasks that we can open source at this point.\n",
    "\n",
    "In the absence of these downstream tasks, we can at least sanity check that these embeddings contain semantic information by doing the following:\n",
    "\n",
    "1. Manually examine similarity between sentences, by supplying a statement and examining if the nearest phrase found is similar. \n",
    "\n",
    "2. Visualize the embeddings.\n",
    "\n",
    "We will do the first approach, and leave the second approach as an exercise for the reader.  **It should be noted that this is only a sanity check -- a more rigorous approach is to measure the impact of these embeddings on a variety of downstream tasks** and use that to form a more objective opinion about the quality of your embeddings.\n",
    "\n",
    "Furthermroe, there are many different ways of constructing a sentence embedding from the language model.  For example, we can take the average, the maximum or even the last value of the hidden states (or concatenate them all together).  **For simplicity, we will only evaluate the sentence embedding that is constructed by taking the average over the hidden states** (and leave other possibilities as an exercise for the reader). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create search index using `nmslib` \n",
    "\n",
    "[nmslib](https://github.com/nmslib/nmslib) is a great library for doing nearest neighbor lookups, which we will use as a search engine for finding nearest neighbors of comments in vector-space.  \n",
    "\n",
    "The convenience function `create_nmslib_search_index` builds this search index given a matrix of vectors as input.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "from general_utils import create_nmslib_search_index\n",
    "import nmslib\n",
    "from lang_model_utils import Query2Emb\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "from lang_model_utils import load_lm_vocab\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load matrix of vectors\n",
    "loadpath = Path('./data/lang_model_emb/')\n",
    "avg_emb_dim500 = np.load(loadpath/'avg_emb_dim500_test_v2.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build search index (takes about an hour on a p3.8xlarge)\n",
    "dim500_avg_searchindex = create_nmslib_search_index(avg_emb_dim500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save search index\n",
    "dim500_avg_searchindex.saveIndex('./data/lang_model_emb/dim500_avg_searchindex.nmslib')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that if you did not train your own language model and are downloading the pre-trained model artifacts instead, you can similarly download the pre-computed search index here: \n",
    "\n",
    "https://storage.googleapis.com/kubeflow-examples/code_search/data/lang_model_emb/dim500_avg_searchindex.nmslib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After you have built this search index with nmslib, you can do fast nearest-neighbor lookups.  We use the `Query2Emb` object to help convert strings to the embeddings: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "dim500_avg_searchindex = nmslib.init(method='hnsw', space='cosinesimil')\n",
    "dim500_avg_searchindex.loadIndex('./data/lang_model_emb/dim500_avg_searchindex.nmslib')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "lang_model = torch.load('./data/lang_model/lang_model_cpu_v2.torch')\n",
    "vocab = load_lm_vocab('./data/lang_model/vocab_v2.cls')\n",
    "\n",
    "q2emb = Query2Emb(lang_model = lang_model.cpu(),\n",
    "                  vocab = vocab)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The method `Query2Emb.emb_mean` will allow us to use the langauge model we trained earlier to generate a sentence embedding given a string.   Here is an example, `emb_mean` will return a numpy array of size (1, 500)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 500)"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = q2emb.emb_mean('Read data into pandas dataframe')\n",
    "query.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Make search engine to inspect semantic similarity of phrases**.  This will take 3 inputs:\n",
    "\n",
    "1. `nmslib_index` - this is the search index we built above.  This object takes a vector and will return the index of the closest vector(s) according to cosine distance.  \n",
    "2. `ref_data` - this is the data for which the index refer to, in this case will be the docstrings. \n",
    "3. `query2emb_func` - this is a function that will convert a string into an embedding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "class search_engine:\n",
    "    def __init__(self, \n",
    "                 nmslib_index, \n",
    "                 ref_data, \n",
    "                 query2emb_func):\n",
    "        \n",
    "        self.search_index = nmslib_index\n",
    "        self.data = ref_data\n",
    "        self.query2emb_func = query2emb_func\n",
    "    \n",
    "    def search(self, str_search, k=3):\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",
    "            print(f'cosine dist:{dist:.4f}\\n---------------\\n', self.data[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "se = search_engine(nmslib_index=dim500_avg_searchindex,\n",
    "                   ref_data = test_raw,\n",
    "                   query2emb_func = q2emb.emb_mean)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Manually Inspect Phrase Similarity\n",
    "\n",
    "Compare a user-supplied query vs. vectorized docstrings on test set.  We can see that similar phrases are not exactly the same, but the nearest neighbors are reasonable.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "logging.getLogger().setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cosine dist:0.0977\n",
      "---------------\n",
      " load csv or json into pandas dataframe\n",
      "\n",
      "cosine dist:0.1167\n",
      "---------------\n",
      " reads swn database into a dictionary\n",
      "\n",
      "cosine dist:0.1187\n",
      "---------------\n",
      " load csv files into a raw object .\n",
      "\n"
     ]
    }
   ],
   "source": [
    "se.search('read csv into pandas dataframe')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cosine dist:0.0561\n",
      "---------------\n",
      " train a classifier\n",
      "\n",
      "cosine dist:0.0607\n",
      "---------------\n",
      " train a network\n",
      "\n",
      "cosine dist:0.0800\n",
      "---------------\n",
      " train a selection class\n",
      "\n"
     ]
    }
   ],
   "source": [
    "se.search('train a random forest')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cosine dist:0.0533\n",
      "---------------\n",
      " download contratos files\n",
      "\n",
      "cosine dist:0.0591\n",
      "---------------\n",
      " download s3 files\n",
      "\n",
      "cosine dist:0.0678\n",
      "---------------\n",
      " download data files .\n",
      "\n"
     ]
    }
   ],
   "source": [
    "se.search('download files')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cosine dist:0.0370\n",
      "---------------\n",
      " start server agent\n",
      "\n",
      "cosine dist:0.0476\n",
      "---------------\n",
      " start openvswitch service\n",
      "\n",
      "cosine dist:0.0571\n",
      "---------------\n",
      " start acestream engine\n",
      "\n"
     ]
    }
   ],
   "source": [
    "se.search('start webserver')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cosine dist:0.0559\n",
      "---------------\n",
      " send email notification .\n",
      "\n",
      "cosine dist:0.0692\n",
      "---------------\n",
      " send out stats request .\n",
      "\n",
      "cosine dist:0.0692\n",
      "---------------\n",
      " send out stats request .\n",
      "\n"
     ]
    }
   ],
   "source": [
    "se.search('send out email notification')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cosine dist:0.0460\n",
      "---------------\n",
      " save json file\n",
      "\n",
      "cosine dist:0.0734\n",
      "---------------\n",
      " save image definitions\n",
      "\n",
      "cosine dist:0.0888\n",
      "---------------\n",
      " save library template\n",
      "\n"
     ]
    }
   ],
   "source": [
    "se.search('save pickle file')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize Embeddings (Optional)\n",
    "\n",
    "We highly recommend using [tensorboard](https://www.tensorflow.org/versions/r1.0/get_started/embedding_viz) as way to visualize embeddings.  Tensorboard contains an interactive search that makes it easy (and fun) to explore embeddings.  We leave this as an exercise to the reader."
   ]
  },
  {
   "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
}