TensorFlow/LanguageModeling/BERT/notebooks/bert_squad_tf_inference_colab.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jDXroBuNw60P"
   },
   "outputs": [],
   "source": [
    "# Copyright 2021 NVIDIA Corporation. 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",
    "# =============================================================================="
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DXgkABw0IUIC"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NVIDIA/DeepLearningExamples/blob/master/TensorFlow/LanguageModeling/BERT/notebooks/bert_squad_tf_inference_colab.ipynb#scrollTo=5hRb96NKE3X0\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "k-XnFINow60d"
   },
   "source": [
    "<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# BERT Question Answering Inference with Mixed Precision\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TfF7V662w60j"
   },
   "source": [
    "## 1. Overview\n",
    "\n",
    "Bidirectional Embedding Representations from Transformers (BERT), is a method of pre-training language representations which obtains state-of-the-art results on a wide array of Natural Language Processing (NLP) tasks. \n",
    "\n",
    "The original paper can be found here: https://arxiv.org/abs/1810.04805.\n",
    "\n",
    "NVIDIA's BERT 19.10 is an optimized version of Google's official implementation, leveraging mixed precision arithmetic and tensor cores on V100 GPUS for faster training times while maintaining target accuracy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ah3Lv9zyw60l"
   },
   "source": [
    "### 1.a Learning objectives\n",
    "\n",
    "This notebook demonstrates:\n",
    "- Inference on QA task with BERT Large model\n",
    "- The use/download of fine-tuned NVIDIA BERT models\n",
    "- Use of Mixed Precision for Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hxNJ8HByw60o"
   },
   "source": [
    "## 2. Requirements\n",
    "\n",
    "### 2.a GPU\n",
    "\n",
    "Before running this notebook, please set the Colab runtime environment to GPU via the menu *Runtime => Change runtime type => GPU*.\n",
    "\n",
    "This demo will work on any NVIDIA GPU with CUDA cores, though for improved FP16 inference, a Volta, Turing or newer generation GPU with Tensor cores is desired.  On Google Colab, this normally means a T4 GPU. If you are assigned an older K80 GPU, another trial at another time might give you a T4 GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "zyMjTVx1Ibt-",
    "outputId": "216eeecc-3785-466b-d126-94f02abe19d8"
   },
   "outputs": [],
   "source": [
    "#Select lower version of tensroflow on Google Colab\n",
    "%tensorflow_version 1.x\n",
    "import tensorflow\n",
    "print(tensorflow.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1PXBHL8KIUIF",
    "outputId": "cb05ba96-4d21-4483-de95-dd724383dd80"
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BgfmN3z6IUIF"
   },
   "source": [
    "### 2.b Download the required files from NVIDIA-Github:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KV_WnOY4zUa_",
    "outputId": "a1b0f548-75f5-432d-8fbd-6fbad5f421a3"
   },
   "outputs": [],
   "source": [
    "!wget -nc -q --show-progress -O ./master.zip \\\n",
    "https://github.com/NVIDIA/DeepLearningExamples/archive/master.zip\n",
    "!unzip -q -n -d . ./master.zip "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5D7i7Pao5qoj",
    "outputId": "d172417e-1cb1-4138-e1f4-1eda0f3d72bd"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "WORKSPACE_DIR='./DeepLearningExamples-master/TensorFlow/LanguageModeling/BERT/'\n",
    "os.chdir(WORKSPACE_DIR)\n",
    "print (os.getcwd())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mjlZbP0dw60r"
   },
   "source": [
    "## 3. BERT Inference: Question Answering\n",
    "\n",
    "We can run inference on a fine-tuned BERT model for tasks like Question Answering.\n",
    "\n",
    "Here we use a BERT model fine-tuned on a [SQuaD 2.0 Dataset](https://rajpurkar.github.io/SQuAD-explorer/) which contains 100,000+ question-answer pairs on 500+ articles combined with over 50,000 new, unanswerable questions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mOc16svBw60t"
   },
   "source": [
    "### 3.a Paragraph and Queries\n",
    "\n",
    "In this example we will ask our BERT model questions related to the following paragraph:\n",
    "\n",
    "**The Apollo Program**\n",
    "_\"The Apollo program, also known as Project Apollo, was the third United States human spaceflight program carried out by the National Aeronautics and Space Administration (NASA), which accomplished landing the first humans on the Moon from 1969 to 1972. First conceived during Dwight D. Eisenhower's administration as a three-man spacecraft to follow the one-man Project Mercury which put the first Americans in space, Apollo was later dedicated to President John F. Kennedy's national goal of landing a man on the Moon and returning him safely to the Earth by the end of the 1960s, which he proposed in a May 25, 1961, address to Congress. Project Mercury was followed by the two-man Project Gemini. The first manned flight of Apollo was in 1968. Apollo ran from 1961 to 1972, and was supported by the two-man Gemini program which ran concurrently with it from 1962 to 1966. Gemini missions developed some of the space travel techniques that were necessary for the success of the Apollo missions. Apollo used Saturn family rockets as launch vehicles. Apollo/Saturn vehicles were also used for an Apollo Applications Program, which consisted of Skylab, a space station that supported three manned missions in 1973-74, and the Apollo-Soyuz Test Project, a joint Earth orbit mission with the Soviet Union in 1975.\"_\n",
    "\n",
    "The questions and relative answers expected are shown below:\n",
    "\n",
    " - **Q1:** \"What project put the first Americans into space?\" \n",
    "  - **A1:** \"Project Mercury\"\n",
    " - **Q2:** \"What program was created to carry out these projects and missions?\"\n",
    "  - **A2:** \"The Apollo program\"\n",
    " - **Q3:** \"What year did the first manned Apollo flight occur?\"\n",
    "  - **A3:** \"1968\"\n",
    " - **Q4:** \"What President is credited with the original notion of putting Americans in space?\"\n",
    "  - **A4:** \"John F. Kennedy\"\n",
    " - **Q5:** \"Who did the U.S. collaborate with on an Earth orbit mission in 1975?\"\n",
    "  - **A5:** \"Soviet Union\"\n",
    " - **Q6:** \"How long did Project Apollo run?\"\n",
    "  - **A6:** \"1961 to 1972\"\n",
    " - **Q7:** \"What program helped develop space travel techniques that Project Apollo used?\"\n",
    "  - **A7:** \"Gemini Mission\"\n",
    " - **Q8:** \"What space station supported three manned missions in 1973-1974?\"\n",
    "  - **A8:** \"Skylab\"\n",
    "  \n",
    "---\n",
    "\n",
    "The paragraph and the questions can be easily customized by changing the code below:\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "srU0TT1Iw60v",
    "outputId": "b5340882-7b24-4343-ebc2-9ff709b489a3"
   },
   "outputs": [],
   "source": [
    "%%writefile input.json\n",
    "{\"data\": \n",
    " [\n",
    "     {\"title\": \"Project Apollo\",\n",
    "      \"paragraphs\": [\n",
    "          {\"context\":\"The Apollo program, also known as Project Apollo, was the third United States human spaceflight program carried out by the National Aeronautics and Space Administration (NASA), which accomplished landing the first humans on the Moon from 1969 to 1972. First conceived during Dwight D. Eisenhower's administration as a three-man spacecraft to follow the one-man Project Mercury which put the first Americans in space, Apollo was later dedicated to President John F. Kennedy's national goal of landing a man on the Moon and returning him safely to the Earth by the end of the 1960s, which he proposed in a May 25, 1961, address to Congress. Project Mercury was followed by the two-man Project Gemini. The first manned flight of Apollo was in 1968. Apollo ran from 1961 to 1972, and was supported by the two man Gemini program which ran concurrently with it from 1962 to 1966. Gemini missions developed some of the space travel techniques that were necessary for the success of the Apollo missions. Apollo used Saturn family rockets as launch vehicles. Apollo/Saturn vehicles were also used for an Apollo Applications Program, which consisted of Skylab, a space station that supported three manned missions in 1973-74, and the Apollo-Soyuz Test Project, a joint Earth orbit mission with the Soviet Union in 1975.\", \n",
    "           \"qas\": [\n",
    "               { \"question\": \"What project put the first Americans into space?\", \n",
    "                 \"id\": \"Q1\"\n",
    "               },\n",
    "               { \"question\": \"What program was created to carry out these projects and missions?\",\n",
    "                 \"id\": \"Q2\"\n",
    "               },\n",
    "               { \"question\": \"What year did the first manned Apollo flight occur?\",\n",
    "                 \"id\": \"Q3\"\n",
    "               },                \n",
    "               { \"question\": \"What President is credited with the original notion of putting Americans in space?\",\n",
    "                 \"id\": \"Q4\"\n",
    "               },\n",
    "               { \"question\": \"Who did the U.S. collaborate with on an Earth orbit mission in 1975?\",\n",
    "                 \"id\": \"Q5\"\n",
    "               },\n",
    "               { \"question\": \"How long did Project Apollo run?\",\n",
    "                 \"id\": \"Q6\"\n",
    "               },               \n",
    "               { \"question\": \"What program helped develop space travel techniques that Project Apollo used?\",\n",
    "                 \"id\": \"Q7\"\n",
    "               },                \n",
    "               {\"question\": \"What space station supported three manned missions in 1973-1974?\",\n",
    "                 \"id\": \"Q8\"\n",
    "               }                \n",
    "]}]}]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ujyka-8Iw603"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "working_dir = os.getcwd();\n",
    "data_dir = os.path.join(working_dir, 'data/download');\n",
    "if working_dir not in sys.path:\n",
    "    sys.path.append(working_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6gA3-6LVw61D"
   },
   "outputs": [],
   "source": [
    "input_file = os.path.join(working_dir, 'input.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "D9p8XaBnw61N"
   },
   "source": [
    "### 3.b Mixed Precision\n",
    "\n",
    "Mixed precision training offers significant computational speedup by performing operations in half-precision format, while storing minimal information in single-precision to retain as much information as possible in critical parts of the network. Since the introduction of tensor cores in the Volta and Turing architectures, significant training speedups are experienced by switching to mixed precision -- up to 3x overall speedup on the most arithmetically intense model architectures.\n",
    "\n",
    "For information about:\n",
    "- How to train using mixed precision, see the [Mixed Precision Training](https://arxiv.org/abs/1710.03740) paper and [Training With Mixed Precision](https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/index.html) documentation.\n",
    "- How to access and enable AMP for TensorFlow, see [Using TF-AMP](https://docs.nvidia.com/deeplearning/dgx/tensorflow-user-guide/index.html#tfamp) from the TensorFlow User Guide.\n",
    "- Techniques used for mixed precision training, see the [Mixed-Precision Training of Deep Neural Networks](https://devblogs.nvidia.com/mixed-precision-training-deep-neural-networks/) blog."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ceeYPqQcw61P"
   },
   "source": [
    "In this notebook we control mixed precision execution with the environmental variable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "k4jIJevFw61R"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"TF_ENABLE_AUTO_MIXED_PRECISION\"] = \"1\" "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iu4Jb5puw61p"
   },
   "source": [
    "## 4. Fine-Tuned NVIDIA BERT TF Models\n",
    "\n",
    "Based on the model size, we have the following two default configurations of BERT.\n",
    "\n",
    "| **Model** | **Hidden layers** | **Hidden unit size** | **Attention heads** | **Feedforward filter size** | **Max sequence length** | **Parameters** |\n",
    "|:---------:|:----------:|:----:|:---:|:--------:|:---:|:----:|\n",
    "|BERTBASE |12 encoder| 768| 12|4 x  768|512|110M|\n",
    "|BERTLARGE|24 encoder|1024| 16|4 x 1024|512|330M|\n",
    "\n",
    "We will take advantage of the fine-tuned models available on NGC (NVIDIA GPU Cluster, https://ngc.nvidia.com).\n",
    "Among the many configurations available we will download these two:\n",
    "\n",
    " - **bert_tf_ckpt_large_qa_squad2_amp_384**\n",
    "\n",
    "Which are trained on the SQuaD 2.0 Dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5JWKZfP8w61t",
    "outputId": "c5b92300-1f61-435e-9f6c-e3b3c116161e"
   },
   "outputs": [],
   "source": [
    "# bert_tf_ckpt_large_qa_squad2_amp_384\n",
    "DATA_DIR_FT = os.path.join(data_dir, 'finetuned_large_model')\n",
    "!mkdir -p $DATA_DIR_FT    \n",
    "!wget --content-disposition -O $DATA_DIR_FT/bert_tf_ckpt_large_qa_squad2_amp_384_19.03.1.zip  \\\n",
    "https://api.ngc.nvidia.com/v2/models/nvidia/bert_tf_ckpt_large_qa_squad2_amp_384/versions/19.03.1/zip \\\n",
    "&& unzip -n -d $DATA_DIR_FT/ $DATA_DIR_FT/bert_tf_ckpt_large_qa_squad2_amp_384_19.03.1.zip \\\n",
    "&& rm $DATA_DIR_FT/bert_tf_ckpt_large_qa_squad2_amp_384_19.03.1.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GrFrZickw61z"
   },
   "source": [
    "In the code that follows we will refer to these models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cU8mGJDa1FfX"
   },
   "source": [
    "Download the Google pretrained weights and vocab file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5hRb96NKE3X0",
    "outputId": "85b50f58-d89f-41b1-8ac1-55c3117743ea"
   },
   "outputs": [],
   "source": [
    "os.chdir(\"./data\");\n",
    "from GooglePretrainedWeightDownloader import GooglePretrainedWeightDownloader\n",
    "gd = GooglePretrainedWeightDownloader(data_dir)\n",
    "gd.download()\n",
    "os.chdir(\"..\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VY1Dipam15DE"
   },
   "source": [
    "We need the horovod package:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jqAJob92C2wA"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    __import__(\"horovod\")\n",
    "except ImportError:\n",
    "    os.system(\"pip install --no-cache-dir horovod\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5NuuGNsDw611"
   },
   "source": [
    "## 5. Running QA task inference\n",
    "\n",
    "In order to run QA inference we will follow step-by-step the flow implemented in run_squad.py.\n",
    "\n",
    "Configuration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_c2qCQ9-w613"
   },
   "outputs": [],
   "source": [
    "import run_squad\n",
    "import json\n",
    "import tensorflow as tf\n",
    "import modeling\n",
    "import tokenization\n",
    "import time\n",
    "import random\n",
    "\n",
    "tf.logging.set_verbosity(tf.logging.INFO)\n",
    "\n",
    "# Create the output directory where all the results are saved.\n",
    "output_dir = os.path.join(working_dir, 'results')\n",
    "tf.gfile.MakeDirs(output_dir)\n",
    "\n",
    "# The config json file corresponding to the pre-trained BERT model.\n",
    "# This specifies the model architecture.\n",
    "bert_config_file = os.path.join(data_dir, 'finetuned_large_model_SQUAD2.0/bert_config.json')\n",
    "\n",
    "# The vocabulary file that the BERT model was trained on.\n",
    "vocab_file = os.path.join(data_dir, 'finetuned_large_model_SQUAD2.0/vocab.txt')\n",
    "\n",
    "# Initiate checkpoint to the fine-tuned BERT Large model\n",
    "init_checkpoint = os.path.join(data_dir, 'finetuned_large_model/model.ckpt')\n",
    "\n",
    "\n",
    "# Whether to lower case the input text. \n",
    "# Should be True for uncased models and False for cased models.\n",
    "do_lower_case = True\n",
    "  \n",
    "# Total batch size for predictions\n",
    "predict_batch_size = 1\n",
    "params = dict([('batch_size', predict_batch_size)])\n",
    "\n",
    "# The maximum total input sequence length after WordPiece tokenization. \n",
    "# Sequences longer than this will be truncated, and sequences shorter than this will be padded.\n",
    "max_seq_length = 384\n",
    "\n",
    "# When splitting up a long document into chunks, how much stride to take between chunks.\n",
    "doc_stride = 128\n",
    "\n",
    "# The maximum number of tokens for the question. \n",
    "# Questions longer than this will be truncated to this length.\n",
    "max_query_length = 64\n",
    "\n",
    "# This is a WA to use flags from here:\n",
    "flags = tf.flags\n",
    "\n",
    "if 'f' not in tf.flags.FLAGS: \n",
    "    tf.app.flags.DEFINE_string('f', '', 'kernel')\n",
    "FLAGS = flags.FLAGS\n",
    "\n",
    "verbose_logging = True\n",
    "# Set to True if the dataset has samples with no answers. For SQuAD 1.1, this is set to False\n",
    "version_2_with_negative = False\n",
    "\n",
    "# The total number of n-best predictions to generate in the nbest_predictions.json output file.\n",
    "n_best_size = 20\n",
    "\n",
    "# The maximum length of an answer that can be generated. \n",
    "# This is needed  because the start and end predictions are not conditioned on one another.\n",
    "max_answer_length = 30"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2h_eLUgPw618"
   },
   "source": [
    "Let's define the tokenizer and create the model for the estimator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RXHdoUb9w619",
    "outputId": "6079b4b3-c535-4b0e-aa17-d8a505cdb3ea"
   },
   "outputs": [],
   "source": [
    "# Validate the casing config consistency with the checkpoint name.\n",
    "tokenization.validate_case_matches_checkpoint(do_lower_case, init_checkpoint)\n",
    "\n",
    "# Create the tokenizer.\n",
    "tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file, do_lower_case=do_lower_case)\n",
    "\n",
    "# Load the configuration from file\n",
    "bert_config = modeling.BertConfig.from_json_file(bert_config_file)\n",
    "\n",
    "def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument\n",
    "    unique_ids = features[\"unique_ids\"]\n",
    "    input_ids = features[\"input_ids\"]\n",
    "    input_mask = features[\"input_mask\"]\n",
    "    segment_ids = features[\"segment_ids\"]\n",
    "\n",
    "    (start_logits, end_logits) = run_squad.create_model(\n",
    "        bert_config=bert_config,\n",
    "        is_training=False,\n",
    "        input_ids=input_ids,\n",
    "        input_mask=input_mask,\n",
    "        segment_ids=segment_ids,\n",
    "        use_one_hot_embeddings=False)\n",
    "\n",
    "    tvars = tf.trainable_variables()\n",
    "\n",
    "    initialized_variable_names = {}\n",
    "    (assignment_map, initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)\n",
    "    tf.train.init_from_checkpoint(init_checkpoint, assignment_map)\n",
    "    output_spec = None\n",
    "    predictions = {\"unique_ids\": unique_ids,\n",
    "                   \"start_logits\": start_logits,\n",
    "                   \"end_logits\": end_logits}\n",
    "    output_spec = tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)\n",
    "    return output_spec\n",
    "\n",
    "config = tf.ConfigProto(log_device_placement=True) \n",
    "\n",
    "run_config = tf.estimator.RunConfig(\n",
    "      model_dir=None,\n",
    "      session_config=config,\n",
    "      save_checkpoints_steps=1000,\n",
    "      keep_checkpoint_max=1)\n",
    "\n",
    "estimator = tf.estimator.Estimator(\n",
    "  model_fn=model_fn,\n",
    "  config=run_config,\n",
    "  params=params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xSKkf4JLw62E"
   },
   "source": [
    "### 5.a Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "3OKhc349w62F",
    "outputId": "07d056d2-6d55-4922-88b3-bd25802f70fd",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "eval_examples = run_squad.read_squad_examples(\n",
    "        input_file=input_file, is_training=False)\n",
    "\n",
    "eval_writer = run_squad.FeatureWriter(\n",
    "    filename=os.path.join(output_dir, \"eval.tf_record\"),\n",
    "    is_training=False)\n",
    "\n",
    "eval_features = []\n",
    "def append_feature(feature):\n",
    "    eval_features.append(feature)\n",
    "    eval_writer.process_feature(feature)\n",
    "\n",
    "\n",
    "# Loads a data file into a list of InputBatch's\n",
    "run_squad.convert_examples_to_features(\n",
    "    examples=eval_examples,\n",
    "    tokenizer=tokenizer,\n",
    "    max_seq_length=max_seq_length,\n",
    "    doc_stride=doc_stride,\n",
    "    max_query_length=max_query_length,\n",
    "    is_training=False,\n",
    "    output_fn=append_feature)\n",
    "\n",
    "eval_writer.close()\n",
    "\n",
    "tf.logging.info(\"***** Running predictions *****\")\n",
    "tf.logging.info(\"  Num orig examples = %d\", len(eval_examples))\n",
    "tf.logging.info(\"  Num split examples = %d\", len(eval_features))\n",
    "tf.logging.info(\"  Batch size = %d\", predict_batch_size)\n",
    "\n",
    "predict_input_fn = run_squad.input_fn_builder(\n",
    "    input_file=eval_writer.filename,\n",
    "    batch_size=predict_batch_size,\n",
    "    seq_length=max_seq_length,\n",
    "    is_training=False,\n",
    "    drop_remainder=False)\n",
    "\n",
    "all_results = []\n",
    "eval_hooks = [run_squad.LogEvalRunHook(predict_batch_size)]\n",
    "eval_start_time = time.time()\n",
    "for result in estimator.predict(\n",
    "        predict_input_fn, yield_single_examples=True, hooks=eval_hooks, checkpoint_path=init_checkpoint):\n",
    "    unique_id = int(result[\"unique_ids\"])\n",
    "    start_logits = [float(x) for x in result[\"start_logits\"].flat]\n",
    "    end_logits = [float(x) for x in result[\"end_logits\"].flat]\n",
    "    all_results.append(\n",
    "      run_squad.RawResult(\n",
    "          unique_id=unique_id,\n",
    "          start_logits=start_logits,\n",
    "          end_logits=end_logits))\n",
    "\n",
    "eval_time_elapsed = time.time() - eval_start_time\n",
    "\n",
    "time_list = eval_hooks[-1].time_list\n",
    "time_list.sort()\n",
    "eval_time_wo_startup = sum(time_list[:int(len(time_list) * 0.99)])\n",
    "num_sentences = eval_hooks[-1].count * predict_batch_size\n",
    "avg_sentences_per_second = num_sentences * 1.0 / eval_time_wo_startup\n",
    "\n",
    "tf.logging.info(\"-----------------------------\")\n",
    "tf.logging.info(\"Total Inference Time = %0.2f Inference Time W/O start up overhead = %0.2f \"\n",
    "                \"Sentences processed = %d\", eval_time_elapsed, eval_time_wo_startup,\n",
    "                num_sentences)\n",
    "tf.logging.info(\"Inference Performance = %0.4f sentences/sec\", avg_sentences_per_second)\n",
    "tf.logging.info(\"-----------------------------\")\n",
    "\n",
    "output_prediction_file = os.path.join(output_dir, \"predictions.json\")\n",
    "output_nbest_file = os.path.join(output_dir, \"nbest_predictions.json\")\n",
    "output_null_log_odds_file = os.path.join(output_dir, \"null_odds.json\")\n",
    "\n",
    "run_squad.write_predictions(eval_examples, eval_features, all_results,\n",
    "                  n_best_size, max_answer_length,\n",
    "                  do_lower_case, output_prediction_file,\n",
    "                  output_nbest_file, output_null_log_odds_file,\n",
    "                  version_2_with_negative, verbose_logging)\n",
    "\n",
    "tf.logging.info(\"Inference Results:\")\n",
    "\n",
    "# Here we show only the prediction results, nbest prediction is also available in the output directory\n",
    "results = \"\"\n",
    "with open(output_prediction_file, 'r') as json_file:\n",
    "    data = json.load(json_file)\n",
    "    for question in eval_examples:\n",
    "        results += \"<tr><td>{}</td><td>{}</td><td>{}</td></tr>\".format(question.qas_id, question.question_text, data[question.qas_id])\n",
    "\n",
    "\n",
    "from IPython.display import display, HTML\n",
    "display(HTML(\"<table><tr><th>Id</th><th>Question</th><th>Answer</th></tr>{}</table>\".format(results)))        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EMT0sKxHw62L"
   },
   "source": [
    "## 6. What's next"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mKBM_UD6w62N"
   },
   "source": [
    "Now that you are familiar with running QA Inference on BERT, using mixed precision, you may want to try\n",
    "your own paragraphs and queries. \n",
    "\n",
    "You may also want to take a look to the notebook __bert_squad_tf_finetuning.ipynb__ on how to run fine-tuning on BERT, available in the same directory."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "Copy of bert_squad_tf_inference.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}