notebooks/1 - Preprocess Data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preprocess Data\n",
    "This notebook contains materials to parse raw python files into function and docstring pairs, tokenize both function and dosctring into tokens, and split these pairs into a train, valid and test set.  \n",
    "\n",
    "*This step is optional, as we provide links to download pre-processed data at various points in the tutorial.  However, you might find it useful to go through these steps in order to understand how the data is prepared.*\n",
    "\n",
    "If you are using the recommended approach of using a `p3.8xlarge` instance for this entire tutorial you can use this docker container to run this notebook: [hamelsmu/ml-gpu](https://hub.docker.com/r/hamelsmu/ml-gpu/).\n",
    "\n",
    "Alternatively, if you wish to speed up *this notebook* by using an instance with lots of cores (because everything in this notebook is CPU bound), you can use this container [hamelsmu/ml-cpu](https://hub.docker.com/r/hamelsmu/ml-gpu/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import ast\n",
    "import glob\n",
    "import re\n",
    "from pathlib import Path\n",
    "\n",
    "import astor\n",
    "import pandas as pd\n",
    "import spacy\n",
    "from tqdm import tqdm\n",
    "from nltk.tokenize import RegexpTokenizer\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from general_utils import apply_parallel, flattenlist\n",
    "\n",
    "EN = spacy.load('en')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Python 3.6.3 :: Anaconda custom (64-bit)\r\n"
     ]
    }
   ],
   "source": [
    "! python -V"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download and read  raw python files\n",
    "\n",
    "The first thing we will want to do is to gather python code.  There is an open dataset that Google hosts on [BigQuery](https://cloud.google.com/bigquery/) that has code from open source projects on Github.  You can use [bigquery](https://cloud.google.com/bigquery/) to get the python files as a tabular dataset by executing the following SQL query in the bigquery console:\n",
    "\n",
    "```{sql}\n",
    "SELECT \n",
    " max(concat(f.repo_name, ' ', f.path)) as repo_path,\n",
    " c.content\n",
    "FROM `bigquery-public-data.github_repos.files` as f\n",
    "JOIN `bigquery-public-data.github_repos.contents` as c on f.id = c.id\n",
    "JOIN (\n",
    "      --this part of the query makes sure repo is watched at least twice since 2017\n",
    "      SELECT repo FROM(\n",
    "        SELECT \n",
    "          repo.name as repo\n",
    "        FROM `githubarchive.year.2017` WHERE type=\"WatchEvent\"\n",
    "        UNION ALL\n",
    "        SELECT \n",
    "          repo.name as repo\n",
    "        FROM `githubarchive.month.2018*` WHERE type=\"WatchEvent\"\n",
    "        )\n",
    "      GROUP BY 1\n",
    "      HAVING COUNT(*) >= 2\n",
    "      ) as r on f.repo_name = r.repo\n",
    "WHERE \n",
    "  f.path like '%.py' and --with python extension\n",
    "  c.size < 15000 and --get rid of ridiculously long files\n",
    "  REGEXP_CONTAINS(c.content, r'def ') --contains function definition\n",
    "group by c.content\n",
    "```\n",
    "\n",
    "\n",
    "Here is a link to the [SQL Query](https://bigquery.cloud.google.com/savedquery/506213277345:009fa66f301240e5ad9e4006c59a4762) incase it is helpful.  The raw data contains approximate 1.2 million distinct python code files.\n",
    "\n",
    "**To make things easier for this tutorial, the folks on the Google [Kubeflow team](https://kubernetes.io/blog/2017/12/introducing-kubeflow-composable/) have hosted the raw data for this tutorial in the form of 10 csv files, available at the url: https://storage.googleapis.com/kubeflow-examples/code_search/raw_data/00000000000{i}.csv as illustrated in the below code:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 1s, sys: 10.6 s, total: 1min 12s\n",
      "Wall time: 1min 31s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Read the data into a pandas dataframe, and parse out some meta-data\n",
    "\n",
    "df = pd.concat([pd.read_csv(f'https://storage.googleapis.com/kubeflow-examples/code_search/raw_data/00000000000{i}.csv') \\\n",
    "                for i in range(10)])\n",
    "\n",
    "df['nwo'] = df['repo_path'].apply(lambda r: r.split()[0])\n",
    "df['path'] = df['repo_path'].apply(lambda r: r.split()[1])\n",
    "df.drop(columns=['repo_path'], inplace=True)\n",
    "df = df[['nwo', 'path', 'content']]\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1241664, 3)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Inspect shape of the raw data\n",
    "df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions to parse data and tokenize\n",
    "\n",
    "Our goal is to parse the python files into (code, docstring) pairs.  Fortunately, the standard library in python comes with the wonderful [ast](https://docs.python.org/3.6/library/ast.html) module which helps us extract code from files as well as extract docstrings.  \n",
    "\n",
    "We also use the [astor](http://astor.readthedocs.io/en/latest/) library to strip the code of comments by doing a round trip of converting the code to an [AST](https://en.wikipedia.org/wiki/Abstract_syntax_tree) and then from AST back to code. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def tokenize_docstring(text):\n",
    "    \"Apply tokenization using spacy to docstrings.\"\n",
    "    tokens = EN.tokenizer(text)\n",
    "    return [token.text.lower() for token in tokens if not token.is_space]\n",
    "\n",
    "\n",
    "def tokenize_code(text):\n",
    "    \"A very basic procedure for tokenizing code strings.\"\n",
    "    return RegexpTokenizer(r'\\w+').tokenize(text)\n",
    "\n",
    "\n",
    "def get_function_docstring_pairs(blob):\n",
    "    \"Extract (function/method, docstring) pairs from a given code blob.\"\n",
    "    pairs = []\n",
    "    try:\n",
    "        module = ast.parse(blob)\n",
    "        classes = [node for node in module.body if isinstance(node, ast.ClassDef)]\n",
    "        functions = [node for node in module.body if isinstance(node, ast.FunctionDef)]\n",
    "        for _class in classes:\n",
    "            functions.extend([node for node in _class.body if isinstance(node, ast.FunctionDef)])\n",
    "\n",
    "        for f in functions:\n",
    "            source = astor.to_source(f)\n",
    "            docstring = ast.get_docstring(f) if ast.get_docstring(f) else ''\n",
    "            function = source.replace(ast.get_docstring(f, clean=False), '') if docstring else source\n",
    "\n",
    "            pairs.append((f.name,\n",
    "                          f.lineno,\n",
    "                          source,\n",
    "                          ' '.join(tokenize_code(function)),\n",
    "                          ' '.join(tokenize_docstring(docstring.split('\\n\\n')[0]))\n",
    "                         ))\n",
    "    except (AssertionError, MemoryError, SyntaxError, UnicodeEncodeError):\n",
    "        pass\n",
    "    return pairs\n",
    "\n",
    "\n",
    "def get_function_docstring_pairs_list(blob_list):\n",
    "    \"\"\"apply the function `get_function_docstring_pairs` on a list of blobs\"\"\"\n",
    "    return [get_function_docstring_pairs(b) for b in blob_list]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The below convience function `apply_parallel` parses the code in parallel using process based threading.  Adjust the `cpu_cores` parameter accordingly to your system resources!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 41.7 s, sys: 20 s, total: 1min 1s\n",
      "Wall time: 4min 36s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "pairs = flattenlist(apply_parallel(get_function_docstring_pairs_list, df.content.tolist(), cpu_cores=32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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>nwo</th>\n",
       "      <th>path</th>\n",
       "      <th>content</th>\n",
       "      <th>pairs</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>fnl/libfnl</td>\n",
       "      <td>src/fnl/nlp/dictionary.py</td>\n",
       "      <td>\"\"\"\\n.. py:module:: fnl.text.dictionary\\n   :s...</td>\n",
       "      <td>[(__init__, 19, def __init__(self, *leafs, **e...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>KivApple/mcu-info-util</td>\n",
       "      <td>mcu_info_util/linker_script.py</td>\n",
       "      <td>from six import iteritems\\n\\n\\ndef generate(op...</td>\n",
       "      <td>[(generate, 4, def generate(options, filename=...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Yelp/pyleus</td>\n",
       "      <td>examples/bandwith_monitoring/bandwith_monitori...</td>\n",
       "      <td>from __future__ import absolute_import, divisi...</td>\n",
       "      <td>[(__init__, 18, def __init__(self, size):\\n   ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>jhuapl-boss/boss-manage</td>\n",
       "      <td>bin/bearer_token.py</td>\n",
       "      <td>#!/usr/bin/env python3\\n\\n# Copyright 2016 The...</td>\n",
       "      <td>[(request, 46, def request(url, params=None, h...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>djfroofy/beatlounge</td>\n",
       "      <td>bl/orchestra/base.py</td>\n",
       "      <td>from itertools import cycle\\n\\nfrom twisted.py...</td>\n",
       "      <td>[(schedule, 149, def schedule(time, func, args...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                       nwo                                               path  \\\n",
       "0               fnl/libfnl                          src/fnl/nlp/dictionary.py   \n",
       "1   KivApple/mcu-info-util                     mcu_info_util/linker_script.py   \n",
       "2              Yelp/pyleus  examples/bandwith_monitoring/bandwith_monitori...   \n",
       "3  jhuapl-boss/boss-manage                                bin/bearer_token.py   \n",
       "4      djfroofy/beatlounge                               bl/orchestra/base.py   \n",
       "\n",
       "                                             content  \\\n",
       "0  \"\"\"\\n.. py:module:: fnl.text.dictionary\\n   :s...   \n",
       "1  from six import iteritems\\n\\n\\ndef generate(op...   \n",
       "2  from __future__ import absolute_import, divisi...   \n",
       "3  #!/usr/bin/env python3\\n\\n# Copyright 2016 The...   \n",
       "4  from itertools import cycle\\n\\nfrom twisted.py...   \n",
       "\n",
       "                                               pairs  \n",
       "0  [(__init__, 19, def __init__(self, *leafs, **e...  \n",
       "1  [(generate, 4, def generate(options, filename=...  \n",
       "2  [(__init__, 18, def __init__(self, size):\\n   ...  \n",
       "3  [(request, 46, def request(url, params=None, h...  \n",
       "4  [(schedule, 149, def schedule(time, func, args...  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert len(pairs) == df.shape[0], f'Row count mismatch. `df` has {df.shape[0]:,} rows; `pairs` has {len(pairs):,} rows.'\n",
    "df['pairs'] = pairs\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Flatten code, docstring pairs and extract meta-data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Flatten (code, docstring) pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6min 6s, sys: 8.88 s, total: 6min 15s\n",
      "Wall time: 6min 15s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# flatten pairs\n",
    "df = df.set_index(['nwo', 'path'])['pairs'].apply(pd.Series).stack()\n",
    "df = df.reset_index()\n",
    "df.columns = ['nwo', 'path', '_', 'pair']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extract meta-data and format dataframe.  \n",
    "\n",
    "We have not optimized this code.  Pull requests are welcome!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4min 38s, sys: 0 ns, total: 4min 38s\n",
      "Wall time: 4min 37s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "df['function_name'] = df['pair'].apply(lambda p: p[0])\n",
    "df['lineno'] = df['pair'].apply(lambda p: p[1])\n",
    "df['original_function'] = df['pair'].apply(lambda p: p[2])\n",
    "df['function_tokens'] = df['pair'].apply(lambda p: p[3])\n",
    "df['docstring_tokens'] = df['pair'].apply(lambda p: p[4])\n",
    "df = df[['nwo', 'path', 'function_name', 'lineno', 'original_function', 'function_tokens', 'docstring_tokens']]\n",
    "df['url'] = df[['nwo', 'path', 'lineno']].apply(lambda x: 'https://github.com/{}/blob/master/{}#L{}'.format(x[0], x[1], x[2]), axis=1)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Remove Duplicates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Removed 1,199,426 duplicate rows\n",
      "CPU times: user 26.5 s, sys: 0 ns, total: 26.5 s\n",
      "Wall time: 26.4 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# remove observations where the same function appears more than once\n",
    "before_dedup = len(df)\n",
    "df = df.drop_duplicates(['original_function', 'function_tokens'])\n",
    "after_dedup = len(df)\n",
    "\n",
    "print(f'Removed {before_dedup - after_dedup:,} duplicate rows')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5413927, 8)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Separate function w/o docstrings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def listlen(x):\n",
    "    if not isinstance(x, list):\n",
    "        return 0\n",
    "    return len(x)\n",
    "\n",
    "# separate functions w/o docstrings\n",
    "# docstrings should be at least 3 words in the docstring to be considered a valid docstring\n",
    "\n",
    "with_docstrings = df[df.docstring_tokens.str.split().apply(listlen) >= 3]\n",
    "without_docstrings = df[df.docstring_tokens.str.split().apply(listlen) < 3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Partition code by repository to minimize leakage between train, valid & test sets. \n",
    "Rough assumption that each repository has its own style.  We want to avoid having code from the same repository in the training set as well as the validation or holdout set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "grouped = with_docstrings.groupby('nwo')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.6/site-packages/sklearn/model_selection/_split.py:2026: FutureWarning: From version 0.21, test_size will always complement train_size unless both are specified.\n",
      "  FutureWarning)\n"
     ]
    }
   ],
   "source": [
    "# train, valid, test splits\n",
    "train, test = train_test_split(list(grouped), train_size=0.87, shuffle=True, random_state=8081)\n",
    "train, valid = train_test_split(train, train_size=0.82, random_state=8081)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train = pd.concat([d for _, d in train]).reset_index(drop=True)\n",
    "valid = pd.concat([d for _, d in valid]).reset_index(drop=True)\n",
    "test = pd.concat([d for _, d in test]).reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train set num rows 1,008,492\n",
      "valid set num rows 215,445\n",
      "test set num rows 181,272\n",
      "without docstring rows 4,008,718\n"
     ]
    }
   ],
   "source": [
    "print(f'train set num rows {train.shape[0]:,}')\n",
    "print(f'valid set num rows {valid.shape[0]:,}')\n",
    "print(f'test set num rows {test.shape[0]:,}')\n",
    "print(f'without docstring rows {without_docstrings.shape[0]:,}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Preview what the training set looks like.  You can start to see how the data looks, the function tokens and docstring tokens are what will be fed downstream into the models.  The other information is important for diagnostics and bookeeping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "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>nwo</th>\n",
       "      <th>path</th>\n",
       "      <th>function_name</th>\n",
       "      <th>lineno</th>\n",
       "      <th>original_function</th>\n",
       "      <th>function_tokens</th>\n",
       "      <th>docstring_tokens</th>\n",
       "      <th>url</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ajalt/python-sha1</td>\n",
       "      <td>sha1.py</td>\n",
       "      <td>_left_rotate</td>\n",
       "      <td>13</td>\n",
       "      <td>def _left_rotate(n, b):\\n    \"\"\"Left rotate a ...</td>\n",
       "      <td>def _left_rotate n b return n b n 32 b 4294967295</td>\n",
       "      <td>left rotate a 32-bit integer n by b bits .</td>\n",
       "      <td>https://github.com/ajalt/python-sha1/blob/mast...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>ajalt/python-sha1</td>\n",
       "      <td>sha1.py</td>\n",
       "      <td>_process_chunk</td>\n",
       "      <td>18</td>\n",
       "      <td>def _process_chunk(chunk, h0, h1, h2, h3, h4):...</td>\n",
       "      <td>def _process_chunk chunk h0 h1 h2 h3 h4 assert...</td>\n",
       "      <td>process a chunk of data and return the new dig...</td>\n",
       "      <td>https://github.com/ajalt/python-sha1/blob/mast...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>ajalt/python-sha1</td>\n",
       "      <td>sha1.py</td>\n",
       "      <td>sha1</td>\n",
       "      <td>146</td>\n",
       "      <td>def sha1(data):\\n    \"\"\"SHA-1 Hashing Function...</td>\n",
       "      <td>def sha1 data return Sha1Hash update data hexd...</td>\n",
       "      <td>sha-1 hashing function</td>\n",
       "      <td>https://github.com/ajalt/python-sha1/blob/mast...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>ajalt/python-sha1</td>\n",
       "      <td>sha1.py</td>\n",
       "      <td>update</td>\n",
       "      <td>90</td>\n",
       "      <td>def update(self, arg):\\n    \"\"\"Update the curr...</td>\n",
       "      <td>def update self arg if isinstance arg bytes by...</td>\n",
       "      <td>update the current digest .</td>\n",
       "      <td>https://github.com/ajalt/python-sha1/blob/mast...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>ajalt/python-sha1</td>\n",
       "      <td>sha1.py</td>\n",
       "      <td>digest</td>\n",
       "      <td>113</td>\n",
       "      <td>def digest(self):\\n    \"\"\"Produce the final ha...</td>\n",
       "      <td>def digest self return b join struct pack b I ...</td>\n",
       "      <td>produce the final hash value ( big - endian ) ...</td>\n",
       "      <td>https://github.com/ajalt/python-sha1/blob/mast...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                 nwo     path   function_name  lineno  \\\n",
       "0  ajalt/python-sha1  sha1.py    _left_rotate      13   \n",
       "1  ajalt/python-sha1  sha1.py  _process_chunk      18   \n",
       "2  ajalt/python-sha1  sha1.py            sha1     146   \n",
       "3  ajalt/python-sha1  sha1.py          update      90   \n",
       "4  ajalt/python-sha1  sha1.py          digest     113   \n",
       "\n",
       "                                   original_function  \\\n",
       "0  def _left_rotate(n, b):\\n    \"\"\"Left rotate a ...   \n",
       "1  def _process_chunk(chunk, h0, h1, h2, h3, h4):...   \n",
       "2  def sha1(data):\\n    \"\"\"SHA-1 Hashing Function...   \n",
       "3  def update(self, arg):\\n    \"\"\"Update the curr...   \n",
       "4  def digest(self):\\n    \"\"\"Produce the final ha...   \n",
       "\n",
       "                                     function_tokens  \\\n",
       "0  def _left_rotate n b return n b n 32 b 4294967295   \n",
       "1  def _process_chunk chunk h0 h1 h2 h3 h4 assert...   \n",
       "2  def sha1 data return Sha1Hash update data hexd...   \n",
       "3  def update self arg if isinstance arg bytes by...   \n",
       "4  def digest self return b join struct pack b I ...   \n",
       "\n",
       "                                    docstring_tokens  \\\n",
       "0         left rotate a 32-bit integer n by b bits .   \n",
       "1  process a chunk of data and return the new dig...   \n",
       "2                             sha-1 hashing function   \n",
       "3                        update the current digest .   \n",
       "4  produce the final hash value ( big - endian ) ...   \n",
       "\n",
       "                                                 url  \n",
       "0  https://github.com/ajalt/python-sha1/blob/mast...  \n",
       "1  https://github.com/ajalt/python-sha1/blob/mast...  \n",
       "2  https://github.com/ajalt/python-sha1/blob/mast...  \n",
       "3  https://github.com/ajalt/python-sha1/blob/mast...  \n",
       "4  https://github.com/ajalt/python-sha1/blob/mast...  "
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Output each set to train/valid/test.function/docstrings/lineage files\n",
    "Original functions are also written to compressed json files. (Raw functions contain `,`, `\\t`, `\\n`, etc., it is less error-prone using json format)\n",
    "\n",
    "`{train,valid,test}.lineage` are files that contain a reference to the original location where the code was retrieved. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def write_to(df, filename, path='./data/processed_data/'):\n",
    "    \"Helper function to write processed files to disk.\"\n",
    "    out = Path(path)\n",
    "    out.mkdir(exist_ok=True)\n",
    "    df.function_tokens.to_csv(out/'{}.function'.format(filename), index=False)\n",
    "    df.original_function.to_json(out/'{}_original_function.json.gz'.format(filename), orient='values', compression='gzip')\n",
    "    if filename != 'without_docstrings':\n",
    "        df.docstring_tokens.to_csv(out/'{}.docstring'.format(filename), index=False)\n",
    "    df.url.to_csv(out/'{}.lineage'.format(filename), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "if not os.path.exists('data/'):\n",
    "    os.makedirs('data/')\n",
    "# write to output files\n",
    "write_to(train, 'train')\n",
    "write_to(valid, 'valid')\n",
    "write_to(test, 'test')\n",
    "write_to(without_docstrings, 'without_docstrings')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 2.6G\r\n",
      "drwxr-xr-x 2 root root 6.0K May 22 00:59 .\r\n",
      "drwxr-xr-x 9 root root 6.0K May 22 00:53 ..\r\n",
      "-rw-r--r-- 1 root root  13M May 22 00:55 test.docstring\r\n",
      "-rw-r--r-- 1 root root  55M May 22 00:55 test.function\r\n",
      "-rw-r--r-- 1 root root  16M May 22 00:55 test.lineage\r\n",
      "-rw-r--r-- 1 root root  25M May 22 00:55 test_original_function.json.gz\r\n",
      "-rw-r--r-- 1 root root  74M May 22 00:55 train.docstring\r\n",
      "-rw-r--r-- 1 root root 312M May 22 00:53 train.function\r\n",
      "-rw-r--r-- 1 root root  89M May 22 00:55 train.lineage\r\n",
      "-rw-r--r-- 1 root root 140M May 22 00:55 train_original_function.json.gz\r\n",
      "-rw-r--r-- 1 root root  15M May 22 00:55 valid.docstring\r\n",
      "-rw-r--r-- 1 root root  67M May 22 00:55 valid.function\r\n",
      "-rw-r--r-- 1 root root  18M May 22 00:55 valid.lineage\r\n",
      "-rw-r--r-- 1 root root  30M May 22 00:55 valid_original_function.json.gz\r\n",
      "-rw-r--r-- 1 root root 1.1G May 22 00:56 without_docstrings.function\r\n",
      "-rw-r--r-- 1 root root 345M May 22 00:59 without_docstrings.lineage\r\n",
      "-rw-r--r-- 1 root root 357M May 22 00:59 without_docstrings_original_function.json.gz\r\n"
     ]
    }
   ],
   "source": [
    "!ls -lah ./data/processed_data/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The pre-processed data is also hosted on Google Cloud, at the following URLs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/test.docstring\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/test.function\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/test.lineage\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/test_original_function.json.gz\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/train.docstring\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/train.function\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/train.lineage\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/train_original_function.json.gz\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/valid.docstring\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/valid.function\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/valid.lineage\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/valid_original_function.json.gz\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/without_docstrings.function\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/without_docstrings.lineage\n",
      "https://storage.googleapis.com/kubeflow-examples/code_search/data/without_docstrings_original_function.json.gz\n"
     ]
    }
   ],
   "source": [
    "# # cool trick to send shell command results into a python variable in a jupyter notebook!\n",
    "# files = ! ls ./data/processed_data/ | grep -E '*.function$|*.docstring$|*.lineage$|*_original_function.json.gz$'\n",
    "\n",
    "# # print the urls\n",
    "# urls = [f'https://storage.googleapis.com/kubeflow-examples/code_search/data/{f}' for f in files]\n",
    "# for s in urls:\n",
    "#     print(s)"
   ]
  }
 ],
 "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
}