notebooks/inference/build_index_and_search.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "name": "build_index_and_search.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "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.7"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Santosh-Gupta/NaturalLanguageRecommendations/blob/master/notebooks/inference/build_index_and_search.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qf93M9RFhUj2",
        "colab_type": "text"
      },
      "source": [
        "### Works Best With TPU Runtime"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "8wAw2JV6F4Am",
        "outputId": "21c4a0e0-49df-4623-9d07-9c7110e0092b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 476
        }
      },
      "source": [
        "!gdown --id \"10LV9QbZOkUyOzR4nh8hxesoKJhpmvpM9\"   # citation vectors\n",
        "!gdown --id \"1-23aNm7j0bnycvyd_OaQfofVYPTewgOI\"   # abstract vectors\n",
        "!gdown --id \"1NyUQwgUNj9bFsiCnZ2TfKmWn5r-Y6wav\"   # TitlesIdAbstractsEmbedIds\n",
        "!wget 'https://s3-us-west-2.amazonaws.com/ai2-s2-research/scibert/huggingface_pytorch/scibert_scivocab_uncased.tar'\n",
        "!tar -xvf 'scibert_scivocab_uncased.tar'"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=10LV9QbZOkUyOzR4nh8hxesoKJhpmvpM9\n",
            "To: /content/CitationSimilarityVectors106Epochs.npy\n",
            "2.59GB [00:33, 76.4MB/s]\n",
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1-23aNm7j0bnycvyd_OaQfofVYPTewgOI\n",
            "To: /content/AbstractSimVectors.npy\n",
            "2.59GB [00:42, 61.5MB/s]\n",
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1NyUQwgUNj9bFsiCnZ2TfKmWn5r-Y6wav\n",
            "To: /content/TitlesIdsAbstractsEmbedIdsCOMPLETE_12-30-19.json.gzip\n",
            "432MB [00:06, 69.3MB/s]\n",
            "--2020-01-02 20:33:41--  https://s3-us-west-2.amazonaws.com/ai2-s2-research/scibert/huggingface_pytorch/scibert_scivocab_uncased.tar\n",
            "Resolving s3-us-west-2.amazonaws.com (s3-us-west-2.amazonaws.com)... 52.218.224.136\n",
            "Connecting to s3-us-west-2.amazonaws.com (s3-us-west-2.amazonaws.com)|52.218.224.136|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 442460160 (422M) [application/x-tar]\n",
            "Saving to: ‘scibert_scivocab_uncased.tar’\n",
            "\n",
            "scibert_scivocab_un 100%[===================>] 421.96M  41.6MB/s    in 9.4s    \n",
            "\n",
            "2020-01-02 20:33:51 (45.0 MB/s) - ‘scibert_scivocab_uncased.tar’ saved [442460160/442460160]\n",
            "\n",
            "scibert_scivocab_uncased/\n",
            "scibert_scivocab_uncased/vocab.txt\n",
            "scibert_scivocab_uncased/pytorch_model.bin\n",
            "scibert_scivocab_uncased/config.json\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JP2AqdEiQhbS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install transformers --quiet"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "SGov6e8uGO3B",
        "outputId": "d655c492-b31b-44f9-d7d8-dfdd15b5a54c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "%tensorflow_version 2.x\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from time import time\n",
        "from tqdm import tqdm_notebook as tqdm\n",
        "from transformers import BertTokenizer\n",
        "import pandas as pd\n",
        "from pprint import pprint\n",
        "import psutil\n",
        "import gc\n",
        "\n",
        "print('TensorFlow:', tf.__version__)"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "TensorFlow 2.x selected.\n",
            "TensorFlow: 2.1.0-rc1\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "iht3WQa3H_g6",
        "outputId": "cbb9e356-74ae-4376-e95c-4b3a72d40042",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 748
        }
      },
      "source": [
        "try:\n",
        "    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n",
        "    print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])\n",
        "except ValueError:\n",
        "    tpu = None\n",
        "\n",
        "if tpu:\n",
        "    tf.config.experimental_connect_to_cluster(tpu)\n",
        "    tf.tpu.experimental.initialize_tpu_system(tpu)\n",
        "    strategy = tf.distribute.experimental.TPUStrategy(tpu)\n",
        "else:\n",
        "    strategy = tf.distribute.MirroredStrategy()\n",
        "\n",
        "print(\"REPLICAS: \", strategy.num_replicas_in_sync)"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:absl:Entering into master device scope: /job:worker/replica:0/task:0/device:CPU:0\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Running on TPU  ['10.108.186.42:8470']\n",
            "INFO:tensorflow:Initializing the TPU system: 10.108.186.42:8470\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Initializing the TPU system: 10.108.186.42:8470\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Clearing out eager caches\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Clearing out eager caches\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Finished initializing TPU system.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Finished initializing TPU system.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Found TPU system:\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Found TPU system:\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores: 8\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores: 8\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Workers: 1\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Workers: 1\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "REPLICAS:  8\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "3ks95sxuIKbR",
        "outputId": "14cd3350-e37c-4853-9eec-4044c8065415",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 54
        }
      },
      "source": [
        "if tpu:\n",
        "  workers = ['/job:worker/replica:0/task:0/device:TPU:'+str(i) for i in range(8)]\n",
        "  print(workers)\n",
        "  gpu =''\n",
        "if tf.config.experimental.list_physical_devices('GPU'):\n",
        "  gpu = 'gpu'\n",
        "  workers = ['/GPU:0']\n",
        "  pprint(workers)"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "['/job:worker/replica:0/task:0/device:TPU:0', '/job:worker/replica:0/task:0/device:TPU:1', '/job:worker/replica:0/task:0/device:TPU:2', '/job:worker/replica:0/task:0/device:TPU:3', '/job:worker/replica:0/task:0/device:TPU:4', '/job:worker/replica:0/task:0/device:TPU:5', '/job:worker/replica:0/task:0/device:TPU:6', '/job:worker/replica:0/task:0/device:TPU:7']\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "uByeXRrVO78_",
        "colab": {}
      },
      "source": [
        "class Index:\n",
        "    def __init__(self, embeddings, worker):\n",
        "        self.embeddings = tf.math.l2_normalize(embeddings, axis=1)\n",
        "        if 'TPU' in worker:\n",
        "          self.embeddings = tf.cast(self.embeddings, dtype=tf.bfloat16)\n",
        "        self.worker = worker\n",
        "\n",
        "    @tf.function\n",
        "    def search(self, query_vector):\n",
        "      with tf.device(self.worker):\n",
        "        dot_product = tf.reduce_sum(tf.multiply(self.embeddings, query_vector), axis=1)\n",
        "        distances = 1 - dot_product\n",
        "        sorted_indices =  tf.argsort(distances)\n",
        "        nearest_distances = tf.cast(tf.gather(distances, sorted_indices), dtype=tf.float32)\n",
        "        return nearest_distances[:20], sorted_indices[:20]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f_YaU7dUB6-P",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "26435eb8-edd7-47ae-c348-92b033358de8"
      },
      "source": [
        "print('System Memory used    : ' + str(psutil.virtual_memory().percent) + '%')"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "System Memory used    : 4.7%\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "DFloyWt7G8ts",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "e635c542-bb67-4e28-d986-4e3298b36a4f"
      },
      "source": [
        "citations_embeddings = np.load('CitationSimilarityVectors106Epochs.npy')\n",
        "abstract_embeddings = np.load('AbstractSimVectors.npy')\n",
        "assert citations_embeddings.shape == abstract_embeddings.shape\n",
        "print('System Memory used    : ' + str(psutil.virtual_memory().percent) + '%')"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "System Memory used    : 18.4%\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "AB8L4sA-H2nf",
        "outputId": "d6fb6357-8895-444c-e1d6-b1fb5a50f056",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "if tpu:\n",
        "  # Discarding last 4 vectors to make number of vectors divisible by 8\n",
        "  citations_embeddings = np.split(citations_embeddings[:-4], 8, axis=0)\n",
        "  abstract_embeddings = np.split(abstract_embeddings[:-4], 8, axis=0)\n",
        "  vecs_per_index = citations_embeddings[0].shape[0]\n",
        "  print('Vectors per index :', vecs_per_index)\n",
        "elif gpu:\n",
        "  vecs_per_index = citations_embeddings.shape[0]\n",
        "  print('Vectors per index :', vecs_per_index)\n",
        "print('System Memory used    : ' + str(psutil.virtual_memory().percent) + '%')"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Vectors per index : 157874\n",
            "System Memory used    : 18.4%\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "HoEA0SqeJW4n",
        "outputId": "df7bb695-37da-49ab-d727-cdd9b46c51e6",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 170
        }
      },
      "source": [
        "citation_indices = []\n",
        "abstract_indices = []\n",
        "if gpu:\n",
        "  for i, worker in enumerate(workers):\n",
        "    with tf.device(worker):\n",
        "      print('Building index with {} vectors on {}'.format(citations_embeddings.shape[0], worker))\n",
        "      citation_indices.append(Index(citations_embeddings, worker))\n",
        "      abstract_indices.append(Index(abstract_embeddings, worker))\n",
        "elif tpu:\n",
        "  ## Place 1/8 of total embeddings on each TPU core\n",
        "  for i, worker in enumerate(workers):\n",
        "    with tf.device(worker):\n",
        "      print('Building index with {} vectors on {}'.format(citations_embeddings[i].shape[0],worker))\n",
        "      citation_indices.append(Index(citations_embeddings[i].copy(), worker))\n",
        "      abstract_indices.append(Index(abstract_embeddings[i].copy(), worker))\n",
        "\n",
        "del citations_embeddings\n",
        "del abstract_embeddings\n",
        "print('System Memory used    : ' + str(psutil.virtual_memory().percent) + '%')"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Building index with 157874 vectors on /job:worker/replica:0/task:0/device:TPU:0\n",
            "Building index with 157874 vectors on /job:worker/replica:0/task:0/device:TPU:1\n",
            "Building index with 157874 vectors on /job:worker/replica:0/task:0/device:TPU:2\n",
            "Building index with 157874 vectors on /job:worker/replica:0/task:0/device:TPU:3\n",
            "Building index with 157874 vectors on /job:worker/replica:0/task:0/device:TPU:4\n",
            "Building index with 157874 vectors on /job:worker/replica:0/task:0/device:TPU:5\n",
            "Building index with 157874 vectors on /job:worker/replica:0/task:0/device:TPU:6\n",
            "Building index with 157874 vectors on /job:worker/replica:0/task:0/device:TPU:7\n",
            "System Memory used    : 6.7%\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "W5r4LmcVqGAZ",
        "colab": {}
      },
      "source": [
        "def search(xq, top_k=10):\n",
        "  cD, cI = [], []\n",
        "  aD, aI = [], []\n",
        "  if tpu:\n",
        "    r = 8\n",
        "  else:\n",
        "    r = 1\n",
        "  for i in range(r):\n",
        "    print('Search running on {}'.format(citation_indices[i].worker))\n",
        "    cd, cidx = citation_indices[i].search(xq)\n",
        "    ad, aidx = abstract_indices[i].search(xq)\n",
        "\n",
        "    cD.extend(cd.numpy())\n",
        "    aD.extend(ad.numpy())\n",
        "\n",
        "    cI.extend(i*vecs_per_index + cidx.numpy())\n",
        "    aI.extend(i*vecs_per_index + aidx.numpy())\n",
        "\n",
        "  cid_sorted = np.argsort(cD)[:top_k]\n",
        "  aid_sorted = np.argsort(aD)[:top_k]\n",
        "\n",
        "  cD = np.array(cD)[cid_sorted]\n",
        "  aD = np.array(aD)[aid_sorted]\n",
        "\n",
        "  cI = np.array(cI)[cid_sorted]\n",
        "  aI = np.array(aI)[aid_sorted]\n",
        "  return cD, aD, cI, aI"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Syc0k8Yi4rcz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = tf.saved_model.load('gs://tfworld/saved_models')\n",
        "tokenizer = BertTokenizer(vocab_file='scibert_scivocab_uncased/vocab.txt')\n",
        "\n",
        "df = pd.read_json('/content/TitlesIdsAbstractsEmbedIdsCOMPLETE_12-30-19.json.gzip', compression = 'gzip')\n",
        "embed2Title = pd.Series(df['title'].values,index=df['EmbeddingID']).to_dict()\n",
        "embed2Abstract = pd.Series(df['paperAbstract'].values,index=df['EmbeddingID']).to_dict()\n",
        "embed2Paper = pd.Series(df['id'].values,index=df['EmbeddingID']).to_dict()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jxkcuP0CDNKM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# embed_id = 70000\n",
        "# title = embed2Title[embed_id]\n",
        "#abstract = embed2Abstract[embed_id]\n",
        "# abstract = '''Forecasting stock market direction is always an amazing but challenging problem in finance. Although many popular shallow computational methods (such as Backpropagation Network and Support Vector Machine) have extensively been proposed, most algorithms have not yet attained a desirable level of applicability. In this paper, we present a deep learning model with strong ability to generate high level feature representations for accurate financial prediction. Precisely, a stacked denoising autoencoder (SDAE) from deep learning is applied to predict the daily CSI 300 index, from Shanghai and Shenzhen Stock Exchanges in China. We use six evaluation criteria to evaluate its performance compared with the back propagation network, support vector machine. The experiment shows that the underlying financial model with deep machine technology has a significant advantage for the prediction of the CSI 300 index.'''\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QFopFGOqit5y",
        "colab_type": "code",
        "outputId": "b4a8328b-b9ad-4084-d44a-fdaf5c8a053f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "abstract = \"neural machine translation\" #@param {type:\"string\"}\n",
        "abstract_encoded = tokenizer.encode(abstract, max_length=512, pad_to_max_length=True)\n",
        "abstract_encoded = tf.constant(abstract_encoded, dtype=tf.int32)[None, :]\n",
        "print('\\nAbstract : ')\n",
        "pprint(abstract)"
      ],
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Abstract : \n",
            "'neural machine translation'\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p-_vAr8z8Hh5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "s = time()\n",
        "bert_output = model(abstract_encoded)\n",
        "xq = tf.nn.l2_normalize(bert_output, axis=1)\n",
        "prediction_time = time() - s"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EhiyvUxrF4xr",
        "colab_type": "code",
        "outputId": "a61d1e46-7ef0-4084-9a7c-edf6d1cb8e1e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "s_s = time()\n",
        "if tpu:\n",
        "  xq = tf.cast(xq, dtype=tf.bfloat16)\n",
        "\n",
        "cD, aD, cI, aI = search(xq, top_k=5)\n",
        "search_time = time() - s_s\n",
        "total_time = prediction_time + search_time\n",
        "print('\\n'*2)\n",
        "print('Prediction time       :', np.round(prediction_time, 3), 'secs')\n",
        "print('Search time           :', np.round(search_time, 3), 'secs')\n",
        "print('Total time            :', np.round(total_time, 3), 'secs')\n",
        "print('System Memory used    : ' + str(psutil.virtual_memory().percent) + '%')\n",
        "\n",
        "print('\\n'*2)\n",
        "print('*'*80)\n",
        "for i in range(len(cI)):\n",
        "  print('Title : ')\n",
        "  pprint(embed2Title[cI[i]])\n",
        "  print('\\n')\n",
        "  pprint('Link: semanticscholar.org/paper/'+embed2Paper[cI[i]])\n",
        "  print('*'*80, )\n",
        "print('\\nNeighbours       :', cI )\n",
        "print('Distances        :', np.round(cD, 4))\n",
        "\n",
        "print('\\n'*4)\n",
        "print('*'*80)\n",
        "for i in range(len(aI)):\n",
        "  print('Abstract : ')\n",
        "  pprint(embed2Abstract[aI[i]])\n",
        "  print('\\n')\n",
        "  pprint('Link: semanticscholar.org/paper/'+embed2Paper[aI[i]])\n",
        "  print('*'*80)\n",
        "print('\\nNeighbours       :', aI )\n",
        "print('Distances        :', np.round(aD, 4))"
      ],
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Search running on /job:worker/replica:0/task:0/device:TPU:0\n",
            "Search running on /job:worker/replica:0/task:0/device:TPU:1\n",
            "Search running on /job:worker/replica:0/task:0/device:TPU:2\n",
            "Search running on /job:worker/replica:0/task:0/device:TPU:3\n",
            "Search running on /job:worker/replica:0/task:0/device:TPU:4\n",
            "Search running on /job:worker/replica:0/task:0/device:TPU:5\n",
            "Search running on /job:worker/replica:0/task:0/device:TPU:6\n",
            "Search running on /job:worker/replica:0/task:0/device:TPU:7\n",
            "\n",
            "\n",
            "\n",
            "Prediction time       : 0.002 secs\n",
            "Search time           : 0.293 secs\n",
            "Total time            : 0.295 secs\n",
            "System Memory used    : 6.5%\n",
            "\n",
            "\n",
            "\n",
            "********************************************************************************\n",
            "Title : \n",
            "'Traffic incident detection algorithm based on non-parameter regression'\n",
            "\n",
            "\n",
            "'Link: semanticscholar.org/paper/f6ddaf3fa9ff849ec863705aa2ddeec246366112'\n",
            "********************************************************************************\n",
            "Title : \n",
            "('Feedback Effects in Plasmonic Slot Waveguides Examined Using a Closed Form '\n",
            " 'Model')\n",
            "\n",
            "\n",
            "'Link: semanticscholar.org/paper/219190aa8b7690ad6524e15b1617f4eedf389de0'\n",
            "********************************************************************************\n",
            "Title : \n",
            "'The gait identification challenge problem: data sets and baseline algorithm'\n",
            "\n",
            "\n",
            "'Link: semanticscholar.org/paper/06637fecb6cd2b600137451f8ea817c6c1f3a5be'\n",
            "********************************************************************************\n",
            "Title : \n",
            "('Effect of cryogenic deformation on the structure and properties of '\n",
            " 'chromium-nickel steels')\n",
            "\n",
            "\n",
            "'Link: semanticscholar.org/paper/c68d5fabaf75237abc88d32f0cadcfc4b21ab852'\n",
            "********************************************************************************\n",
            "Title : \n",
            "('A 30 V/1 MHz AC/AC converter for high frequency AC distributed power system '\n",
            " 'applications')\n",
            "\n",
            "\n",
            "'Link: semanticscholar.org/paper/6d6a7eb62299c12e400357332bcd2b22dd084bdd'\n",
            "********************************************************************************\n",
            "\n",
            "Neighbours       : [  96818   25686  438411  111103 1159377]\n",
            "Distances        : [0.6914 0.6953 0.6992 0.7031 0.707 ]\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "********************************************************************************\n",
            "Abstract : \n",
            "('This paper discusses machine translation with regards to its '\n",
            " 'development,accuracy,readability, classification,application,and existing '\n",
            " 'problems.At the same time,it has analyzed several viewpoints held by people '\n",
            " 'towards machine translation.')\n",
            "\n",
            "\n",
            "'Link: semanticscholar.org/paper/311b26f623a2c61f029a1f0fd62bd34531b47aec'\n",
            "********************************************************************************\n",
            "Abstract : \n",
            "('The emergence of neural machine translation (NMT) has revolutionized the '\n",
            " 'filed of machine translation. In the first section, we introduce the '\n",
            " 'fundamental of NMT models. Then we study the advantages of NMT over '\n",
            " 'traditional statistical machine translations (SMT), some of its existing '\n",
            " 'challenges and recent research efforts. Finally, we summarize our four works '\n",
            " 'that address several existing problems in NMT. In the second section, we '\n",
            " 'follow attentional neural machine translation to use mathematical formula to '\n",
            " 'define the overall procedure for NMT.')\n",
            "\n",
            "\n",
            "'Link: semanticscholar.org/paper/6798b1c6e0b208f3e3f567e5e6172ffd0f6a52d7'\n",
            "********************************************************************************\n",
            "Abstract : \n",
            "('In the last decade, the statistical approach has found widespread use in '\n",
            " 'machine translation both for written and spoken language and has had a major '\n",
            " 'impact on the translation accuracy. The goal of this paper is to cover the '\n",
            " 'state of the art in statistical machine translation. We would re-visit the '\n",
            " 'underlying principles of the statistical approach to machine translation and '\n",
            " 'summarize the progress that has been made over the last decade')\n",
            "\n",
            "\n",
            "'Link: semanticscholar.org/paper/06c2f8e67421b1bce89d6872148bf8a4c6943368'\n",
            "********************************************************************************\n",
            "Abstract : \n",
            "('The use of neural machine translation (NMT) in a professional scenario '\n",
            " 'implies a number of challenges despite growing evidence that, in language '\n",
            " 'combinations such as English to Spanish, NMT output quality has already '\n",
            " 'outperformed statistical machine translation in terms of automatic metric '\n",
            " 'scores. This article presents the result of an empirical test that aims to '\n",
            " 'shed light on the differences between NMT post-editing and translation with '\n",
            " 'the aid of a translation memory (TM). The results show that NMT post-editing '\n",
            " 'involves less editing than TM segments, but this editing appears to take '\n",
            " 'more time, with the consequence that NMT post-editing does not seem to '\n",
            " 'improve productivity as may have been expected. This might be due to the '\n",
            " 'fact that NMT segments show a higher variability in terms of quality and '\n",
            " 'time invested in post-editing than TM segments that are ‘more similar’ on '\n",
            " 'average. Finally, results show that translators who perceive that NMT boosts '\n",
            " 'their productivity actually performed faster than those who perceive that '\n",
            " 'NMT slows them down.')\n",
            "\n",
            "\n",
            "'Link: semanticscholar.org/paper/92cbb1590f1568b328c68ad9e8625029ea7ccdcc'\n",
            "********************************************************************************\n",
            "Abstract : \n",
            "('Sentence-level parallel data is essential for training machine translation '\n",
            " 'systems. However, existing parallel data is extremely limited for thousands '\n",
            " 'of languages. In order to increase the available parallel data for a '\n",
            " 'low-resource language we borrow parallel data from a higher-resource closely '\n",
            " 'related language (RL). In so doing we propose a method for translating texts '\n",
            " 'from RL to the low-resource language without requiring any parallel data '\n",
            " 'between them. We use this method to convert RL/English parallel data and use '\n",
            " 'it as an extra resource for machine translation. We show that this extra '\n",
            " 'parallel data highly helps the BLEU score.')\n",
            "\n",
            "\n",
            "'Link: semanticscholar.org/paper/d0465fe6230dfba48f0d05a189da35fffbc48faf'\n",
            "********************************************************************************\n",
            "\n",
            "Neighbours       : [542103 266804 901541 148399 747599]\n",
            "Distances        : [0.418  0.4609 0.5234 0.5234 0.5312]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l-H5itqmRMq7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}