demo/BERT/notebooks/Q-and-A.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ceramic-encoding",
   "metadata": {},
   "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",
    "# ==============================================================================\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "legendary-prairie",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# TensorRT: Q&A with BERT\n",
    "\n",
    "This notebook provides a playground for testing various BERT QA models on the CPU and GPU.\n",
    "\n",
    "For \"CPU - Framework (PyTorch)\" and \"GPU - Framework (PyTorch)\", a SpanBERT large model from HuggingFace model repository is employed. Inference is carried out with PyTorch in FP32 precision. All models run with batch size 1.\n",
    "Average run time across 10 runs is reported.\n",
    "\n",
    "**Notes**: \n",
    " - Prior to running this notebook, run [BERT-TRT-FP16.ipynb](BERT-TRT-FP16.ipynb) and [BERT-TRT-INT8-QAT-sparse.ipynb](BERT-TRT-INT8-QAT-sparse.ipynb) to generate the TensorRT engines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "outdoor-indonesia",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import sys\n",
    "sys.path.append('/workspace/TensorRT/demo/BERT')\n",
    "\n",
    "import ipywidgets as widgets\n",
    "import tensorrt as trt;\n",
    "TRT_VERSION = trt.__version__\n",
    "print(\"TensorRT version: \", TRT_VERSION)\n",
    "\n",
    "import time\n",
    "import json\n",
    "import ctypes\n",
    "import argparse\n",
    "import collections\n",
    "import numpy as np\n",
    "import tensorrt as trt\n",
    "import pycuda.driver as cuda\n",
    "import pycuda.autoinit\n",
    "\n",
    "from helpers import tokenization as tokenization\n",
    "from helpers import data_processing as dp\n",
    "\n",
    "TRT_LOGGER = trt.Logger(trt.Logger.INFO)\n",
    "\n",
    "################################################## PyTorch inference #######################################################\n",
    "# Install a customized version of HuggingFace, adding DL model inference timing\n",
    "#!pip3 install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio===0.8.1 -f https://download.pytorch.org/whl/torch_stable.html\n",
    "#!rm -rf /tmp/transformers\n",
    "#!cd /tmp && git clone https://github.com/vinhngx/transformers && cd transformers && pip install .\n",
    "# SpanBERT large model (340M params): https://github.com/facebookresearch/SpanBERT#finetuned-models-squad-1120-relation-extraction-coreference-resolution\n",
    "from transformers import BertForQuestionAnswering, AutoTokenizer\n",
    "\n",
    "#modelname = 'deepset/bert-base-cased-squad2'\n",
    "modelname = 'mrm8488/spanbert-large-finetuned-squadv2'\n",
    "model = BertForQuestionAnswering.from_pretrained(modelname)\n",
    "\n",
    "from transformers import pipeline\n",
    "nlp = pipeline('question-answering', model=model, tokenizer=\"SpanBERT/spanbert-large-cased\")\n",
    "\n",
    "\n",
    "model_gpu = BertForQuestionAnswering.from_pretrained(modelname).cuda()\n",
    "nlp_gpu = pipeline('question-answering', model=model_gpu, tokenizer=\"SpanBERT/spanbert-large-cased\", device=0)\n",
    "\n",
    "################################################## TensorRT inference #######################################################\n",
    "def inference_FP16(trt_context, d_inputs, h_output, d_output, features, tokens):\n",
    "    #global h_output\n",
    "    context = trt_context\n",
    "    \n",
    "    _NetworkOutput = collections.namedtuple(  # pylint: disable=invalid-name\n",
    "            \"NetworkOutput\",\n",
    "            [\"start_logits\", \"end_logits\", \"feature_index\"])\n",
    "    networkOutputs = []\n",
    "\n",
    "    eval_time_elapsed = 0\n",
    "    for feature_index, feature in enumerate(features):\n",
    "        # Copy inputs\n",
    "        input_ids_batch = np.repeat(np.expand_dims(feature.input_ids, 0), 1, axis=0)\n",
    "        segment_ids_batch = np.repeat(np.expand_dims(feature.segment_ids, 0), 1, axis=0)\n",
    "        input_mask_batch = np.repeat(np.expand_dims(feature.input_mask, 0), 1, axis=0)\n",
    "\n",
    "        input_ids = cuda.register_host_memory(np.ascontiguousarray(input_ids_batch.ravel()))\n",
    "        segment_ids = cuda.register_host_memory(np.ascontiguousarray(segment_ids_batch.ravel()))\n",
    "        input_mask = cuda.register_host_memory(np.ascontiguousarray(input_mask_batch.ravel()))\n",
    "\n",
    "        eval_start_time = time.time()\n",
    "        cuda.memcpy_htod_async(d_inputs[0], input_ids, stream)\n",
    "        cuda.memcpy_htod_async(d_inputs[1], segment_ids, stream)\n",
    "        cuda.memcpy_htod_async(d_inputs[2], input_mask, stream)\n",
    "\n",
    "        # Setup tensor address\n",
    "        bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n",
    "\n",
    "        for i in range(engine.num_io_tensors):\n",
    "            context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n",
    "\n",
    "        # Run inference\n",
    "        trt_context.execute_async_v3(stream_handle=stream.handle)\n",
    "        # Synchronize the stream\n",
    "        stream.synchronize()\n",
    "        eval_time_elapsed += (time.time() - eval_start_time)\n",
    "\n",
    "        # Transfer predictions back from GPU\n",
    "        cuda.memcpy_dtoh_async(h_output, d_output, stream)\n",
    "        stream.synchronize()\n",
    "\n",
    "        for index, batch in enumerate(h_output):\n",
    "            # Data Post-processing\n",
    "            networkOutputs.append(_NetworkOutput(\n",
    "                start_logits = np.array(batch.squeeze()[:, 0]),\n",
    "                end_logits = np.array(batch.squeeze()[:, 1]),\n",
    "                feature_index = feature_index\n",
    "                ))\n",
    "\n",
    "    eval_time_elapsed /= len(features)\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. This is needed \n",
    "    #  because the start and end predictions are not conditioned on one another\n",
    "    max_answer_length = 30\n",
    "\n",
    "    prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features,\n",
    "        networkOutputs, n_best_size, max_answer_length)\n",
    "\n",
    "    return eval_time_elapsed, prediction, nbest_json\n",
    "\n",
    "def inference_INT8(trt_context, d_inputs, h_output, d_output, features, tokens):\n",
    "    #global h_output\n",
    "    context = trt_context\n",
    "    \n",
    "    _NetworkOutput = collections.namedtuple(  # pylint: disable=invalid-name\n",
    "            \"NetworkOutput\",\n",
    "            [\"start_logits\", \"end_logits\", \"feature_index\"])\n",
    "    networkOutputs = []\n",
    "\n",
    "    eval_time_elapsed = 0\n",
    "    for feature_index, feature in enumerate(features):\n",
    "        # Copy inputs\n",
    "        B = 1\n",
    "        S = np.sum(feature.input_mask)\n",
    "        input_ids = feature.input_ids[0:S]\n",
    "        segment_ids = feature.segment_ids[0:S]\n",
    "        cu_seq_lens = np.array([0, S], dtype=np.int32)\n",
    "\n",
    "        first_tensor_name = engine.get_tensor_name(0)\n",
    "        second_tensor_name = engine.get_tensor_name(1)\n",
    "        third_tensor_name = engine.get_tensor_name(2)\n",
    "        forth_tensor_name = engine.get_tensor_name(3)\n",
    "\n",
    "        if context.get_tensor_shape(first_tensor_name)[0] != S:\n",
    "            context.set_input_shape(first_tensor_name, (S,))\n",
    "        if context.get_tensor_shape(second_tensor_name)[0] != S:\n",
    "            context.set_input_shape(second_tensor_name, (S,))\n",
    "        if context.get_tensor_shape(third_tensor_name)[0] != 2:\n",
    "            context.set_input_shape(third_tensor_name, (2,))\n",
    "        if context.get_tensor_shape(forth_tensor_name)[0] != S:\n",
    "            context.set_input_shapee(forth_tensor_name, (S,))\n",
    "\n",
    "        h_input_ids = cuda.register_host_memory(np.ascontiguousarray(input_ids.ravel()))\n",
    "        h_segment_ids = cuda.register_host_memory(np.ascontiguousarray(segment_ids.ravel()))\n",
    "        h_cu_seq_lens = cuda.register_host_memory(np.ascontiguousarray(cu_seq_lens.ravel()))\n",
    "\n",
    "        eval_start_time = time.time()\n",
    "        cuda.memcpy_htod_async(d_inputs[0], h_input_ids, INT8_stream)\n",
    "        cuda.memcpy_htod_async(d_inputs[1], h_segment_ids, INT8_stream)\n",
    "        cuda.memcpy_htod_async(d_inputs[2], h_cu_seq_lens, INT8_stream)\n",
    "\n",
    "        # Setup tensor address\n",
    "        bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n",
    "\n",
    "        for i in range(engine.num_io_tensors):\n",
    "            context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n",
    "\n",
    "        # Run inference\n",
    "        trt_context.execute_async_v3(stream_handle=INT8_stream.handle)\n",
    "        # Synchronize the stream\n",
    "        INT8_stream.synchronize()\n",
    "        eval_time_elapsed += (time.time() - eval_start_time)\n",
    "\n",
    "        # Transfer predictions back from GPU\n",
    "        cuda.memcpy_dtoh_async(h_output, d_output, INT8_stream)\n",
    "        INT8_stream.synchronize()\n",
    "\n",
    "        # Only retrieve and post-process the first batch\n",
    "        networkOutputs.append(_NetworkOutput(\n",
    "            start_logits = np.array(h_output[0:S]),\n",
    "            end_logits = np.array(h_output[S:S*2]),\n",
    "            feature_index = feature_index\n",
    "            ))\n",
    "\n",
    "    eval_time_elapsed /= len(features)\n",
    "\n",
    "    # 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. This is needed\n",
    "    # because the start and end predictions are not conditioned on one another\n",
    "    max_answer_length = 30\n",
    "\n",
    "    prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features,\n",
    "            networkOutputs, n_best_size, max_answer_length)\n",
    "\n",
    "    return eval_time_elapsed, prediction, nbest_json\n",
    "        \n",
    "def print_single_query(eval_time_elapsed, prediction, nbest_json):\n",
    "    print(\"Answer: '{}'\".format(prediction))\n",
    "    print(\"With probability: {:.2f}%\".format(nbest_json[0]['probability'] * 100.0))\n",
    "    \n",
    "def question_features(tokens, question):\n",
    "    # Extract features from the paragraph and question\n",
    "    return dp.convert_example_to_features(tokens, question, tokenizer, max_seq_length, doc_stride, max_query_length)\n",
    "\n",
    "doc_stride = 128\n",
    "max_query_length = 64\n",
    "\n",
    "vocab_file = \"models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt\"\n",
    "tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file, do_lower_case=True)\n",
    "\n",
    "### FP16 TRT model\n",
    "engine_path = \"engines_{}/bert_large_384.engine\".format(TRT_VERSION)\n",
    "max_seq_length = 384\n",
    "batch_size = 1\n",
    "\n",
    "runtime = trt.Runtime(TRT_LOGGER)\n",
    "engine = runtime.deserialize_cuda_engine(open(engine_path, 'rb') .read()) \n",
    "context = engine.create_execution_context()\n",
    "\n",
    " # We always use batch size 1.\n",
    "input_shape = (1, max_seq_length)\n",
    "input_nbytes = trt.volume(input_shape) * trt.int32.itemsize\n",
    "\n",
    "# Allocate device memory for inputs.\n",
    "d_inputs = [cuda.mem_alloc(input_nbytes) for binding in range(3)]\n",
    "\n",
    "# Specify input shapes. These must be within the min/max bounds of the active profile (0th profile in this case)\n",
    "# Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape.\n",
    "for binding in range(3):\n",
    "    tensor_name = engine.get_tensor_name(binding)\n",
    "    context.set_input_shape(tensor_name, input_shape)\n",
    "assert context.all_binding_shapes_specified\n",
    "\n",
    "# Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n",
    "h_output = cuda.pagelocked_empty(tuple(context.get_tensor_shape(engine.get_tensor_name(3))), dtype=np.float32)\n",
    "d_output = cuda.mem_alloc(h_output.nbytes)\n",
    "\n",
    "# Create a stream in which to copy inputs/outputs and run inference.\n",
    "stream = cuda.Stream()\n",
    "\n",
    "### INT8 TRT model\n",
    "engine_path = \"engines_%s/megatron_large_seqlen384_int8qat_sparse.engine\"%TRT_VERSION\n",
    "max_seq_length = 384\n",
    "\n",
    "INT8_runtime = trt.Runtime(TRT_LOGGER)\n",
    "INT8_engine = INT8_runtime.deserialize_cuda_engine(open(engine_path, 'rb') .read()) \n",
    "INT8_context = INT8_engine.create_execution_context()\n",
    "\n",
    "# select engine profile\n",
    "INT8_context.set_optimization_profile_async(0, stream.handle)\n",
    "\n",
    "input_nbytes = max_seq_length * trt.int32.itemsize\n",
    "\n",
    "# Allocate device memory for inputs.\n",
    "INT8_d_inputs = [cuda.mem_alloc(input_nbytes) for binding in range(4)]\n",
    "\n",
    "# Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n",
    "INT8_h_output = cuda.pagelocked_empty((2 * max_seq_length), dtype=np.float32)\n",
    "INT8_d_output = cuda.mem_alloc(INT8_h_output.nbytes)\n",
    "\n",
    "# Create a stream in which to copy inputs/outputs and run inference.\n",
    "INT8_stream = cuda.Stream()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pressed-adobe",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "device = widgets.RadioButtons(\n",
    "    options=['CPU - Framework (PyTorch)', \n",
    "             'GPU - Framework (PyTorch)', \n",
    "             'GPU - TensorRT FP16',\n",
    "             'GPU - TensorRT INT8'],\n",
    "    description='Device:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "paragraph_text = widgets.Textarea(\n",
    "    value='TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps'\\\n",
    "'such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops'\\\n",
    "'and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep'\\\n",
    "'learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps.',\n",
    "    placeholder='Type something',\n",
    "    description='Passage:',\n",
    "    disabled=False,\n",
    "    layout=widgets.Layout(width=\"auto\"),\n",
    "    rows=10,  \n",
    ")\n",
    "\n",
    "question_text = widgets.Textarea(\n",
    "    value='What is TensorRT?',\n",
    "    placeholder='Type something',\n",
    "    description='Question:',\n",
    "    disabled=False,\n",
    "    layout=widgets.Layout(width=\"auto\"),\n",
    "    rows=2,\n",
    ")\n",
    "display(paragraph_text)\n",
    "display(question_text)\n",
    "\n",
    "from IPython.display import display\n",
    "box_layout = widgets.Layout(display='flex',\n",
    "                flex_flow='column',\n",
    "                align_items='center',\n",
    "                width='100%')\n",
    "\n",
    "button = widgets.Button(description=\"Answer Me!\")\n",
    "output = widgets.Output()\n",
    "box = widgets.HBox(children=[button],layout=box_layout)\n",
    "\n",
    "N_RUN = 10\n",
    "\n",
    "def answer(b):\n",
    "    progress_bar.value = 0\n",
    "    inference_time_arr = []\n",
    "    with output:\n",
    "        if device.value == 'GPU - TensorRT FP16':\n",
    "            output.clear_output()\n",
    "            for _ in range(N_RUN):\n",
    "                doc_tokens = dp.convert_doc_tokens(paragraph_text.value)\n",
    "                features = question_features(doc_tokens, question_text.value)\n",
    "                eval_time_elapsed, prediction, nbest_json = inference_FP16(context, d_inputs, h_output, d_output, features, doc_tokens)\n",
    "                progress_bar.value += 1                \n",
    "                inference_time_arr.append(eval_time_elapsed)\n",
    "\n",
    "            print_single_query(eval_time_elapsed, prediction, nbest_json)\n",
    "            print(\"Average inference time (over {} runs): {:.2f} ms\".format(N_RUN, 1000*np.mean(inference_time_arr)))   \n",
    "        elif device.value == 'GPU - TensorRT INT8':\n",
    "            output.clear_output()\n",
    "            for _ in range(N_RUN):\n",
    "                doc_tokens = dp.convert_doc_tokens(paragraph_text.value)\n",
    "                features = question_features(doc_tokens, question_text.value)\n",
    "                eval_time_elapsed, prediction, nbest_json = inference_INT8(INT8_context, INT8_d_inputs, INT8_h_output, INT8_d_output, features, doc_tokens)\n",
    "                progress_bar.value += 1                \n",
    "                inference_time_arr.append(eval_time_elapsed)\n",
    "\n",
    "            print_single_query(eval_time_elapsed, prediction, nbest_json)\n",
    "            print(\"Average inference time (over {} runs): {:.2f} ms\".format(N_RUN, 1000*np.mean(inference_time_arr)))   \n",
    " \n",
    "        elif device.value == 'CPU - Framework (PyTorch)':\n",
    "            output.clear_output()\n",
    "            for _ in range(N_RUN):\n",
    "                inference_time = time.time()\n",
    "                answer = nlp({\n",
    "                        'question': question_text.value,\n",
    "                        'context': paragraph_text.value\n",
    "                        })\n",
    "                progress_bar.value += 1                \n",
    "                inference_time_arr.append(time.time() - inference_time)\n",
    "                \n",
    "            print(\"Answer: '{}'\".format(answer['answer']))\n",
    "            print(\"With probability: {:.2f}%\".format(answer['score']*100))\n",
    "            print(\"Average inference time (over {} runs): {:.2f} ms\".format(N_RUN, 1000*np.mean(inference_time_arr)))   \n",
    "        elif  device.value == 'GPU - Framework (PyTorch)':  \n",
    "            output.clear_output()\n",
    "            for _ in range(N_RUN):\n",
    "                inference_time = time.time()\n",
    "                answer = nlp_gpu({\n",
    "                        'question': question_text.value,\n",
    "                        'context': paragraph_text.value\n",
    "                        })\n",
    "                progress_bar.value += 1                \n",
    "                inference_time_arr.append(time.time() - inference_time)\n",
    "                \n",
    "            print(\"Answer: '{}'\".format(answer['answer']))\n",
    "            print(\"With probability: {:.2f}%\".format(answer['score']*100))\n",
    "            print(\"Average inference time (over {} runs): {:.2f} ms\".format(N_RUN, 1000*np.mean(inference_time_arr)))           \n",
    "            \n",
    "button.on_click(answer)\n",
    "display(device, box, output)\n",
    "\n",
    "progress_bar = widgets.IntProgress(\n",
    "    value=0,\n",
    "    min=0,\n",
    "    max=N_RUN,\n",
    "    description='Progress:',\n",
    "    bar_style='', # 'success', 'info', 'warning', 'danger' or ''\n",
    "    style={'bar_color': 'green'},\n",
    "    orientation='horizontal', \n",
    "    layout=widgets.Layout(width='100%', height='50px')\n",
    ")\n",
    "display(progress_bar)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "musical-right",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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": 5
}