hub2ann_tf2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "hub2ann-tf2.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "ACbjNjyO4f_8"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ACbjNjyO4f_8",
        "colab_type": "text"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Hub Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MCM50vaM4jiK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Copyright 2018 The TensorFlow Hub Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9qOVy-_vmuUP",
        "colab_type": "text"
      },
      "source": [
        "# Approximate Embeddings Similarity Matching\n",
        "\n",
        "This tutorial illustrates how to gnerate embeddings from a [TensorFlow Hub](https://www.tensorflow.org/hub) (TF-Hb) module given input data, and build an approximate nearest neighbours (ANN) index using the extracted embeddings. The index can then be used for real-time similarity matching an retreival.\n",
        "\n",
        "The steps of this tutorial are:\n",
        "1. Download sample data.\n",
        "2. Generate embeddings for the data using a TF-Hub module\n",
        "3. Build an ANN index for the embeddings\n",
        "4. Use the index for similarity matching\n",
        "\n",
        "We use [Apache Beam](https://beam.apache.org/documentation/programming-guide/) with to generate the embeddings from the TF-Hub module. We also use Spotify's [ANNOY](https://github.com/spotify/annoy) library to build the approximate nearest neighbours index.\n",
        "\n",
        "This tutorial uses with **TensorFlow 2.0**, and works only with tf2 TF-Hub modules (**SavedModel 2.0**)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q0jr0QK9qO5P",
        "colab_type": "text"
      },
      "source": [
        "## Getting Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "whMRj9qeqed4",
        "colab_type": "text"
      },
      "source": [
        "Install the required libraries. Make sure to **restart the runtime** after installtion is completed."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qmXkLPoaqS--",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip3 install -U tensorflow\n",
        "!pip3 install apache_beam[gcp]\n",
        "!pip3 install annoy"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A-vBZiCCqld0",
        "colab_type": "text"
      },
      "source": [
        "Import the required libraries"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6NTYbdWcseuK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import sys\n",
        "import multiprocessing\n",
        "import pickle\n",
        "from datetime import datetime\n",
        "import numpy as np\n",
        "import apache_beam as beam\n",
        "from apache_beam.transforms import util\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "import annoy"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tx0SZa6-7b-f",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print('TF version: {}'.format(tf.__version__))\n",
        "print('TF-Hub version: {}'.format(hub.__version__))\n",
        "print('Apache Beam version: {}'.format(beam.__version__))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P6Imq876rLWx",
        "colab_type": "text"
      },
      "source": [
        "## 1. Download Sample Data\n",
        "\n",
        "[A Million News Headlines](https://dataverse.harvard.edu/dataset.xhtml?persistentId=doi:10.7910/DVN/SYBGZL#) contains data of news headlines published over a period of 15 years. Sourced from the reputable Australian news source Australian Broadcasting Corp. (ABC). This this news dataset as a summarised historical record of noteworthy events in the globe from early-2003 to end-2017 with a more granular focus on Australia. \n",
        "\n",
        "This includes the entire corpus of articles published by the ABC website in the given time range. With a volume of 200 articles per day and a good focus on international news, events of significance has been captured here. Digging into the keywords, one can see all the important episodes shaping the last decade and how they evolved over time. Ex: financial crisis, iraq war, multiple US elections, ecological disasters, terrorism, famous people, Australian crimes etc.\n",
        "\n",
        "**Format**: Tab-separated two-column data: 1) publication date and 2) headline text. We are only interested in the headline text.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "OpF57n8e5C9D",
        "colab": {}
      },
      "source": [
        "!wget https://dataverse.harvard.edu/api/access/datafile/3450625?format=tab&gbrecs=true\n",
        "!mv 3450625?format=tab raw.tsv\n",
        "!wc -l raw.tsv\n",
        "!head raw.tsv"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Reeoc9z0zTxJ",
        "colab_type": "text"
      },
      "source": [
        "For simplicity, we only keep the headline text and remove the publication date"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "INPWa4upv_yJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!rm -r corpus\n",
        "!mkdir corpus\n",
        "\n",
        "with open('corpus/text.txt', 'w') as out_file:\n",
        "  with open('raw.tsv', 'r') as in_file:\n",
        "    for line in in_file:\n",
        "      headline = line.split('\\t')[1].strip().strip('\"')\n",
        "      out_file.write(headline+\"\\n\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5-oedX40z6o2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!tail corpus/text.txt"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2AngMtH50jNb",
        "colab_type": "text"
      },
      "source": [
        "## 2. Generate Embeddings for the Data.\n",
        "\n",
        "In this tutorial, we use the [Neural Network Language Model (NNLM)](https://tfhub.dev/google/tf2-preview/nnlm-en-dim50/1) to generate emebeddings for the headlines data. The sentence embeddings can then be easily used to compute sentence level meaning similarity. We run the embeddings generation process using Apache Beam."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F_DvXnDB1pEX",
        "colab_type": "text"
      },
      "source": [
        "### Embeddings extraction pipeline"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yL7OEY1E0A35",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "embed_fn = None\n",
        "\n",
        "def generate_embeddings(text, module_url, random_projection_matrix=None):\n",
        "  import tensorflow_hub as hub\n",
        "  global embed_fn\n",
        "  if embed_fn is None:\n",
        "    embed_fn = hub.load(module_url)\n",
        "  embedding = embed_fn(text).numpy()\n",
        "  if random_projection_matrix is not None:\n",
        "    embedding = embedding.dot(random_projection_matrix)\n",
        "  return text, embedding\n",
        "\n",
        "def to_tf_example(entries):\n",
        "  examples = []\n",
        "\n",
        "  text_list, embedding_list = entries\n",
        "  for i in range(len(text_list)):\n",
        "    text = text_list[i]\n",
        "    embedding = embedding_list[i]\n",
        "\n",
        "    features = {\n",
        "        'text': tf.train.Feature(\n",
        "            bytes_list=tf.train.BytesList(value=[text.encode('utf-8')])),\n",
        "        'embedding': tf.train.Feature(\n",
        "            float_list=tf.train.FloatList(value=embedding.tolist()))\n",
        "    }\n",
        "  \n",
        "    example = tf.train.Example(\n",
        "        features=tf.train.Features(\n",
        "            feature=features)).SerializeToString(deterministic=True)\n",
        "  \n",
        "    examples.append(example)\n",
        "  \n",
        "  return examples"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jCGUIB172m2G",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def run_hub2emb(args):\n",
        "  '''Runs the embedding generation pipeline'''\n",
        "  source_data_location = args['source_data_location']\n",
        "  sink_data_location = args['sink_data_location']\n",
        "  runner = args['runner']\n",
        "  temporary_dir = args['temporary_location']\n",
        "  module_url = args['module_url']\n",
        "  original_dim = args['original_dim']\n",
        "  projected_dim = args['projected_dim']\n",
        "  batch_size = args['batch_size']\n",
        "\n",
        "  pipeline_options = beam.options.pipeline_options.PipelineOptions(**args)\n",
        "  \n",
        "  random_projection_matrix = None\n",
        "  if projected_dim and original_dim > projected_dim:\n",
        "    # Creating a random projection matrix\n",
        "    random_projection_matrix = np.random.uniform(\n",
        "        low=-1, high=1, size=(original_dim, projected_dim))\n",
        "    print('Storing random projection matrix to disk...')\n",
        "    with open('random_projection_matrix', 'wb') as handle:\n",
        "      pickle.dump(random_projection_matrix, \n",
        "                  handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
        "\n",
        "  print(\"Starting the Beam pipeline...\")\n",
        "  with beam.Pipeline(runner, options=pipeline_options) as pipeline:\n",
        "    (\n",
        "        pipeline\n",
        "        | 'Read sentences from files' >> beam.io.ReadFromText(\n",
        "            file_pattern=source_data_location)\n",
        "        |'Batch elements' >> util.BatchElements(\n",
        "            min_batch_size=batch_size, max_batch_size=batch_size)\n",
        "        | 'Generate emebddings' >> beam.Map(\n",
        "            generate_embeddings, module_url, random_projection_matrix)\n",
        "        | 'Encode to tf example' >> beam.FlatMap(to_tf_example)\n",
        "        | 'Write to TFRecords files' >> beam.io.WriteToTFRecord(\n",
        "            file_path_prefix='{}/emb'.format(sink_data_location),\n",
        "            file_name_suffix='.tfrecords')\n",
        "    )"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "On-MbzD922kb",
        "colab_type": "text"
      },
      "source": [
        "### Run pipeline"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y3I1Wv4i21yY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "runner = 'DirectRunner'\n",
        "job_name = 'hub2emb-{}'.format(datetime.utcnow().strftime('%y%m%d-%H%M%S'))\n",
        "input_data = 'corpus/*.txt'\n",
        "output_dir = 'embeds'\n",
        "temporary_dir = 'tmp'\n",
        "module_url = 'https://tfhub.dev/google/tf2-preview/nnlm-en-dim128/1'\n",
        "\n",
        "original_dim = 128\n",
        "projected_dim = 64\n",
        "\n",
        "args = {\n",
        "    'job_name': job_name,\n",
        "    'runner': runner,\n",
        "    'batch_size': 1024,\n",
        "    'source_data_location': input_data,\n",
        "    'sink_data_location': output_dir,\n",
        "    'temporary_location': temporary_dir,\n",
        "    'module_url': module_url,\n",
        "    'original_dim': original_dim,\n",
        "    'projected_dim': projected_dim,\n",
        "}\n",
        "\n",
        "print(\"Pipeline args are set.\")\n",
        "args"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iS9obmeP4ZOA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!rm -r {output_dir}\n",
        "!rm -r {temporary_dir}\n",
        "!rm random_projection_matrix\n",
        "\n",
        "print(\"Running pipeline...\")\n",
        "%time run_hub2emb(args)\n",
        "print(\"Pipeline is done.\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JAwOo7gQWvVd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!ls {output_dir}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HVnee4e6U90u",
        "colab_type": "text"
      },
      "source": [
        "Read some of the generated embeddings..."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-K7pGXlXOj1N",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "embed_file = '{}/emb-00000-of-00001.tfrecords'.format(output_dir)\n",
        "sample = 5\n",
        "\n",
        "# Create a description of the features.\n",
        "feature_description = {\n",
        "    'text': tf.io.FixedLenFeature([], tf.string),\n",
        "    'embedding': tf.io.FixedLenFeature([projected_dim], tf.float32)\n",
        "}\n",
        "\n",
        "def _parse_example(example):\n",
        "  # Parse the input `tf.Example` proto using the dictionary above.\n",
        "  return tf.io.parse_single_example(example, feature_description)\n",
        "\n",
        "dataset = tf.data.TFRecordDataset(embed_file)\n",
        "for record in dataset.take(sample).map(_parse_example):\n",
        "  print(\"{}:{}\".format(record['text'].numpy().decode('utf-8'), record['embedding'].numpy()[:10]))\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "agGoaMSgY8wN",
        "colab_type": "text"
      },
      "source": [
        "## 3. Build the ANN Index for the Embeddings"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UcPDspU3WjgH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def build_index(embedding_files_pattern, index_filename, vector_length, \n",
        "    metric='angular', num_trees=100):\n",
        "  '''Builds an ANNOY index'''\n",
        "\n",
        "  annoy_index = annoy.AnnoyIndex(vector_length, metric=metric)\n",
        "  # Mapping between the item and its identifier in the index\n",
        "  mapping = {}\n",
        "\n",
        "  embed_files = tf.io.gfile.glob(embedding_files_pattern)\n",
        "  num_files = len(embed_files)\n",
        "  print('Found {} embedding file(s).'.format(num_files))\n",
        "\n",
        "  item_counter = 0\n",
        "  for i, embed_file in enumerate(embed_files):\n",
        "    print('Loading embeddings in file {} of {}...'.format(i+1, num_files))\n",
        "    dataset = tf.data.TFRecordDataset(embed_file)\n",
        "    for record in dataset.map(_parse_example):\n",
        "      text = record['text'].numpy().decode(\"utf-8\")\n",
        "      embedding = record['embedding'].numpy()\n",
        "      mapping[item_counter] = text\n",
        "      annoy_index.add_item(item_counter, embedding)\n",
        "      item_counter += 1\n",
        "      if item_counter % 200000 == 0:\n",
        "        print('{} items loaded to the index'.format(item_counter))\n",
        "\n",
        "  print('A total of {} items added to the index'.format(item_counter))\n",
        "\n",
        "  print('Building the index with {} trees...'.format(num_trees))\n",
        "  annoy_index.build(n_trees=num_trees)\n",
        "  print('Index is successfully built.')\n",
        "  \n",
        "  print('Saving index to disk...')\n",
        "  annoy_index.save(index_filename)\n",
        "  print('Index is saved to disk.')\n",
        "  print(\"Index file size: {} GB\".format(\n",
        "    round(os.path.getsize(index_filename) / float(1024 ** 3), 2)))\n",
        "  annoy_index.unload()\n",
        "\n",
        "  print('Saving mapping to disk...')\n",
        "  with open(index_filename + '.mapping', 'wb') as handle:\n",
        "    pickle.dump(mapping, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
        "  print('Mapping is saved to disk.')\n",
        "  print(\"Mapping file size: {} MB\".format(\n",
        "    round(os.path.getsize(index_filename + '.mapping') / float(1024 ** 2), 2)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AgyOQhUq6FNE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "embedding_files = \"{}/emb-*.tfrecords\".format(output_dir)\n",
        "embedding_dimension = projected_dim\n",
        "index_filename = \"index\"\n",
        "\n",
        "!rm {index_filename}\n",
        "!rm {index_filename}.mapping\n",
        "\n",
        "%time build_index(embedding_files, index_filename, embedding_dimension)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ic31Tm5cgAd5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!ls"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "maGxDl8ufP-p",
        "colab_type": "text"
      },
      "source": [
        "## 4. Use the Index for Similarity Matching"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_dIs8W78fYPp",
        "colab_type": "text"
      },
      "source": [
        "### Load the index and the mapping files"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jlTTrbQHayvb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "index = annoy.AnnoyIndex(embedding_dimension)\n",
        "index.load(index_filename, prefault=True)\n",
        "print('Annoy index is loaded.')\n",
        "with open(index_filename + '.mapping', 'rb') as handle:\n",
        "  mapping = pickle.load(handle)\n",
        "print('Mapping file is loaded.')\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y6liFMSUh08J",
        "colab_type": "text"
      },
      "source": [
        "### Similarity matching method"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mUxjTag8hc16",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def find_similar_items(embedding, num_matches=5):\n",
        "  '''Finds similar items to a given embedding in the ANN index'''\n",
        "  ids = index.get_nns_by_vector(\n",
        "  embedding, num_matches, search_k=-1, include_distances=False)\n",
        "  items = [mapping[i] for i in ids]\n",
        "  return items"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hjerNpmZja0A",
        "colab_type": "text"
      },
      "source": [
        "### Extract embedding from a given query"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a0IIXzfBjZ19",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Load the TF-Hub module\n",
        "print(\"Loading the TF-Hub module...\")\n",
        "embed_fn = hub.load(module_url)\n",
        "print(\"TF-hub module is loaded.\")\n",
        "\n",
        "random_projection_matrix = None\n",
        "if os.path.exists('random_projection_matrix'):\n",
        "  print(\"Loading random projection matrix...\")\n",
        "  with open('random_projection_matrix', 'rb') as handle:\n",
        "    random_projection_matrix = pickle.load(handle)\n",
        "  print('random projection matrix is loaded.')\n",
        "\n",
        "def extract_embeddings(query):\n",
        "  '''Generates the embedding for the query'''\n",
        "  query_embedding =  embed_fn([query])[0].numpy()\n",
        "  if random_projection_matrix is not None:\n",
        "    query_embedding = query_embedding.dot(random_projection_matrix)\n",
        "  return query_embedding\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kCoCNROujEIO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "extract_embeddings(\"Hello Machine Learning!\")[:10]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TaQ0IRhukCRx",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "query = \"confronting global challenges\" #@param {type:\"string\"}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wC0uLjvfk5nB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(\"Generating embedding for the query...\")\n",
        "%time query_embedding = extract_embeddings(query)\n",
        "\n",
        "print(\"\")\n",
        "print(\"Finding relevant items in the index...\")\n",
        "%time items = find_similar_items(query_embedding, 10)\n",
        "\n",
        "print(\"\")\n",
        "print(\"Results:\")\n",
        "print(\"=========\")\n",
        "for item in items:\n",
        "  print(item)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ss9W2EBCx-0r",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}