<jupyter_start><jupyter_text>CS Local search  Outline1. Export model2. Run large-batch inference of code embeddings, write to CSV.3. Load code embeddings into nmslib search index.4. Using tf.Eager, perform inference for a query in notebook.5. Search that query embedding against the search index, listing hits in notebook.  Use tf.Eager This will be far slower in part because for now I'm not able to use batch sizes > 1 (for what ever reason). But - I know this will for sure work. And it will probably take less than like 5min to compute embeddings for like 10k examples.... 200 takes about 10-20s... ok so more like 15min.... an index of 3k examples in 5min... so maybe could build an index using a set of things that are or aren't known to go together, e.g. two very distinct codebases instead of a random sampling of all... maybe just use the t2t library?...<jupyter_code>import numpy as np

import functools

import tensorflow as tf

from tensor2tensor.data_generators import problem
from tensor2tensor.layers import common_layers
from tensor2tensor.models import transformer
from tensor2tensor.utils import registry
from tensor2tensor.utils import t2t_model
from tensor2tensor.models.transformer import Transformer

from tensor2tensor.models.transformer import transformer_base

from tensorflow.contrib.eager.python import tfe
tfe.enable_eager_execution()
Modes = tf.estimator.ModeKeys


from tk.models import similarity_transformer
from tk.data_generators import function_docstring

import numpy as np; np.random.seed(0)
import matplotlib.pyplot as plt
import seaborn as sns; sns.set()
%matplotlib inline


mp_constrained_embedding = function_docstring.GithubConstrainedEmbedding()

data_dir = "/mnt/nfs-east1-d/data"

hparams = similarity_transformer.similarity_transformer_tiny()
hparams.data_dir = data_dir

p_hparams = mp_constrained_embedding.get_hparams(hparams)

model = similarity_transformer.ConstrainedEmbeddingTransformer(
    hparams, tf.estimator.ModeKeys.PREDICT, p_hparams
)

# Get the encoders from the problem
encoders = mp_constrained_embedding.feature_encoders(data_dir)

# Setup helper functions for encoding and decoding
def encode(input_str, output_str=None):
  """Input str to features dict, ready for inference"""
  inputs = encoders["inputs"].encode(input_str) + [1]  # add EOS id
  batch_inputs = tf.reshape(inputs, [1, -1, 1])  # Make it 3D.
  return {"inputs": batch_inputs}

def decode(integers):
  """List of ints to str
  
  For decoding an integer encoding to its string representation,
  not for decoding an embedding vector into the same.
  """
  integers = list(np.squeeze(integers))
  if 1 in integers:
    integers = integers[:integers.index(1)]
  return encoders["inputs"].decode(np.squeeze(integers))

batch_size = 1
train_dataset = mp_constrained_embedding.dataset(Modes.PREDICT, data_dir)
train_dataset = train_dataset.repeat(None).batch(batch_size)

iterator = tfe.Iterator(train_dataset)


def embed_many(dataset_iterator, model, ckpt_path, num=100):
    embeddings = []
    num_examples = 100

    with tfe.restore_variables_on_create(ckpt_path):

      for i in range(0, num):
        example = dataset_iterator.next()

        #doc_emb, _ = model({"inputs": example["docstring"]})
        code_emb, _ = model({"inputs": example["code"]})

        #embeddings.append([doc_emb, code_emb, decode(example["docstring"]), decode(example["code"])])
        embeddings.append([code_emb, decode(example["code"]), decode(example["docstring"])])
        
        if i % 100 == 0:
          print("Processing step %s" % i)
    
    print("Finished embedding %s" % num)

    return embeddings

ckpt_path = "gs://kubeflow-rl-checkpoints/comparisons/cs-v7-lvslicenet-ts1-tcs1-tm1-exl1-ntm1/cs-v7-lvslicenet-ts1-tcs1-tm1-exl1-ntm1-j1026-1730-f4d7/output/model.ckpt-215387"
import datetime
print datetime.datetime.now()
embeddings = embed_many(iterator, model, ckpt_path, 10000)
print datetime.datetime.now()

numpy_embedding_vectors = np.asarray([thing[0].numpy()[0] for thing in embeddings])

embedding_data_path = "/mnt/nfs-east1-d/tmp/embeddings.csv"

np.savetxt(embedding_data_path, np.asarray([thing[0].numpy()[0] for thing in embeddings]), delimiter=",")<jupyter_output><empty_output><jupyter_text>Build an index<jupyter_code>import nmslib
import pprint

del index

def build_index(index_save_path, index_data):
  index = nmslib.init(method='hnsw', space='cosinesimil')
  index.addDataPointBatch(index_data)
  index.createIndex({'post': 2}, print_progress=True)
  index.saveIndex(index_save_path)
  return index

def embed_query(model, ckpt_path, query):
  with tfe.restore_variables_on_create(ckpt_path):
    return model(encode(query))[0].numpy()[0]

ckpt_path = "gs://kubeflow-rl-checkpoints/comparisons/cs-v7-lvslicenet-ts1-tcs1-tm1-exl1-ntm1/cs-v7-lvslicenet-ts1-tcs1-tm1-exl1-ntm1-j1026-1730-f4d7/output/model.ckpt-215387"
query_fn = functools.partial(embed_query, model, ckpt_path)

local_index_save_path = "/mnt/nfs-east1-d/tmp/index-004"

k = 10000

def results_for_query(query):
  query_embedding = query_fn(query)

  idxs, dists = index.knnQuery(query_embedding, k=k)

  hits = []
    
  for (i, d) in zip(idxs, dists):
    hits.append((embeddings[i][2], embeddings[i][1], d))
  return hits

def print_hits(hit_subset):
  for hit in hit_subset:
    print(hit[0])
    print(hit[1])
    print(hit[2])
    print("\n")
  print("--")


index = build_index(local_index_save_path, numpy_embedding_vectors)


hits = results_for_query("fetch the pagination marker field from flask.request.args")


# Most similar
print_hits(hits[0:5])

# Least similar
print_hits(hits[-5:])

# Nope.

for etuple in embeddings:
  if "lask" in etuple[2]:
    print etuple[2]

# That exact doc string is present in the dataset.


for etuple in hits:
  if "lask" in etuple[0]:
    print etuple
    

np.median([hit[2] for hit in hits])

[hit[2] for hit in hits]

# Not a highly ranked hit<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10819451455493859, 'fraction_numerical': 0.023725391216557295, 'mean_word_length': 4.8046875, 'pattern_counts': {'":': 3, '<': 9, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '166566', 'license': 'Apache-2.0', 'n_tokens_mistral': 2135, 'n_tokens_neox': 2025, 'n_words': 517, 'path': 'docs/cs_local_search.ipynb', 'repo_name': 'cwbeitel/tk', 'size': '5943'}