<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Code Search on Kubeflow
This notebook implements an end-to-end Semantic Code Search on top of Kubeflow - given an input query string, get a list of code snippets semantically similar to the query string.
NOTE
Step1:  Install Pip Packages
Step2:  Configure Variables
This involves setting up the Ksonnet application as well as utility environment variables for various CLI steps.
Set the following variables
PROJECT
Step3:  Setup Authorization
In a Kubeflow cluster on GKE, we already have the Google Application Credentials mounted onto each Pod. We can simply point gcloud to activate that service account.
Step4:  Additionally, to interact with the underlying cluster, we configure kubectl.
Step5:  Collectively, these allow us to interact with Google Cloud Services as well as the Kubernetes Cluster directly to submit TFJobs and execute Dataflow pipelines.
Setup Ksonnet Application
We now point the Ksonnet application to the underlying Kubernetes cluster.
Step7:  View Github Files
This is the query that is run as the first step of the Pre-Processing pipeline and is sent through a set of transformations. This is illustrative of the rows being processed in the pipeline we trigger next.
WARNING
Step8:  Define an experiment
This solution consists of multiple jobs and servers that need to share parameters
To facilitate this we use experiments.libsonnet to define sets of parameters
Each set of parameters has a name corresponding to a key in the dictionary defined in experiments.libsonnet
You configure an experiment by defining a set of experiments in experiments.libsonnet
You then set the global parameter experiment to the name of the experiment defining the parameters you want
  to use.
