deployment/tensorflow/nearest-neighbors-tf-glove.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Nearest Neighbor Search with Tensorflow and Glove Embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example features:\n",
    "- Using Glove embeddings to turn IMDB words into vectors\n",
    "- Using a simple TF model to compute representations of IMDB reviews by averaging word vectors\n",
    "- Using the annoy library to build an index of word vectors\n",
    "- Search similar reviews via embeddings and the annoy library\n",
    "\n",
    "\n",
    "Versions:\n",
    "- Python: 2.7\n",
    "- TF: {1.14.0, 1.15.2}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import six\n",
    "\n",
    "import json\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import annoy\n",
    "\n",
    "from verta import Client\n",
    "from verta.utils import ModelAPI, TFSavedModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "REPRESENTATION_LENGTH = 25\n",
    "MAX_INPUT_LENGTH = 50\n",
    "\n",
    "HOST = 'app.verta.ai'\n",
    "DATA_DIR = ''\n",
    "DATA_FILE = DATA_DIR + 'imdb_master.csv'\n",
    "EMBEDDING_FILE = DATA_DIR + 'glove.twitter.27B/glove.twitter.27B.25d.txt'\n",
    "SAVED_MODEL_DIR = 'saved-model'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the Glove embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the glove embeddings\n",
    "embeddings_index = dict()\n",
    "with open(EMBEDDING_FILE) as f:\n",
    "    for line in f:\n",
    "        values = line.split()\n",
    "        word = values[0]\n",
    "        coefs = np.asarray(values[1:], dtype='float32')\n",
    "        embeddings_index[word] = coefs\n",
    "embeddings_index['UNK'] = [0.0] * REPRESENTATION_LENGTH\n",
    "print('Loaded %s word vectors.' % len(embeddings_index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_matrix = np.zeros((len(embeddings_index), REPRESENTATION_LENGTH))\n",
    "word_to_index = {}\n",
    "ctr = 0\n",
    "UNK_INDEX = -1\n",
    "for word, embedding in embeddings_index.items():\n",
    "    if word == 'UNK':\n",
    "        UNK_INDEX = ctr\n",
    "    word_to_index[word] = ctr\n",
    "    embedding_matrix[ctr, :] = embedding\n",
    "    ctr += 1\n",
    "VOCABULARY_SIZE = embedding_matrix.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "reviews = pd.read_csv(DATA_FILE)['review'].values.tolist()\n",
    "\n",
    "input_data = reviews[:1000]\n",
    "input_data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build Simple TF model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil; shutil.rmtree(SAVED_MODEL_DIR, ignore_errors=True)\n",
    "\n",
    "batch_indexes = tf.placeholder(tf.int32, shape=[50], name=\"input\")\n",
    "\n",
    "tf_embedding = tf.Variable(\n",
    "    tf.constant(0.0, shape=[VOCABULARY_SIZE, REPRESENTATION_LENGTH]),\n",
    "    trainable=False,\n",
    "    name=\"Embedding\",\n",
    ")\n",
    "\n",
    "tf_embedding_placeholder = tf.placeholder(\n",
    "    tf.float32,\n",
    "    [VOCABULARY_SIZE, REPRESENTATION_LENGTH]\n",
    ")\n",
    "tf_embedding_init = tf_embedding.assign(tf_embedding_placeholder)\n",
    "\n",
    "embedding_list = tf.nn.embedding_lookup(\n",
    "    params=tf_embedding, ids=batch_indexes)\n",
    "\n",
    "concatenated_embedding = tf.concat(embedding_list, -1)\n",
    "\n",
    "embedding = tf.reduce_mean(concatenated_embedding, axis=0, name=\"output\")\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "\n",
    "    _ = sess.run(\n",
    "        tf_embedding_init,\n",
    "        feed_dict={\n",
    "            tf_embedding_placeholder: embedding_matrix\n",
    "        }\n",
    "    )\n",
    "\n",
    "    # save the model for use later\n",
    "    tf.saved_model.simple_save(\n",
    "        sess,\n",
    "        SAVED_MODEL_DIR,\n",
    "        {'batch_indexes': batch_indexes},\n",
    "        {'embedding': embedding}\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from annoy import AnnoyIndex\n",
    "\n",
    "model = TFSavedModel(SAVED_MODEL_DIR)\n",
    "t = AnnoyIndex(REPRESENTATION_LENGTH, 'angular')  # Length of item vector that will be indexed\n",
    "for i in range(len(reviews)):\n",
    "    review = reviews[i]\n",
    "    words = review.split()[:MAX_INPUT_LENGTH]\n",
    "    batch_indexes = [word_to_index.get(w.lower(), word_to_index['UNK']) for w in words]\n",
    "    batch_indexes += [UNK_INDEX] * (MAX_INPUT_LENGTH - len(batch_indexes))\n",
    "        \n",
    "    # calculate embedding with TF\n",
    "    embedding = model.predict(batch_indexes=batch_indexes)['embedding']\n",
    "    t.add_item(i, embedding)\n",
    "\n",
    "t.build(10) # 10 trees\n",
    "t.save('reviews.ann')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instantiate Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(HOST)\n",
    "client.set_project(\"TF\")\n",
    "client.set_experiment(\"SavedModel\")\n",
    "run = client.set_experiment_run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Log Artifacts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "run.log_artifact(\"saved_model\", SAVED_MODEL_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"word_to_index.json\", 'w') as f:\n",
    "    json.dump(word_to_index, f)\n",
    "run.log_artifact(\"word_to_index\", \"word_to_index.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "run.log_artifact(\"reviews_index\", \"reviews.ann\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build a NearestNeighbor Search class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "class TextNNSearch(object):\n",
    "    def __init__(self, artifacts):\n",
    "        self.session = tf.Session()\n",
    "        tf.compat.v1.saved_model.load(self.session, ['serve'], artifacts['saved_model'])\n",
    "        \n",
    "        with open(artifacts['word_to_index'], 'r') as f:\n",
    "            self.word_to_index = json.load(f)\n",
    "        \n",
    "        self.index = annoy.AnnoyIndex(REPRESENTATION_LENGTH, \"angular\")\n",
    "        self.index.load(artifacts['reviews_index'])\n",
    "\n",
    "    def predict(self, input_strs):\n",
    "        predictions = []\n",
    "        for input_str in input_strs:\n",
    "            words = input_str.split()[:MAX_INPUT_LENGTH]\n",
    "            batch_indexes = [self.word_to_index.get(w.lower(), self.word_to_index['UNK']) for w in words]\n",
    "            batch_indexes += [UNK_INDEX] * (MAX_INPUT_LENGTH - len(batch_indexes))\n",
    "        \n",
    "            # calculate embedding with TF\n",
    "            embedding = self.session.run(\"output:0\", {\"input:0\": batch_indexes})\n",
    "            \n",
    "            # find embedding vectors of ten nearest neighbors\n",
    "            predictions.append({\n",
    "                input_str: self.index.get_nns_by_vector(embedding, 10)\n",
    "            })\n",
    "        return predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run some simple tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "artifacts = run.fetch_artifacts([\"saved_model\", \"word_to_index\", \"reviews_index\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = TextNNSearch(artifacts)\n",
    "\n",
    "prediction = model.predict([\"omg I love this film\"])\n",
    "similar_reviews = [reviews[i] for i in prediction[0].values()[0]]\n",
    "print(similar_reviews[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deploy Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "run.log_model(\n",
    "    TextNNSearch,\n",
    "    custom_modules=[],\n",
    "    model_api=ModelAPI(input_data, model.predict(input_data)),\n",
    "    artifacts=[\"saved_model\", \"word_to_index\", \"reviews_index\"],\n",
    ")\n",
    "run.log_requirements([\"tensorflow\", \"annoy==1.15.2\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deploy the model and make predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove reviews with weird bytes\n",
    "bad_reviews = []\n",
    "for i, review in enumerate(input_data):\n",
    "    try:\n",
    "        unicode(review, 'utf-8')\n",
    "    except UnicodeDecodeError:\n",
    "        bad_reviews.append(i)\n",
    "    else:\n",
    "        pass\n",
    "for i in sorted(bad_reviews, reverse=True):\n",
    "    del input_data[i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from verta.deployment import DeployedModel\n",
    "\n",
    "embeddings = DeployedModel(HOST, run.id).predict(input_data[:1000], compress=True)\n",
    "\n",
    "embeddings[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}