To get started define your experiment
Create a new entry containing a set of values to be used for your experiment
Pick a suitable name for your experiment.
Set the following values
outputDir
Step9:  Submit the Dataflow Job
Step11:  When completed successfully, this should create a dataset in BigQuery named target_dataset. Additionally, it also dumps CSV files into data_dir which contain training samples (pairs of function and docstrings) for our Tensorflow Model. A representative set of results can be viewed using the following query.
Step12:  This pipeline also writes a set of CSV files which contain function and docstring pairs delimited by a comma. Here, we list a subset of them.
Step13:  Prepare Dataset for Training
We will use t2t-datagen to convert the transformed data above into the TFRecord format.
TIP
Step14:  Once this job finishes, the data directory should have a vocabulary file and a list of TFRecords prefixed by the problem name which in our case is github_function_docstring_extended. Here, we list a subset of them.
Step15:  Execute Tensorflow Training
Once, the TFRecords are generated, we will use t2t-trainer to execute the training.
Step16:  This will generate TensorFlow model checkpoints which is illustrated below.
Step17:  Export Tensorflow Model
We now use t2t-exporter to export the TFModel.
Step18:  Once completed, this will generate a TensorFlow SavedModel which we will further use for both online (via TF Serving) and offline inference (via Kubeflow Batch Prediction).
Step19:  Compute Function Embeddings
In this step, we will use the exported model above to compute function embeddings via another Dataflow pipeline. A Python 2 module code_search.dataflow.cli.create_function_embeddings has been provided for this purpose. A list of all possible arguments can be seen below.
Configuration
First, select a Exported Model version from the ${WORKING_DIR}/output/export/Servo as seen above. This should be name of a folder with UNIX Seconds Timestamp like 1533685294. Below, we automatically do that by selecting the folder which represents the latest timestamp.
Step20:  Modify experiments.libsonnet and set modelDir to the directory computed above
Run the Dataflow Job for Function Embeddings
Step22:  When completed successfully, this should create another table in the same BigQuery dataset which contains the function embeddings for each existing data sample available from the previous Dataflow Job. Additionally, it also dumps a CSV file containing metadata for each of the function and its embeddings. A representative query result is shown below.
Step23:  The pipeline also generates a set of CSV files which will be useful to generate the search index.
Step24:  Create Search Index
We now create the Search Index from the computed embeddings. This facilitates k-Nearest Neighbor search to for semantically similar results.
Step25:  Using the CSV files generated from the previous step, this creates an index using NMSLib. A unified CSV file containing all the code examples for a human-readable reverse lookup during the query, is also created.
Step26:  Deploy the Web App
The included web app provides a simple way for users to submit queries
The web app includes two pieces
A Flask app that serves a simple UI for sending queries
The flask app also uses nmslib to provide fast lookups
A TF Serving instance to compute the embeddings for search queries
The ksonnet components for the web app are in a separate ksonnet application ks-web-app
A second web app is used so that we can optionally use ArgoCD to keep the serving components up to date.
Deploy an Inference Server
We've seen offline inference during the computation of embeddings. For online inference, we deploy the exported Tensorflow model above using Tensorflow Serving.
You need to set the parameter modelBasePath to the GCS directory where the model was exported
This will be a directory produced by the export model step
e.g gs
Step27:  Deploy Search UI
We finally deploy the Search UI which allows the user to input arbitrary strings and see a list of results corresponding to semantically similar Python functions. This internally uses the inference server we just deployed.
We need to configure the index server to use the lookup file and CSV produced by the create search index step above
The values will be the values of the parameters lookupFile and indexFile that you set in experiments.libsonnet
<ASSISTANT_TASK:>
Python Code:
%%bash
echo "Pip Version Info: " && python2 --version && python2 -m pip --version && echo
echo "Google Cloud SDK Info: " && gcloud --version && echo
echo "Ksonnet Version Info: " && ks version && echo
echo "Kubectl Version Info: " && kubectl version
Explanation: Code Search on Kubeflow
This notebook implements an end-to-end Semantic Code Search on top of Kubeflow - given an input query string, get a list of code snippets semantically similar to the query string.
NOTE: If you haven't already, see kubeflow/examples/code_search for instructions on how to get this notebook,.
Install dependencies
Let us install all the Python dependencies. Note that everything must be done with Python 2. This will take a while the first time.
Verify Version Information
End of explanation
! python2 -m pip install -U pip
# Code Search dependencies
! python2 -m pip install --user https://github.com/kubeflow/batch-predict/tarball/master
! python2 -m pip install --user -r src/requirements.txt
# BigQuery Cell Dependencies
! python2 -m pip install --user pandas-gbq
# NOTE: The RuntimeWarnings (if any) are harmless. See ContinuumIO/anaconda-issues#6678.
from pandas.io import gbq
Explanation: Install Pip Packages
End of explanation
import getpass
import subprocess
# Configuration Variables. Modify as desired.
PROJECT = subprocess.check_output(["gcloud", "config", "get-value", "project"]).strip()
# Dataflow Related Variables.
TARGET_DATASET = 'code_search'
WORKING_DIR = 'gs://{0}_code_search/workingDir'.format(PROJECT)
KS_ENV=getpass.getuser()
# DO NOT MODIFY. These are environment variables to be used in a bash shell.
%env PROJECT $PROJECT
%env TARGET_DATASET $TARGET_DATASET
%env WORKING_DIR $WORKING_DIR
Explanation: Configure Variables
This involves setting up the Ksonnet application as well as utility environment variables for various CLI steps.
Set the following variables
PROJECT: Set this to the GCP project you want to use
If gcloud has a project set this will be used by default
To use a different project or if gcloud doesn't have a project set you will need to configure one explicitly
WORKING_DIR: Override this if you don't want to use the default configured below     
KS_ENV: Set this to the name of the ksonnet environment you want to create
End of explanation
%%bash
# Activate Service Account provided by Kubeflow.
gcloud auth activate-service-account --key-file=${GOOGLE_APPLICATION_CREDENTIALS}
Explanation: Setup Authorization
In a Kubeflow cluster on GKE, we already have the Google Application Credentials mounted onto each Pod. We can simply point gcloud to activate that service account.
End of explanation
%%bash
kubectl config set-cluster kubeflow --server=https://kubernetes.default --certificate-authority=/var/run/secrets/kubernetes.io/serviceaccount/ca.crt
kubectl config set-credentials jupyter --token "$(cat /var/run/secrets/kubernetes.io/serviceaccount/token)"
kubectl config set-context kubeflow --cluster kubeflow --user jupyter
kubectl config use-context kubeflow
Explanation: Additionally, to interact with the underlying cluster, we configure kubectl.
End of explanation
%%bash
cd kubeflow
# Update Ksonnet to point to the Kubernetes Cluster
ks env add code-search
# Update Ksonnet to use the namespace where kubeflow is deployed. By default it's 'kubeflow'
ks env set code-search --namespace=kubeflow
# Update the Working Directory of the application
sed -i'' "s,gs://example/prefix,${WORKING_DIR}," components/params.libsonnet
# FIXME(sanyamkapoor): This command completely replaces previous configurations.
# Hence, using string replacement in file.
# ks param set t2t-code-search workingDir ${WORKING_DIR}
Explanation: Collectively, these allow us to interact with Google Cloud Services as well as the Kubernetes Cluster directly to submit TFJobs and execute Dataflow pipelines.
Setup Ksonnet Application
We now point the Ksonnet application to the underlying Kubernetes cluster.
End of explanation
query = 
  SELECT
    MAX(CONCAT(f.repo_name, ' ', f.path)) AS repo_path,
    c.content
  FROM
    `bigquery-public-data.github_repos.files` AS f
  JOIN
    `bigquery-public-data.github_repos.contents` AS c
  ON
    f.id = c.id
  JOIN (
      --this part of the query makes sure repo is watched at least twice since 2017
    SELECT
      repo
    FROM (
      SELECT
        repo.name AS repo
      FROM
        `githubarchive.year.2017`
      WHERE
        type="WatchEvent"
      UNION ALL
      SELECT
        repo.name AS repo
      FROM
        `githubarchive.month.2018*`
      WHERE
        type="WatchEvent" )
    GROUP BY
      1
    HAVING
      COUNT(*) >= 2 ) AS r
  ON
    f.repo_name = r.repo
  WHERE
    f.path LIKE '%.py' AND --with python extension
    c.size < 15000 AND --get rid of ridiculously long files
    REGEXP_CONTAINS(c.content, r'def ') --contains function definition
  GROUP BY
    c.content
  LIMIT
    10
gbq.read_gbq(query, dialect='standard', project_id=PROJECT)
Explanation: View Github Files
This is the query that is run as the first step of the Pre-Processing pipeline and is sent through a set of transformations. This is illustrative of the rows being processed in the pipeline we trigger next.
WARNING: The table is large and the query can take a few minutes to complete.
End of explanation
%%bash
bq mk ${PROJECT}:${TARGET_DATASET}
Explanation: Define an experiment
This solution consists of multiple jobs and servers that need to share parameters
To facilitate this we use experiments.libsonnet to define sets of parameters
Each set of parameters has a name corresponding to a key in the dictionary defined in experiments.libsonnet
You configure an experiment by defining a set of experiments in experiments.libsonnet
You then set the global parameter experiment to the name of the experiment defining the parameters you want
  to use.
To get started define your experiment
Create a new entry containing a set of values to be used for your experiment
Pick a suitable name for your experiment.
Set the following values
outputDir: The GCS directory where the output should be written
train_steps: Numer oftraining steps
eval_steps: Number of steps to be used for eval
hparams_set: The set of hyperparameters to use; see some suggestions here.
transformer_tiny can be used to train a very small model suitable for ensuring the code works.
project: Set this to the GCP project you can use
modelDir: 
After training your model set this to a GCS directory containing the export model
e.g gs://code-search-demo/models/20181107-dist-sync-gpu/export/1541712907/
problem: set this to "kf_github_function_docstring",
model: set this  "kf_similarity_transformer",
lookupFile: set this to the GCS location of the CSV produced by the job to create the nmslib index of the embeddings for all GitHub data
indexFile: set this to the GCS location of the nmslib index for all the data in GitHub
Configure your ksonnet environment to use your experiment
Open kubeflow/environments/${ENVIRONMENT}/globals.libsonnet
Define the following global parameters
experiment: Name of your experiment; should correspond to a key defined in experiments.libsonnet
project: Set this to the GCP project you can use
dataDir: The data directory to be used by T2T  
workingDir: Working directory
Here's an example of what the contents should look like
workingDir: "gs://code-search-demo/20181104",
  dataDir: "gs://code-search-demo/20181104/data",
  project: "code-search-demo",
  experiment: "demo-trainer-11-07-dist-sync-gpu",
Pre-Processing Github Files
In this step, we use Google Cloud Dataflow to preprocess the data.
We use a K8s Job to run a python program code_search.dataflow.cli.preprocess_github_dataset that submits the Dataflow job
Once the job has been created it can be monitored using the Dataflow console
The parameter target_dataset specifies a BigQuery dataset to write the data to
Create the BigQuery dataset
End of explanation
%%bash
cd kubeflow
ks param set --env=code-search submit-preprocess-job targetDataset ${TARGET_DATASET}
ks apply code-search -c submit-preprocess-job
Explanation: Submit the Dataflow Job
End of explanation
query = 
  SELECT * 
  FROM 
    {}.token_pairs
  LIMIT
    10
.format(TARGET_DATASET)
gbq.read_gbq(query, dialect='standard', project_id=PROJECT)
Explanation: When completed successfully, this should create a dataset in BigQuery named target_dataset. Additionally, it also dumps CSV files into data_dir which contain training samples (pairs of function and docstrings) for our Tensorflow Model. A representative set of results can be viewed using the following query.
End of explanation
%%bash
LIMIT=10
gsutil ls ${WORKING_DIR}/data/*.csv | head -n ${LIMIT}
Explanation: This pipeline also writes a set of CSV files which contain function and docstring pairs delimited by a comma. Here, we list a subset of them.
End of explanation
%%bash
cd kubeflow
ks show code-search -c t2t-code-search-datagen
%%bash
cd kubeflow
ks apply code-search -c t2t-code-search-datagen
Explanation: Prepare Dataset for Training
We will use t2t-datagen to convert the transformed data above into the TFRecord format.
TIP: Use ks show to view the Resource Spec submitted.
End of explanation
%%bash
LIMIT=10
gsutil ls ${WORKING_DIR}/data/vocab*
gsutil ls ${WORKING_DIR}/data/*train* | head -n ${LIMIT}
Explanation: Once this job finishes, the data directory should have a vocabulary file and a list of TFRecords prefixed by the problem name which in our case is github_function_docstring_extended. Here, we list a subset of them.
End of explanation
%%bash
cd kubeflow
ks show code-search -c t2t-code-search-trainer
%%bash
cd kubeflow
ks apply code-search -c t2t-code-search-trainer
Explanation: Execute Tensorflow Training
Once, the TFRecords are generated, we will use t2t-trainer to execute the training.
End of explanation
%%bash
gsutil ls ${WORKING_DIR}/output/*ckpt*
Explanation: This will generate TensorFlow model checkpoints which is illustrated below.
End of explanation
%%bash
cd kubeflow
ks show code-search -c t2t-code-search-exporter
%%bash
cd kubeflow
ks apply code-search -c t2t-code-search-exporter
Explanation: Export Tensorflow Model
We now use t2t-exporter to export the TFModel.
End of explanation
%%bash
gsutil ls ${WORKING_DIR}/output/export/Servo
Explanation: Once completed, this will generate a TensorFlow SavedModel which we will further use for both online (via TF Serving) and offline inference (via Kubeflow Batch Prediction).
End of explanation
%%bash --out EXPORT_DIR_LS
gsutil ls ${WORKING_DIR}/output/export/Servo | grep -oE "([0-9]+)/$"
# WARNING: This routine will fail if no export has been completed successfully.
MODEL_VERSION = max([int(ts[:-1]) for ts in EXPORT_DIR_LS.split('\n') if ts])
# DO NOT MODIFY. These are environment variables to be used in a bash shell.
%env MODEL_VERSION $MODEL_VERSION
Explanation: Compute Function Embeddings
In this step, we will use the exported model above to compute function embeddings via another Dataflow pipeline. A Python 2 module code_search.dataflow.cli.create_function_embeddings has been provided for this purpose. A list of all possible arguments can be seen below.
Configuration
First, select a Exported Model version from the ${WORKING_DIR}/output/export/Servo as seen above. This should be name of a folder with UNIX Seconds Timestamp like 1533685294. Below, we automatically do that by selecting the folder which represents the latest timestamp.
End of explanation
%%bash
cd kubeflow
ks apply code-search -c submit-code-embeddings-job        
Explanation: Modify experiments.libsonnet and set modelDir to the directory computed above
Run the Dataflow Job for Function Embeddings
End of explanation
query = 
  SELECT * 
  FROM 
    {}.function_embeddings
  LIMIT
    10
.format(TARGET_DATASET)
gbq.read_gbq(query, dialect='standard', project_id=PROJECT)
Explanation: When completed successfully, this should create another table in the same BigQuery dataset which contains the function embeddings for each existing data sample available from the previous Dataflow Job. Additionally, it also dumps a CSV file containing metadata for each of the function and its embeddings. A representative query result is shown below.
End of explanation
%%bash
LIMIT=10
gsutil ls ${WORKING_DIR}/data/*index*.csv | head -n ${LIMIT}
Explanation: The pipeline also generates a set of CSV files which will be useful to generate the search index.
End of explanation
%%bash
cd kubeflow
ks show code-search -c search-index-creator
%%bash
cd kubeflow
ks apply code-search -c search-index-creator
Explanation: Create Search Index
We now create the Search Index from the computed embeddings. This facilitates k-Nearest Neighbor search to for semantically similar results.
End of explanation
%%bash
gsutil ls ${WORKING_DIR}/code_search_index*
Explanation: Using the CSV files generated from the previous step, this creates an index using NMSLib. A unified CSV file containing all the code examples for a human-readable reverse lookup during the query, is also created.
End of explanation
%%bash
cd ks-web-app
ks param set --env=code-search modelBasePath ${MODEL_BASE_PATH}
ks show code-search -c query-embed-server
%%bash
cd ks-web-app
ks apply code-search -c query-embed-server
Explanation: Deploy the Web App
The included web app provides a simple way for users to submit queries
The web app includes two pieces
A Flask app that serves a simple UI for sending queries
The flask app also uses nmslib to provide fast lookups
A TF Serving instance to compute the embeddings for search queries
The ksonnet components for the web app are in a separate ksonnet application ks-web-app
A second web app is used so that we can optionally use ArgoCD to keep the serving components up to date.
Deploy an Inference Server
We've seen offline inference during the computation of embeddings. For online inference, we deploy the exported Tensorflow model above using Tensorflow Serving.
You need to set the parameter modelBasePath to the GCS directory where the model was exported
This will be a directory produced by the export model step
e.g gs://code-search-demo/models/20181107-dist-sync-gpu/export/
Here are sample contents
  ```
  gs://code-search-demo/models/20181107-dist-sync-gpu/export/:
  gs://code-search-demo/models/20181107-dist-sync-gpu/export/
gs://code-search-demo/models/20181107-dist-sync-gpu/export/1541712907/:
  gs://code-search-demo/models/20181107-dist-sync-gpu/export/1541712907/
  gs://code-search-demo/models/20181107-dist-sync-gpu/export/1541712907/saved_model.pbtxt
gs://code-search-demo/models/20181107-dist-sync-gpu/export/1541712907/variables/:
  gs://code-search-demo/models/20181107-dist-sync-gpu/export/1541712907/variables/
  gs://code-search-demo/models/20181107-dist-sync-gpu/export/1541712907/variables/variables.data-00000-of-00001
  gs://code-search-demo/models/20181107-dist-sync-gpu/export/1541712907/variables/variables.index
```
TFServing expects the modelBasePath to consist of numeric subdirectories corresponding to different versions of the model.
Each subdirectory will contain the saved model in protocol buffer along with the weights.
End of explanation
%%bash
cd ks-web-app
ks param set --env=code-search search-index-server lookupFile ${LOOKUP_FILE}
ks param set --env=code-search search-index-server indexFile ${INDEX_FILE}
ks show code-search -c search-index-server
%%bash
cd kubeflow
ks apply code-search -c search-index-server
Explanation: Deploy Search UI
We finally deploy the Search UI which allows the user to input arbitrary strings and see a list of results corresponding to semantically similar Python functions. This internally uses the inference server we just deployed.
We need to configure the index server to use the lookup file and CSV produced by the create search index step above
The values will be the values of the parameters lookupFile and indexFile that you set in experiments.libsonnet
End of explanation
<END_TASK>