Bert_QA.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aAUq7duyG5J0"
   },
   "source": [
    "# Importing dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9XYlThAmGZg1",
    "outputId": "f48cfeb3-afde-4157-f1a5-822bb18f2b94"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting sentencepiece\n",
      "  Downloading sentencepiece-0.1.96-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.2 MB)\n",
      "\u001b[K     |████████████████████████████████| 1.2 MB 4.2 MB/s eta 0:00:01\n",
      "\u001b[?25hInstalling collected packages: sentencepiece\n",
      "Successfully installed sentencepiece-0.1.96\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n",
      "\u001b[33mWARNING: You are using pip version 21.2.4; however, version 22.0.3 is available.\n",
      "You should consider upgrading via the '/usr/bin/python -m pip install --upgrade pip' command.\u001b[0m\n",
      "Collecting tf-models-official\n",
      "  Downloading tf_models_official-2.8.0-py2.py3-none-any.whl (2.2 MB)\n",
      "\u001b[K     |████████████████████████████████| 2.2 MB 23.8 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: tensorflow-addons in /usr/local/lib/python3.8/dist-packages (from tf-models-official) (0.13.1)\n",
      "Collecting gin-config\n",
      "  Downloading gin_config-0.5.0-py3-none-any.whl (61 kB)\n",
      "\u001b[K     |████████████████████████████████| 61 kB 16.1 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting opencv-python-headless\n",
      "  Downloading opencv_python_headless-4.5.5.62-cp36-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (47.7 MB)\n",
      "\u001b[K     |████████████████████████████████| 47.7 MB 72.2 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting tensorflow-hub>=0.6.0\n",
      "  Downloading tensorflow_hub-0.12.0-py2.py3-none-any.whl (108 kB)\n",
      "\u001b[K     |████████████████████████████████| 108 kB 76.1 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting pandas>=0.22.0\n",
      "  Downloading pandas-1.4.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (11.7 MB)\n",
      "\u001b[K     |████████████████████████████████| 11.7 MB 53.1 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: pyyaml<6.0,>=5.1 in /usr/local/lib/python3.8/dist-packages (from tf-models-official) (5.4.1)\n",
      "Collecting Pillow\n",
      "  Downloading Pillow-9.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.3 MB)\n",
      "\u001b[K     |████████████████████████████████| 4.3 MB 65.7 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting tensorflow-text~=2.8.0\n",
      "  Downloading tensorflow_text-2.8.1-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (4.9 MB)\n",
      "\u001b[K     |████████████████████████████████| 4.9 MB 64.4 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: scipy>=0.19.1 in /usr/local/lib/python3.8/dist-packages (from tf-models-official) (1.4.1)\n",
      "Collecting kaggle>=1.3.9\n",
      "  Downloading kaggle-1.5.12.tar.gz (58 kB)\n",
      "\u001b[K     |████████████████████████████████| 58 kB 12.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting seqeval\n",
      "  Downloading seqeval-1.2.2.tar.gz (43 kB)\n",
      "\u001b[K     |████████████████████████████████| 43 kB 4.0 MB/s  eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: sentencepiece in /usr/local/lib/python3.8/dist-packages (from tf-models-official) (0.1.96)\n",
      "Collecting py-cpuinfo>=3.3.0\n",
      "  Downloading py-cpuinfo-8.0.0.tar.gz (99 kB)\n",
      "\u001b[K     |████████████████████████████████| 99 kB 13.9 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting sacrebleu\n",
      "  Downloading sacrebleu-2.0.0-py3-none-any.whl (90 kB)\n",
      "\u001b[K     |████████████████████████████████| 90 kB 13.1 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting pycocotools\n",
      "  Downloading pycocotools-2.0.4.tar.gz (106 kB)\n",
      "\u001b[K     |████████████████████████████████| 106 kB 51.1 MB/s eta 0:00:01\n",
      "\u001b[?25h  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h    Preparing wheel metadata ... \u001b[?25ldone\n",
      "\u001b[?25hCollecting google-api-python-client>=1.6.7\n",
      "  Downloading google_api_python_client-2.38.0-py2.py3-none-any.whl (8.2 MB)\n",
      "\u001b[K     |████████████████████████████████| 8.2 MB 51.0 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting tf-slim>=1.1.0\n",
      "  Downloading tf_slim-1.1.0-py2.py3-none-any.whl (352 kB)\n",
      "\u001b[K     |████████████████████████████████| 352 kB 80.6 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting matplotlib\n",
      "  Downloading matplotlib-3.5.1-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl (11.3 MB)\n",
      "\u001b[K     |████████████████████████████████| 11.3 MB 69.2 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: numpy>=1.15.4 in /usr/local/lib/python3.8/dist-packages (from tf-models-official) (1.19.4)\n",
      "Collecting Cython\n",
      "  Using cached Cython-0.29.28-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_24_x86_64.whl (1.9 MB)\n",
      "Requirement already satisfied: tensorflow-datasets in /usr/local/lib/python3.8/dist-packages (from tf-models-official) (3.2.1)\n",
      "Requirement already satisfied: six in /usr/local/lib/python3.8/dist-packages (from tf-models-official) (1.15.0)\n",
      "Collecting tensorflow~=2.8.0\n",
      "  Downloading tensorflow-2.8.0-cp38-cp38-manylinux2010_x86_64.whl (497.6 MB)\n",
      "\u001b[K     |████████████████████████████████| 497.6 MB 2.6 kB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting tensorflow-model-optimization>=0.4.1\n",
      "  Downloading tensorflow_model_optimization-0.7.1-py2.py3-none-any.whl (234 kB)\n",
      "\u001b[K     |████████████████████████████████| 234 kB 59.1 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting oauth2client\n",
      "  Downloading oauth2client-4.1.3-py2.py3-none-any.whl (98 kB)\n",
      "\u001b[K     |████████████████████████████████| 98 kB 15.2 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: psutil>=5.4.3 in /usr/local/lib/python3.8/dist-packages (from tf-models-official) (5.7.0)\n",
      "Requirement already satisfied: google-auth<3.0.0dev,>=1.16.0 in /usr/local/lib/python3.8/dist-packages (from google-api-python-client>=1.6.7->tf-models-official) (1.35.0)\n",
      "Collecting google-api-core<3.0.0dev,>=1.21.0\n",
      "  Downloading google_api_core-2.5.0-py2.py3-none-any.whl (111 kB)\n",
      "\u001b[K     |████████████████████████████████| 111 kB 57.6 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting google-auth-httplib2>=0.1.0\n",
      "  Downloading google_auth_httplib2-0.1.0-py2.py3-none-any.whl (9.3 kB)\n",
      "Collecting httplib2<1dev,>=0.15.0\n",
      "  Downloading httplib2-0.20.4-py3-none-any.whl (96 kB)\n",
      "\u001b[K     |████████████████████████████████| 96 kB 10.9 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting uritemplate<5,>=3.0.1\n",
      "  Downloading uritemplate-4.1.1-py2.py3-none-any.whl (10 kB)\n",
      "Requirement already satisfied: requests<3.0.0dev,>=2.18.0 in /usr/local/lib/python3.8/dist-packages (from google-api-core<3.0.0dev,>=1.21.0->google-api-python-client>=1.6.7->tf-models-official) (2.26.0)\n",
      "Requirement already satisfied: protobuf>=3.12.0 in /usr/local/lib/python3.8/dist-packages (from google-api-core<3.0.0dev,>=1.21.0->google-api-python-client>=1.6.7->tf-models-official) (3.17.3)\n",
      "Requirement already satisfied: googleapis-common-protos<2.0dev,>=1.52.0 in /usr/local/lib/python3.8/dist-packages (from google-api-core<3.0.0dev,>=1.21.0->google-api-python-client>=1.6.7->tf-models-official) (1.53.0)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.8/dist-packages (from google-auth<3.0.0dev,>=1.16.0->google-api-python-client>=1.6.7->tf-models-official) (4.7.2)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.8/dist-packages (from google-auth<3.0.0dev,>=1.16.0->google-api-python-client>=1.6.7->tf-models-official) (0.2.8)\n",
      "Requirement already satisfied: setuptools>=40.3.0 in /usr/local/lib/python3.8/dist-packages (from google-auth<3.0.0dev,>=1.16.0->google-api-python-client>=1.6.7->tf-models-official) (58.1.0)\n",
      "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from google-auth<3.0.0dev,>=1.16.0->google-api-python-client>=1.6.7->tf-models-official) (4.2.2)\n",
      "Requirement already satisfied: pyparsing!=3.0.0,!=3.0.1,!=3.0.2,!=3.0.3,<4,>=2.4.2 in /usr/local/lib/python3.8/dist-packages (from httplib2<1dev,>=0.15.0->google-api-python-client>=1.6.7->tf-models-official) (2.4.7)\n",
      "Requirement already satisfied: certifi in /usr/local/lib/python3.8/dist-packages (from kaggle>=1.3.9->tf-models-official) (2021.5.30)\n",
      "Requirement already satisfied: python-dateutil in /usr/local/lib/python3.8/dist-packages (from kaggle>=1.3.9->tf-models-official) (2.8.2)\n",
      "Requirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from kaggle>=1.3.9->tf-models-official) (4.62.3)\n",
      "Collecting python-slugify\n",
      "  Downloading python_slugify-6.1.1-py2.py3-none-any.whl (9.1 kB)\n",
      "Requirement already satisfied: urllib3 in /usr/local/lib/python3.8/dist-packages (from kaggle>=1.3.9->tf-models-official) (1.26.7)\n",
      "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.8/dist-packages (from pandas>=0.22.0->tf-models-official) (2021.1)\n",
      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.8/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3.0.0dev,>=1.16.0->google-api-python-client>=1.6.7->tf-models-official) (0.4.8)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.8/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core<3.0.0dev,>=1.21.0->google-api-python-client>=1.6.7->tf-models-official) (2.0.6)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core<3.0.0dev,>=1.21.0->google-api-python-client>=1.6.7->tf-models-official) (3.2)\n",
      "Requirement already satisfied: flatbuffers>=1.12 in /usr/local/lib/python3.8/dist-packages (from tensorflow~=2.8.0->tf-models-official) (1.12)\n",
      "Collecting keras<2.9,>=2.8.0rc0\n",
      "  Downloading keras-2.8.0-py2.py3-none-any.whl (1.4 MB)\n",
      "\u001b[K     |████████████████████████████████| 1.4 MB 72.9 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.8/dist-packages (from tensorflow~=2.8.0->tf-models-official) (1.39.0)\n",
      "Collecting tf-estimator-nightly==2.8.0.dev2021122109\n",
      "  Downloading tf_estimator_nightly-2.8.0.dev2021122109-py2.py3-none-any.whl (462 kB)\n",
      "\u001b[K     |████████████████████████████████| 462 kB 70.1 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.8/dist-packages (from tensorflow~=2.8.0->tf-models-official) (3.3.0)\n",
      "Requirement already satisfied: keras-preprocessing>=1.1.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow~=2.8.0->tf-models-official) (1.1.2)\n",
      "Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow~=2.8.0->tf-models-official) (0.2.0)\n",
      "Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow~=2.8.0->tf-models-official) (1.6.3)\n",
      "Collecting numpy>=1.15.4\n",
      "  Downloading numpy-1.22.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (16.8 MB)\n",
      "\u001b[K     |████████████████████████████████| 16.8 MB 37.3 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting tensorboard<2.9,>=2.8\n",
      "  Downloading tensorboard-2.8.0-py3-none-any.whl (5.8 MB)\n",
      "\u001b[K     |████████████████████████████████| 5.8 MB 75.0 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow~=2.8.0->tf-models-official) (3.1.0)\n",
      "Requirement already satisfied: typing-extensions>=3.6.6 in /usr/local/lib/python3.8/dist-packages (from tensorflow~=2.8.0->tf-models-official) (3.7.4.3)\n",
      "Requirement already satisfied: gast>=0.2.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow~=2.8.0->tf-models-official) (0.4.0)\n",
      "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow~=2.8.0->tf-models-official) (1.1.0)\n",
      "Requirement already satisfied: wrapt>=1.11.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow~=2.8.0->tf-models-official) (1.12.1)\n",
      "Collecting tensorflow-io-gcs-filesystem>=0.23.1\n",
      "  Downloading tensorflow_io_gcs_filesystem-0.24.0-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (2.1 MB)\n",
      "\u001b[K     |████████████████████████████████| 2.1 MB 70.6 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: absl-py>=0.4.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow~=2.8.0->tf-models-official) (0.12.0)\n",
      "Collecting libclang>=9.0.1\n",
      "  Downloading libclang-13.0.0-py2.py3-none-manylinux1_x86_64.whl (14.5 MB)\n",
      "\u001b[K     |████████████████████████████████| 14.5 MB 58.1 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.8/dist-packages (from astunparse>=1.6.0->tensorflow~=2.8.0->tf-models-official) (0.37.0)\n",
      "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.9,>=2.8->tensorflow~=2.8.0->tf-models-official) (1.8.0)\n",
      "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.9,>=2.8->tensorflow~=2.8.0->tf-models-official) (2.0.1)\n",
      "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.9,>=2.8->tensorflow~=2.8.0->tf-models-official) (3.3.4)\n",
      "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.9,>=2.8->tensorflow~=2.8.0->tf-models-official) (0.6.1)\n",
      "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.9,>=2.8->tensorflow~=2.8.0->tf-models-official) (0.4.6)\n",
      "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.8/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.9,>=2.8->tensorflow~=2.8.0->tf-models-official) (1.3.0)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.8/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.9,>=2.8->tensorflow~=2.8.0->tf-models-official) (3.1.1)\n",
      "Collecting dm-tree~=0.1.1\n",
      "  Downloading dm_tree-0.1.6-cp38-cp38-manylinux_2_24_x86_64.whl (94 kB)\n",
      "\u001b[K     |████████████████████████████████| 94 kB 8.2 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting kiwisolver>=1.0.1\n",
      "  Downloading kiwisolver-1.3.2-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl (1.2 MB)\n",
      "\u001b[K     |████████████████████████████████| 1.2 MB 67.3 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting fonttools>=4.22.0\n",
      "  Downloading fonttools-4.29.1-py3-none-any.whl (895 kB)\n",
      "\u001b[K     |████████████████████████████████| 895 kB 44.0 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting cycler>=0.10\n",
      "  Downloading cycler-0.11.0-py3-none-any.whl (6.4 kB)\n",
      "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.8/dist-packages (from matplotlib->tf-models-official) (21.0)\n",
      "Collecting text-unidecode>=1.3\n",
      "  Downloading text_unidecode-1.3-py2.py3-none-any.whl (78 kB)\n",
      "\u001b[K     |████████████████████████████████| 78 kB 13.3 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting colorama\n",
      "  Downloading colorama-0.4.4-py2.py3-none-any.whl (16 kB)\n",
      "Collecting tabulate>=0.8.9\n",
      "  Downloading tabulate-0.8.9-py3-none-any.whl (25 kB)\n",
      "Collecting regex\n",
      "  Downloading regex-2022.1.18-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (764 kB)\n",
      "\u001b[K     |████████████████████████████████| 764 kB 68.1 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting portalocker\n",
      "  Downloading portalocker-2.4.0-py2.py3-none-any.whl (16 kB)\n",
      "Collecting scikit-learn>=0.21.3\n",
      "  Downloading scikit_learn-1.0.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (26.7 MB)\n",
      "\u001b[K     |████████████████████████████████| 26.7 MB 20.2 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting joblib>=0.11\n",
      "  Downloading joblib-1.1.0-py2.py3-none-any.whl (306 kB)\n",
      "\u001b[K     |████████████████████████████████| 306 kB 45.8 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting threadpoolctl>=2.0.0\n",
      "  Downloading threadpoolctl-3.1.0-py3-none-any.whl (14 kB)\n",
      "Requirement already satisfied: typeguard>=2.7 in /usr/local/lib/python3.8/dist-packages (from tensorflow-addons->tf-models-official) (2.12.1)\n",
      "Requirement already satisfied: attrs>=18.1.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow-datasets->tf-models-official) (21.2.0)\n",
      "Requirement already satisfied: future in /usr/local/lib/python3.8/dist-packages (from tensorflow-datasets->tf-models-official) (0.18.2)\n",
      "Requirement already satisfied: promise in /usr/local/lib/python3.8/dist-packages (from tensorflow-datasets->tf-models-official) (2.3)\n",
      "Requirement already satisfied: tensorflow-metadata in /usr/local/lib/python3.8/dist-packages (from tensorflow-datasets->tf-models-official) (1.2.0)\n",
      "Requirement already satisfied: dill in /usr/local/lib/python3.8/dist-packages (from tensorflow-datasets->tf-models-official) (0.3.4)\n",
      "Building wheels for collected packages: kaggle, py-cpuinfo, pycocotools, seqeval\n",
      "  Building wheel for kaggle (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for kaggle: filename=kaggle-1.5.12-py3-none-any.whl size=73051 sha256=213732b4b7dada33aa2246c7c8ff4e33d472e05cf993d4758743dc0684c1e4dd\n",
      "  Stored in directory: /root/.cache/pip/wheels/29/da/11/144cc25aebdaeb4931b231e25fd34b394e6a5725cbb2f50106\n",
      "  Building wheel for py-cpuinfo (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for py-cpuinfo: filename=py_cpuinfo-8.0.0-py3-none-any.whl size=22258 sha256=ec47cf8d19456680d2fc5de447035cafcead304fcac0689b9b67954e0dea331c\n",
      "  Stored in directory: /root/.cache/pip/wheels/57/cb/6d/bab2257f26c5be4a96ff65c3d2a7122c96529b73773ee37f36\n",
      "  Building wheel for pycocotools (PEP 517) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for pycocotools: filename=pycocotools-2.0.4-cp38-cp38-linux_x86_64.whl size=418804 sha256=fe5fa84fd4dabb8d35d79c706b566f2b271e584ff210ea51de45bbdf3baee2e0\n",
      "  Stored in directory: /root/.cache/pip/wheels/dd/e2/43/3e93cd653b3346b3d702bb0509bc611189f95d60407bff1484\n",
      "  Building wheel for seqeval (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for seqeval: filename=seqeval-1.2.2-py3-none-any.whl size=16181 sha256=c62efd8b98837bde502ac29c0a3ecdb44162105232a5d580c4388338a67663ad\n",
      "  Stored in directory: /root/.cache/pip/wheels/ad/5c/ba/05fa33fa5855777b7d686e843ec07452f22a66a138e290e732\n",
      "Successfully built kaggle py-cpuinfo pycocotools seqeval\n",
      "Installing collected packages: numpy, threadpoolctl, tf-estimator-nightly, text-unidecode, tensorflow-io-gcs-filesystem, tensorboard, Pillow, libclang, kiwisolver, keras, joblib, httplib2, fonttools, cycler, uritemplate, tensorflow-hub, tensorflow, tabulate, scikit-learn, regex, python-slugify, portalocker, matplotlib, google-auth-httplib2, google-api-core, dm-tree, colorama, tf-slim, tensorflow-text, tensorflow-model-optimization, seqeval, sacrebleu, pycocotools, py-cpuinfo, pandas, opencv-python-headless, oauth2client, kaggle, google-api-python-client, gin-config, Cython, tf-models-official\n",
      "  Attempting uninstall: numpy\n",
      "    Found existing installation: numpy 1.19.4\n",
      "    Uninstalling numpy-1.19.4:\n",
      "      Successfully uninstalled numpy-1.19.4\n",
      "  Attempting uninstall: tensorboard\n",
      "    Found existing installation: tensorboard 2.6.0\n",
      "    Uninstalling tensorboard-2.6.0:\n",
      "      Successfully uninstalled tensorboard-2.6.0\n",
      "  Attempting uninstall: keras\n",
      "    Found existing installation: keras 2.6.0\n",
      "    Uninstalling keras-2.6.0:\n",
      "      Successfully uninstalled keras-2.6.0\n",
      "  Attempting uninstall: tensorflow\n",
      "    Found existing installation: tensorflow 2.6.0+nv\n",
      "    Uninstalling tensorflow-2.6.0+nv:\n",
      "      Successfully uninstalled tensorflow-2.6.0+nv\n",
      "Successfully installed Cython-0.29.28 Pillow-9.0.1 colorama-0.4.4 cycler-0.11.0 dm-tree-0.1.6 fonttools-4.29.1 gin-config-0.5.0 google-api-core-2.5.0 google-api-python-client-2.38.0 google-auth-httplib2-0.1.0 httplib2-0.20.4 joblib-1.1.0 kaggle-1.5.12 keras-2.8.0 kiwisolver-1.3.2 libclang-13.0.0 matplotlib-3.5.1 numpy-1.22.2 oauth2client-4.1.3 opencv-python-headless-4.5.5.62 pandas-1.4.1 portalocker-2.4.0 py-cpuinfo-8.0.0 pycocotools-2.0.4 python-slugify-6.1.1 regex-2022.1.18 sacrebleu-2.0.0 scikit-learn-1.0.2 seqeval-1.2.2 tabulate-0.8.9 tensorboard-2.8.0 tensorflow-2.8.0 tensorflow-hub-0.12.0 tensorflow-io-gcs-filesystem-0.24.0 tensorflow-model-optimization-0.7.1 tensorflow-text-2.8.1 text-unidecode-1.3 tf-estimator-nightly-2.8.0.dev2021122109 tf-models-official-2.8.0 tf-slim-1.1.0 threadpoolctl-3.1.0 uritemplate-4.1.1\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n",
      "\u001b[33mWARNING: You are using pip version 21.2.4; however, version 22.0.3 is available.\n",
      "You should consider upgrading via the '/usr/bin/python -m pip install --upgrade pip' command.\u001b[0m\n",
      "Collecting tf-nightly\n",
      "  Downloading tf_nightly-2.9.0.dev20220228-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (498.9 MB)\n",
      "\u001b[K     |████████████████████████████████| 498.9 MB 24 kB/s s eta 0:00:01     |▍                               | 6.5 MB 17.8 MB/s eta 0:00:28██████████████▋           | 322.1 MB 59.7 MB/s eta 0:00:03��███████████▋| 492.5 MB 59.2 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: packaging in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (21.0)\n",
      "Collecting tf-estimator-nightly~=2.9.0.dev\n",
      "  Downloading tf_estimator_nightly-2.9.0.dev2022022809-py2.py3-none-any.whl (438 kB)\n",
      "\u001b[K     |████████████████████████████████| 438 kB 45.0 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: setuptools in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (58.1.0)\n",
      "Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (0.24.0)\n",
      "Collecting tb-nightly~=2.9.0.a\n",
      "  Downloading tb_nightly-2.9.0a20220227-py3-none-any.whl (5.8 MB)\n",
      "\u001b[K     |████████████████████████████████| 5.8 MB 23.4 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (3.1.0)\n",
      "Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (1.6.3)\n",
      "Collecting absl-py>=1.0.0\n",
      "  Downloading absl_py-1.0.0-py3-none-any.whl (126 kB)\n",
      "\u001b[K     |████████████████████████████████| 126 kB 47.8 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (1.15.0)\n",
      "Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (0.2.0)\n",
      "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (1.1.0)\n",
      "Requirement already satisfied: flatbuffers>=1.12 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (1.12)\n",
      "Requirement already satisfied: typing-extensions>=3.6.6 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (3.7.4.3)\n",
      "Collecting keras-nightly~=2.9.0.dev\n",
      "  Downloading keras_nightly-2.9.0.dev2022022808-py2.py3-none-any.whl (1.6 MB)\n",
      "\u001b[K     |████████████████████████████████| 1.6 MB 33.6 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: wrapt>=1.11.0 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (1.12.1)\n",
      "Requirement already satisfied: keras-preprocessing>=1.1.1 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (1.1.2)\n",
      "Requirement already satisfied: numpy>=1.20 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (1.22.2)\n",
      "Requirement already satisfied: protobuf>=3.9.2 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (3.17.3)\n",
      "Requirement already satisfied: gast<=0.4.0,>=0.2.1 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (0.4.0)\n",
      "Requirement already satisfied: libclang>=13.0.0 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (13.0.0)\n",
      "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (3.3.0)\n",
      "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.8/dist-packages (from tf-nightly) (1.39.0)\n",
      "Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.8/dist-packages (from astunparse>=1.6.0->tf-nightly) (0.37.0)\n",
      "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.8/dist-packages (from tb-nightly~=2.9.0.a->tf-nightly) (2.26.0)\n",
      "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.8/dist-packages (from tb-nightly~=2.9.0.a->tf-nightly) (1.35.0)\n",
      "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.8/dist-packages (from tb-nightly~=2.9.0.a->tf-nightly) (0.6.1)\n",
      "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.8/dist-packages (from tb-nightly~=2.9.0.a->tf-nightly) (1.8.0)\n",
      "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.8/dist-packages (from tb-nightly~=2.9.0.a->tf-nightly) (2.0.1)\n",
      "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.8/dist-packages (from tb-nightly~=2.9.0.a->tf-nightly) (3.3.4)\n",
      "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.8/dist-packages (from tb-nightly~=2.9.0.a->tf-nightly) (0.4.6)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tb-nightly~=2.9.0.a->tf-nightly) (4.7.2)\n",
      "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tb-nightly~=2.9.0.a->tf-nightly) (4.2.2)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tb-nightly~=2.9.0.a->tf-nightly) (0.2.8)\n",
      "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.8/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tb-nightly~=2.9.0.a->tf-nightly) (1.3.0)\n",
      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.8/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tb-nightly~=2.9.0.a->tf-nightly) (0.4.8)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tb-nightly~=2.9.0.a->tf-nightly) (1.26.7)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tb-nightly~=2.9.0.a->tf-nightly) (2021.5.30)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tb-nightly~=2.9.0.a->tf-nightly) (3.2)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tb-nightly~=2.9.0.a->tf-nightly) (2.0.6)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.8/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tb-nightly~=2.9.0.a->tf-nightly) (3.1.1)\n",
      "Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.8/dist-packages (from packaging->tf-nightly) (2.4.7)\n",
      "Installing collected packages: absl-py, tf-estimator-nightly, tb-nightly, keras-nightly, tf-nightly\n",
      "  Attempting uninstall: absl-py\n",
      "    Found existing installation: absl-py 0.12.0\n",
      "    Uninstalling absl-py-0.12.0:\n",
      "      Successfully uninstalled absl-py-0.12.0\n",
      "  Attempting uninstall: tf-estimator-nightly\n",
      "    Found existing installation: tf-estimator-nightly 2.8.0.dev2021122109\n",
      "    Uninstalling tf-estimator-nightly-2.8.0.dev2021122109:\n",
      "      Successfully uninstalled tf-estimator-nightly-2.8.0.dev2021122109\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "tensorflow 2.8.0 requires tf-estimator-nightly==2.8.0.dev2021122109, but you have tf-estimator-nightly 2.9.0.dev2022022809 which is incompatible.\n",
      "tensorflow-metadata 1.2.0 requires absl-py<0.13,>=0.9, but you have absl-py 1.0.0 which is incompatible.\u001b[0m\n",
      "Successfully installed absl-py-1.0.0 keras-nightly-2.9.0.dev2022022808 tb-nightly-2.9.0a20220227 tf-estimator-nightly-2.9.0.dev2022022809 tf-nightly-2.9.0.dev20220228\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n",
      "\u001b[33mWARNING: You are using pip version 21.2.4; however, version 22.0.3 is available.\n",
      "You should consider upgrading via the '/usr/bin/python -m pip install --upgrade pip' command.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install sentencepiece\n",
    "!pip install tf-models-official\n",
    "#!pip install tf-models-nightly # better to install the version in development\n",
    "!pip install tf-nightly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "eDdYqvxPHnI8"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "zFS9XSASHvQi",
    "outputId": "a781a3bf-cb7f-4e59-cc09-d7d7ba2d21b1"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.9.0-dev20220228'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Fv38eJzSH00S",
    "outputId": "f4500f8a-5907-431b-ba96-b3fe9f008290"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/tensorflow_addons/utils/ensure_tf_install.py:37: UserWarning: You are currently using a nightly version of TensorFlow (2.9.0-dev20220228). \n",
      "TensorFlow Addons offers no support for the nightly versions of TensorFlow. Some things might work, some other might not. \n",
      "If you encounter a bug, do not file an issue on GitHub.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import tensorflow_hub as hub\n",
    "\n",
    "from official.nlp.bert.tokenization import FullTokenizer\n",
    "from official.nlp.bert.input_pipeline import create_squad_dataset\n",
    "from official.nlp.data.squad_lib import generate_tf_record_from_json_file\n",
    "\n",
    "from official.nlp import optimization\n",
    "\n",
    "from official.nlp.data.squad_lib import read_squad_examples\n",
    "from official.nlp.data.squad_lib import FeatureWriter\n",
    "from official.nlp.data.squad_lib import convert_examples_to_features\n",
    "from official.nlp.data.squad_lib import write_predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "WjzrCZJMJN1N"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import math\n",
    "import random\n",
    "import time\n",
    "import json\n",
    "import collections\n",
    "import os\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iSGw1I_zHAb5"
   },
   "source": [
    "# Data preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "E_PcOeemKcYq"
   },
   "outputs": [],
   "source": [
    "input_meta_data = generate_tf_record_from_json_file(\n",
    "    \"train-v1.1.json\",\n",
    "    \"vocab.txt\",\n",
    "    \"train-v1.1.tf_record\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "7XQGdnANLziM"
   },
   "outputs": [],
   "source": [
    "with tf.io.gfile.GFile(\"train_meta_data\", \"w\") as writer:\n",
    "    writer.write(json.dumps(input_meta_data, indent=4) + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "295nG2zQMYSU"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-02-28 15:59:33.471417: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:969] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-02-28 15:59:33.483021: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:969] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-02-28 15:59:33.483998: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:969] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-02-28 15:59:33.486582: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2022-02-28 15:59:33.487764: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:969] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-02-28 15:59:33.488884: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:969] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-02-28 15:59:33.489869: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:969] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-02-28 15:59:34.799513: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:969] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-02-28 15:59:34.800230: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:969] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-02-28 15:59:34.800774: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:969] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2022-02-28 15:59:34.801404: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 7441 MB memory:  -> device: 0, name: Quadro M4000, pci bus id: 0000:00:05.0, compute capability: 5.2\n"
     ]
    }
   ],
   "source": [
    "BATCH_SIZE = 4\n",
    "\n",
    "train_dataset = create_squad_dataset(\n",
    "    \"train-v1.1.tf_record\",\n",
    "    input_meta_data['max_seq_length'], # 384\n",
    "    BATCH_SIZE,\n",
    "    is_training=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F4939RJqHRUs"
   },
   "source": [
    "# Model building"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lcQcFi8kOc6K"
   },
   "source": [
    "## Squad layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "gjmLeQjiaOo5"
   },
   "outputs": [],
   "source": [
    "class BertSquadLayer(tf.keras.layers.Layer):\n",
    "\n",
    "  def __init__(self):\n",
    "    super(BertSquadLayer, self).__init__()\n",
    "    self.final_dense = tf.keras.layers.Dense(\n",
    "        units=2,\n",
    "        kernel_initializer=tf.keras.initializers.TruncatedNormal(stddev=0.02))\n",
    "\n",
    "  def call(self, inputs):\n",
    "    logits = self.final_dense(inputs) # (batch_size, seq_len, 2)\n",
    "\n",
    "    logits = tf.transpose(logits, [2, 0, 1]) # (2, batch_size, seq_len)\n",
    "    unstacked_logits = tf.unstack(logits, axis=0) # [(batch_size, seq_len), (batch_size, seq_len)] \n",
    "    return unstacked_logits[0], unstacked_logits[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aQbQFKjUOeyf"
   },
   "source": [
    "## Whole model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "KgkIFb1GMT81"
   },
   "outputs": [],
   "source": [
    "class BERTSquad(tf.keras.Model):\n",
    "    \n",
    "    def __init__(self,\n",
    "                 name=\"bert_squad\"):\n",
    "        super(BERTSquad, self).__init__(name=name)\n",
    "        \n",
    "        self.bert_layer = hub.KerasLayer(\n",
    "            \"https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/1\",\n",
    "            trainable=True)\n",
    "        \n",
    "        self.squad_layer = BertSquadLayer()\n",
    "    \n",
    "    def apply_bert(self, inputs):\n",
    "\n",
    "        _ , sequence_output = self.bert_layer([inputs[\"input_word_ids\"],\n",
    "                                               inputs[\"input_mask\"],\n",
    "                                               inputs[\"input_type_ids\"]])\n",
    "        return sequence_output\n",
    "\n",
    "    def call(self, inputs):\n",
    "        seq_output = self.apply_bert(inputs)\n",
    "\n",
    "        start_logits, end_logits = self.squad_layer(seq_output)\n",
    "        \n",
    "        return start_logits, end_logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lBmSU2RnHV_a"
   },
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WnA3WHwRIHAZ"
   },
   "source": [
    "## Creating the AI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "JEUxomxENRoJ"
   },
   "outputs": [],
   "source": [
    "TRAIN_DATA_SIZE = 88641\n",
    "NB_BATCHES_TRAIN = 2000\n",
    "BATCH_SIZE = 4\n",
    "NB_EPOCHS = 3\n",
    "INIT_LR = 5e-5\n",
    "WARMUP_STEPS = int(NB_BATCHES_TRAIN * 0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "5Pg6EKe2daFy"
   },
   "outputs": [],
   "source": [
    "train_dataset_light = train_dataset.take(NB_BATCHES_TRAIN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "eHd5MzTdNIZq"
   },
   "outputs": [],
   "source": [
    "bert_squad = BERTSquad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "G0cvDjBm_KXT"
   },
   "outputs": [],
   "source": [
    "optimizer = optimization.create_optimizer(\n",
    "    init_lr=INIT_LR,\n",
    "    num_train_steps=NB_BATCHES_TRAIN,\n",
    "    num_warmup_steps=WARMUP_STEPS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "I6kG-HpzVK7v"
   },
   "outputs": [],
   "source": [
    "def squad_loss_fn(labels, model_outputs):\n",
    "    start_positions = labels['start_positions']\n",
    "    end_positions = labels['end_positions']\n",
    "    start_logits, end_logits = model_outputs\n",
    "\n",
    "    start_loss = tf.keras.backend.sparse_categorical_crossentropy(\n",
    "        start_positions, start_logits, from_logits=True)\n",
    "    end_loss = tf.keras.backend.sparse_categorical_crossentropy(\n",
    "        end_positions, end_logits, from_logits=True)\n",
    "    \n",
    "    total_loss = (tf.reduce_mean(start_loss) + tf.reduce_mean(end_loss)) / 2\n",
    "\n",
    "    return total_loss\n",
    "\n",
    "train_loss = tf.keras.metrics.Mean(name=\"train_loss\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "HN_C_WA5R_Cb",
    "outputId": "064fede8-db83-4af3-d035-df014cb1810c"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "({'input_word_ids': <tf.Tensor: shape=(4, 384), dtype=int32, numpy=\n",
       "  array([[ 101, 1999, 2297, ...,    0,    0,    0],\n",
       "         [ 101, 2054, 7017, ...,    0,    0,    0],\n",
       "         [ 101, 2129, 2116, ...,    0,    0,    0],\n",
       "         [ 101, 2054, 2003, ...,    0,    0,    0]], dtype=int32)>,\n",
       "  'input_mask': <tf.Tensor: shape=(4, 384), dtype=int32, numpy=\n",
       "  array([[1, 1, 1, ..., 0, 0, 0],\n",
       "         [1, 1, 1, ..., 0, 0, 0],\n",
       "         [1, 1, 1, ..., 0, 0, 0],\n",
       "         [1, 1, 1, ..., 0, 0, 0]], dtype=int32)>,\n",
       "  'input_type_ids': <tf.Tensor: shape=(4, 384), dtype=int32, numpy=\n",
       "  array([[0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0]], dtype=int32)>},\n",
       " {'end_positions': <tf.Tensor: shape=(4,), dtype=int32, numpy=array([ 54,  74,  32, 125], dtype=int32)>,\n",
       "  'start_positions': <tf.Tensor: shape=(4,), dtype=int32, numpy=array([ 53,  71,  29, 118], dtype=int32)>})"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(iter(train_dataset_light))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "V-iE2QFC_KRI"
   },
   "outputs": [],
   "source": [
    "bert_squad.compile(optimizer,\n",
    "                   squad_loss_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "Hjp-_4OyTbuK"
   },
   "outputs": [],
   "source": [
    "checkpoint_path = \"./drive/MyDrive/projects/BERT/ckpt_bert_squad/\"\n",
    "\n",
    "ckpt = tf.train.Checkpoint(bert_squad=bert_squad)\n",
    "\n",
    "ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=1)\n",
    "\n",
    "if ckpt_manager.latest_checkpoint:\n",
    "    ckpt.restore(ckpt_manager.latest_checkpoint)\n",
    "    print(\"Latest checkpoint restored!!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bDgEq09xIOOl"
   },
   "source": [
    "## Custom training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5ywelW3uaSbT",
    "outputId": "0e3b348c-603d-4849-98ff-46eddfc78513"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start of epoch 1\n",
      "Epoch 1 Batch 0 Loss 5.9414\n",
      "Epoch 1 Batch 50 Loss 5.7448\n",
      "Epoch 1 Batch 100 Loss 5.2245\n",
      "Epoch 1 Batch 150 Loss 4.6390\n",
      "Epoch 1 Batch 200 Loss 4.2028\n",
      "Epoch 1 Batch 250 Loss 3.8329\n",
      "Epoch 1 Batch 300 Loss 3.5339\n",
      "Epoch 1 Batch 350 Loss 3.3688\n",
      "Epoch 1 Batch 400 Loss 3.1758\n",
      "Epoch 1 Batch 450 Loss 2.9752\n",
      "Epoch 1 Batch 500 Loss 2.8488\n",
      "Epoch 1 Batch 550 Loss 2.7235\n",
      "Epoch 1 Batch 600 Loss 2.6238\n",
      "Epoch 1 Batch 650 Loss 2.5511\n",
      "Epoch 1 Batch 700 Loss 2.4791\n",
      "Epoch 1 Batch 750 Loss 2.4036\n",
      "Epoch 1 Batch 800 Loss 2.3369\n",
      "Epoch 1 Batch 850 Loss 2.2973\n",
      "Epoch 1 Batch 900 Loss 2.2435\n",
      "Epoch 1 Batch 950 Loss 2.1888\n",
      "Epoch 1 Batch 1000 Loss 2.1274\n",
      "Epoch 1 Batch 1050 Loss 2.0777\n",
      "Epoch 1 Batch 1100 Loss 2.0264\n",
      "Epoch 1 Batch 1150 Loss 1.9716\n",
      "Epoch 1 Batch 1200 Loss 1.9426\n",
      "Epoch 1 Batch 1250 Loss 1.9165\n",
      "Epoch 1 Batch 1300 Loss 1.9068\n",
      "Epoch 1 Batch 1350 Loss 1.8900\n",
      "Epoch 1 Batch 1400 Loss 1.8727\n",
      "Epoch 1 Batch 1450 Loss 1.8450\n",
      "Epoch 1 Batch 1500 Loss 1.8203\n",
      "Epoch 1 Batch 1550 Loss 1.8046\n",
      "Epoch 1 Batch 1600 Loss 1.7825\n",
      "Epoch 1 Batch 1650 Loss 1.7812\n",
      "Epoch 1 Batch 1700 Loss 1.7723\n",
      "Epoch 1 Batch 1750 Loss 1.7539\n",
      "Epoch 1 Batch 1800 Loss 1.7392\n",
      "Epoch 1 Batch 1850 Loss 1.7183\n",
      "Epoch 1 Batch 1900 Loss 1.6922\n",
      "Epoch 1 Batch 1950 Loss 1.6758\n",
      "Time taken for 1 epoch: 1797.133960723877 secs\n",
      "\n",
      "Start of epoch 2\n",
      "Epoch 2 Batch 0 Loss 0.8552\n",
      "Epoch 2 Batch 50 Loss 1.0771\n",
      "Epoch 2 Batch 100 Loss 1.1816\n",
      "Epoch 2 Batch 150 Loss 1.1037\n",
      "Epoch 2 Batch 200 Loss 1.0746\n",
      "Epoch 2 Batch 250 Loss 1.0183\n",
      "Epoch 2 Batch 300 Loss 0.9815\n",
      "Epoch 2 Batch 350 Loss 0.9515\n",
      "Epoch 2 Batch 400 Loss 0.9228\n",
      "Epoch 2 Batch 450 Loss 0.8729\n",
      "Epoch 2 Batch 500 Loss 0.8495\n",
      "Epoch 2 Batch 550 Loss 0.8287\n",
      "Epoch 2 Batch 600 Loss 0.8106\n",
      "Epoch 2 Batch 650 Loss 0.7908\n",
      "Epoch 2 Batch 700 Loss 0.7639\n",
      "Epoch 2 Batch 750 Loss 0.7390\n",
      "Epoch 2 Batch 800 Loss 0.7207\n",
      "Epoch 2 Batch 850 Loss 0.7036\n",
      "Epoch 2 Batch 900 Loss 0.6926\n",
      "Epoch 2 Batch 950 Loss 0.6826\n",
      "Epoch 2 Batch 1000 Loss 0.6660\n",
      "Epoch 2 Batch 1050 Loss 0.6481\n",
      "Epoch 2 Batch 1100 Loss 0.6351\n",
      "Epoch 2 Batch 1150 Loss 0.6197\n",
      "Epoch 2 Batch 1200 Loss 0.6114\n",
      "Epoch 2 Batch 1250 Loss 0.6087\n",
      "Epoch 2 Batch 1300 Loss 0.6168\n",
      "Epoch 2 Batch 1350 Loss 0.6274\n",
      "Epoch 2 Batch 1400 Loss 0.6230\n",
      "Epoch 2 Batch 1450 Loss 0.6159\n",
      "Epoch 2 Batch 1500 Loss 0.6154\n",
      "Epoch 2 Batch 1550 Loss 0.6161\n",
      "Epoch 2 Batch 1600 Loss 0.6160\n",
      "Epoch 2 Batch 1650 Loss 0.6220\n",
      "Epoch 2 Batch 1700 Loss 0.6304\n",
      "Epoch 2 Batch 1750 Loss 0.6300\n",
      "Epoch 2 Batch 1800 Loss 0.6336\n",
      "Epoch 2 Batch 1850 Loss 0.6291\n",
      "Epoch 2 Batch 1900 Loss 0.6217\n",
      "Epoch 2 Batch 1950 Loss 0.6290\n",
      "Time taken for 1 epoch: 1771.54221200943 secs\n",
      "\n",
      "Start of epoch 3\n",
      "Epoch 3 Batch 0 Loss 1.7705\n",
      "Epoch 3 Batch 50 Loss 1.1217\n",
      "Epoch 3 Batch 100 Loss 1.1862\n",
      "Epoch 3 Batch 150 Loss 1.0839\n",
      "Epoch 3 Batch 200 Loss 1.0823\n",
      "Epoch 3 Batch 250 Loss 1.0159\n",
      "Epoch 3 Batch 300 Loss 0.9724\n",
      "Epoch 3 Batch 350 Loss 0.9733\n",
      "Epoch 3 Batch 400 Loss 0.9348\n",
      "Epoch 3 Batch 450 Loss 0.8767\n",
      "Epoch 3 Batch 500 Loss 0.8398\n",
      "Epoch 3 Batch 550 Loss 0.8241\n",
      "Epoch 3 Batch 600 Loss 0.8119\n",
      "Epoch 3 Batch 650 Loss 0.7931\n",
      "Epoch 3 Batch 700 Loss 0.7675\n",
      "Epoch 3 Batch 750 Loss 0.7367\n",
      "Epoch 3 Batch 800 Loss 0.7181\n",
      "Epoch 3 Batch 850 Loss 0.7073\n",
      "Epoch 3 Batch 900 Loss 0.6957\n",
      "Epoch 3 Batch 950 Loss 0.6846\n",
      "Epoch 3 Batch 1000 Loss 0.6667\n",
      "Epoch 3 Batch 1050 Loss 0.6530\n",
      "Epoch 3 Batch 1100 Loss 0.6379\n",
      "Epoch 3 Batch 1150 Loss 0.6224\n",
      "Epoch 3 Batch 1200 Loss 0.6129\n",
      "Epoch 3 Batch 1250 Loss 0.6088\n",
      "Epoch 3 Batch 1300 Loss 0.6159\n",
      "Epoch 3 Batch 1350 Loss 0.6249\n",
      "Epoch 3 Batch 1400 Loss 0.6241\n",
      "Epoch 3 Batch 1450 Loss 0.6184\n",
      "Epoch 3 Batch 1500 Loss 0.6169\n",
      "Epoch 3 Batch 1550 Loss 0.6130\n",
      "Epoch 3 Batch 1600 Loss 0.6134\n",
      "Epoch 3 Batch 1650 Loss 0.6229\n",
      "Epoch 3 Batch 1700 Loss 0.6294\n",
      "Epoch 3 Batch 1750 Loss 0.6305\n",
      "Epoch 3 Batch 1800 Loss 0.6324\n",
      "Epoch 3 Batch 1850 Loss 0.6282\n",
      "Epoch 3 Batch 1900 Loss 0.6223\n",
      "Epoch 3 Batch 1950 Loss 0.6299\n",
      "Time taken for 1 epoch: 1821.9350378513336 secs\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(NB_EPOCHS):\n",
    "    print(\"Start of epoch {}\".format(epoch+1))\n",
    "    start = time.time()\n",
    "    \n",
    "    train_loss.reset_states()\n",
    "    \n",
    "    for (batch, (inputs, targets)) in enumerate(train_dataset_light):\n",
    "        with tf.GradientTape() as tape:\n",
    "            model_outputs = bert_squad(inputs)\n",
    "            loss = squad_loss_fn(targets, model_outputs)\n",
    "        \n",
    "        gradients = tape.gradient(loss, bert_squad.trainable_variables)\n",
    "        optimizer.apply_gradients(zip(gradients, bert_squad.trainable_variables))\n",
    "        \n",
    "        train_loss(loss)\n",
    "        \n",
    "        if batch % 50 == 0:\n",
    "            print(\"Epoch {} Batch {} Loss {:.4f}\".format(\n",
    "                epoch+1, batch, train_loss.result()))\n",
    "        \n",
    "        #if batch % 500 == 0:\n",
    "            #ckpt_save_path = ckpt_manager.save()\n",
    "            #print(\"Saving checkpoint for epoch {} at {}\".format(epoch+1, ckpt_save_path))\n",
    "    print(\"Time taken for 1 epoch: {} secs\\n\".format(time.time() - start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v6WquTCiIR7t"
   },
   "source": [
    "# Stage 5: Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qDIlHd5Tos6C"
   },
   "source": [
    "## Prepare evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kU7_AyTIpTTJ"
   },
   "source": [
    "Get the dev set in the session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "VGKl84s5WrhD"
   },
   "outputs": [],
   "source": [
    "eval_examples = read_squad_examples(\n",
    "    \"dev-v1.1.json\",\n",
    "    is_training=False,\n",
    "    version_2_with_negative=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DAEUcZDSpYLD"
   },
   "source": [
    "Define the function that will write the tf_record file for the dev set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "id": "yCVmIgnEo83o"
   },
   "outputs": [],
   "source": [
    "eval_writer = FeatureWriter(\n",
    "    filename=os.path.join(\"\",\n",
    "                     \"eval.tf_record\"),\n",
    "    is_training=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C8aSLFdmp71I"
   },
   "source": [
    "Create a tokenizer for future information needs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "oH5exQ7KwnuH"
   },
   "outputs": [],
   "source": [
    "my_bert_layer = hub.KerasLayer(\n",
    "    \"https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/1\",\n",
    "    trainable=False)\n",
    "vocab_file = my_bert_layer.resolved_object.vocab_file.asset_path.numpy()\n",
    "do_lower_case = my_bert_layer.resolved_object.do_lower_case.numpy()\n",
    "tokenizer = FullTokenizer(vocab_file, do_lower_case)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vdHudjJ_qAzo"
   },
   "source": [
    "Define the function that add the features (feature is a protocol in tensorflow) to our eval_features list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "id": "bmQ5GtoTxRjU"
   },
   "outputs": [],
   "source": [
    "def _append_feature(feature, is_padding):\n",
    "    if not is_padding:\n",
    "        eval_features.append(feature)\n",
    "    eval_writer.process_feature(feature)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HLAcwCiaqHi_"
   },
   "source": [
    "Create the eval features and the writes the tf.record file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "id": "mz7kGYmUwGQb"
   },
   "outputs": [],
   "source": [
    "eval_features = []\n",
    "dataset_size = convert_examples_to_features(\n",
    "    examples=eval_examples,\n",
    "    tokenizer=tokenizer,\n",
    "    max_seq_length=384,\n",
    "    doc_stride=128,\n",
    "    max_query_length=64,\n",
    "    is_training=False,\n",
    "    output_fn=_append_feature,\n",
    "    batch_size=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "WpZfwPEwMabx"
   },
   "outputs": [],
   "source": [
    "eval_writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TbKhx3zuq844"
   },
   "source": [
    "Load the ready-to-be-used dataset to our session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "id": "eUqYvG5TxctF"
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 4\n",
    "\n",
    "eval_dataset = create_squad_dataset(\n",
    "    \"eval.tf_record\",\n",
    "    384,#input_meta_data['max_seq_length'],\n",
    "    BATCH_SIZE,\n",
    "    is_training=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tRbKrFYoo8e8"
   },
   "source": [
    "## Making the predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KyckEWDbrLEX"
   },
   "source": [
    "Defines a certain type of collection (like a dictionary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "tyWOUKaDP0H0"
   },
   "outputs": [],
   "source": [
    "RawResult = collections.namedtuple(\"RawResult\",\n",
    "                                   [\"unique_id\", \"start_logits\", \"end_logits\"])\n",
    "# id of the question, output of our model for the beginning of the ques, end of the answer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "28abKVvqrRa4"
   },
   "source": [
    "Returns each element of batched output at a time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "id": "BScaA0SZQgQW"
   },
   "outputs": [],
   "source": [
    "def get_raw_results(predictions):\n",
    "    for unique_ids, start_logits, end_logits in zip(predictions['unique_ids'],\n",
    "                                                    predictions['start_logits'],\n",
    "                                                    predictions['end_logits']):\n",
    "        yield RawResult(\n",
    "            unique_id=unique_ids.numpy(),\n",
    "            start_logits=start_logits.numpy().tolist(),\n",
    "            end_logits=end_logits.numpy().tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JiLOOmnLre5C"
   },
   "source": [
    "Let's make our predictions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "id": "qqD78Xdjrvpn"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0/2709\n",
      "100/2709\n",
      "200/2709\n",
      "300/2709\n",
      "400/2709\n",
      "500/2709\n",
      "600/2709\n",
      "700/2709\n",
      "800/2709\n",
      "900/2709\n",
      "1000/2709\n",
      "1100/2709\n",
      "1200/2709\n",
      "1300/2709\n",
      "1400/2709\n",
      "1500/2709\n",
      "1600/2709\n",
      "1700/2709\n",
      "1800/2709\n",
      "1900/2709\n",
      "2000/2709\n",
      "2100/2709\n",
      "2200/2709\n",
      "2300/2709\n",
      "2400/2709\n",
      "2500/2709\n",
      "2600/2709\n",
      "2700/2709\n"
     ]
    }
   ],
   "source": [
    "all_results = []\n",
    "for count, inputs in enumerate(eval_dataset):\n",
    "    x, _ = inputs\n",
    "    unique_ids = x.pop(\"unique_ids\")\n",
    "    start_logits, end_logits = bert_squad(x, training=False)\n",
    "    output_dict = dict(\n",
    "        unique_ids=unique_ids,\n",
    "        start_logits=start_logits,\n",
    "        end_logits=end_logits)\n",
    "    for result in get_raw_results(output_dict):\n",
    "        all_results.append(result)\n",
    "    if count % 100 == 0:\n",
    "        print(\"{}/{}\".format(count, 2709))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PjQ6kIqGriHr"
   },
   "source": [
    "Write the predictions in a json file that will work with the evaluation script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "id": "esLdRf7uM3Lz"
   },
   "outputs": [],
   "source": [
    "output_prediction_file = \"predictions.json\"\n",
    "output_nbest_file = \"nbest_predictions.json\"\n",
    "output_null_log_odds_file = \"null_odds.json\"\n",
    "\n",
    "write_predictions(\n",
    "    eval_examples,\n",
    "    eval_features,\n",
    "    all_results,\n",
    "    20,\n",
    "    30,\n",
    "    True,\n",
    "    output_prediction_file,\n",
    "    output_nbest_file,\n",
    "    output_null_log_odds_file,\n",
    "    verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-eaIHyDIYHHx"
   },
   "source": [
    "## Home-made prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P0F4l5h8Zdha"
   },
   "source": [
    "### Input dict creation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IrhHzx7ycXXo"
   },
   "source": [
    "We will concatenate the question and the context, separated by a `[\"SEP\"]`, after tokenization, as it has been made for the training set.\n",
    "\n",
    "The important thing is that we want our answer to start with a real word and to end with a real word. The word \"ecologically\" being tokenized as `[\"ecological\", \"##ly\"]`, if the ending token is `[\"ecological\"]` we want to use \"ecologically\" as the ending word (same thing if the ending token is `[\"##ly\"]`). That is why we first split our context into words, and then into tokens, remembering to which word each token belongs to (see the `tokenize_context()` function)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_tuNXt98Zm4u"
   },
   "source": [
    "#### Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "id": "OBjGoQ_wfmml"
   },
   "outputs": [],
   "source": [
    "my_bert_layer = hub.KerasLayer(\n",
    "    \"https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/1\",\n",
    "    trainable=False)\n",
    "vocab_file = my_bert_layer.resolved_object.vocab_file.asset_path.numpy()\n",
    "do_lower_case = my_bert_layer.resolved_object.do_lower_case.numpy()\n",
    "tokenizer = FullTokenizer(vocab_file, do_lower_case)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "id": "8f_fCe_hLC12"
   },
   "outputs": [],
   "source": [
    "def is_whitespace(c):\n",
    "    '''\n",
    "    Tell if a chain of characters corresponds to a whitespace or not.\n",
    "    '''\n",
    "    if c == \" \" or c == \"\\t\" or c == \"\\r\" or c == \"\\n\" or ord(c) == 0x202F:\n",
    "        return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "id": "4ZA4TYnxLGVT"
   },
   "outputs": [],
   "source": [
    "def whitespace_split(text):\n",
    "    '''\n",
    "    Take a text and return a list of \"words\" by splitting it according to\n",
    "    whitespaces.\n",
    "    '''\n",
    "    doc_tokens = []\n",
    "    prev_is_whitespace = True\n",
    "    for c in text:\n",
    "        if is_whitespace(c):\n",
    "            prev_is_whitespace = True\n",
    "        else:\n",
    "            if prev_is_whitespace:\n",
    "                doc_tokens.append(c)\n",
    "            else:\n",
    "                doc_tokens[-1] += c\n",
    "            prev_is_whitespace = False\n",
    "    return doc_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "id": "Fsfzt3GUNWQK"
   },
   "outputs": [],
   "source": [
    "def tokenize_context(text_words):\n",
    "    '''\n",
    "    Take a list of words (returned by whitespace_split()) and tokenize each word\n",
    "    one by one. Also keep track, for each new token, of its original word in the\n",
    "    text_words parameter.\n",
    "    '''\n",
    "    text_tok = []\n",
    "    tok_to_word_id = []\n",
    "    for word_id, word in enumerate(text_words):\n",
    "        word_tok = tokenizer.tokenize(word)\n",
    "        text_tok += word_tok\n",
    "        tok_to_word_id += [word_id]*len(word_tok)\n",
    "    return text_tok, tok_to_word_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "id": "c8qreqEURUOP"
   },
   "outputs": [],
   "source": [
    "def get_ids(tokens):\n",
    "    return tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "def get_mask(tokens):\n",
    "    return np.char.not_equal(tokens, \"[PAD]\").astype(int)\n",
    "\n",
    "def get_segments(tokens):\n",
    "    seg_ids = []\n",
    "    current_seg_id = 0\n",
    "    for tok in tokens:\n",
    "        seg_ids.append(current_seg_id)\n",
    "        if tok == \"[SEP]\":\n",
    "            current_seg_id = 1-current_seg_id # turns 1 into 0 and vice versa\n",
    "    return seg_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "id": "P2sPGXxsYUsY"
   },
   "outputs": [],
   "source": [
    "def create_input_dict(question, context):\n",
    "    '''\n",
    "    Take a question and a context as strings and return a dictionary with the 3\n",
    "    elements needed for the model. Also return the context_words, the\n",
    "    context_tok to context_word ids correspondance and the length of\n",
    "    question_tok that we will need later.\n",
    "    '''\n",
    "    question_tok = tokenizer.tokenize(my_question)\n",
    "\n",
    "    context_words = whitespace_split(context)\n",
    "    context_tok, context_tok_to_word_id = tokenize_context(context_words)\n",
    "\n",
    "    input_tok = question_tok + [\"[SEP]\"] + context_tok + [\"[SEP]\"]\n",
    "    input_tok += [\"[PAD]\"]*(384-len(input_tok)) # in our case the model has been\n",
    "                                                # trained to have inputs of length max 384\n",
    "    input_dict = {}\n",
    "    input_dict[\"input_word_ids\"] = tf.expand_dims(tf.cast(get_ids(input_tok), tf.int32), 0)\n",
    "    input_dict[\"input_mask\"] = tf.expand_dims(tf.cast(get_mask(input_tok), tf.int32), 0)\n",
    "    input_dict[\"input_type_ids\"] = tf.expand_dims(tf.cast(get_segments(input_tok), tf.int32), 0)\n",
    "\n",
    "    return input_dict, context_words, context_tok_to_word_id, len(question_tok)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QAnaCZWTZpWT"
   },
   "source": [
    "#### Creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "id": "WQlM-M8rMklA"
   },
   "outputs": [],
   "source": [
    "my_context = '''Neoclassical economics views inequalities in the distribution of income as arising from differences in value added by labor, capital and land. Within labor income distribution is due to differences in value added by different classifications of workers. In this perspective, wages and profits are determined by the marginal value added of each economic actor (worker, capitalist/business owner, landlord). Thus, in a market economy, inequality is a reflection of the productivity gap between highly-paid professions and lower-paid professions.'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YL_cE-o0U8mx"
   },
   "source": [
    "Neoclassical economics views inequalities in the distribution of income as arising from differences in value added by labor, capital and land. Within labor income distribution is due to differences in value added by different classifications of workers. In this perspective, wages and profits are determined by the marginal value added of each economic actor (worker, capitalist/business owner, landlord). Thus, in a market economy, inequality is a reflection of the productivity gap between highly-paid professions and lower-paid professions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "id": "SeB29SCNNQ1M"
   },
   "outputs": [],
   "source": [
    "#my_question = '''What philosophy of thought addresses wealth inequality?'''\n",
    "my_question = '''What are examples of economic actors?'''\n",
    "#my_question = '''In a market economy, what is inequality a reflection of?'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "id": "-v5nLMNjZufe"
   },
   "outputs": [],
   "source": [
    "my_input_dict, my_context_words, context_tok_to_word_id, question_tok_len = create_input_dict(my_question, my_context)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dT066rMtZ65X"
   },
   "source": [
    "### Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "id": "CcJgDa9gVShl"
   },
   "outputs": [],
   "source": [
    "start_logits, end_logits = bert_squad(my_input_dict, training=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rhdGlIo5Z9IZ"
   },
   "source": [
    "### Interpretation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZQJMBkVLd9wp"
   },
   "source": [
    "We remove the ids corresponding to the question and the `[\"SEP\"]` token:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "id": "TfwBJfsSTwRn"
   },
   "outputs": [],
   "source": [
    "start_logits_context = start_logits.numpy()[0, question_tok_len+1:]\n",
    "end_logits_context = end_logits.numpy()[0, question_tok_len+1:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "te1u6iZAawYf"
   },
   "source": [
    "First easy interpretation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "id": "lQJ8tp-1WvI4"
   },
   "outputs": [],
   "source": [
    "start_word_id = context_tok_to_word_id[np.argmax(start_logits_context)]\n",
    "end_word_id = context_tok_to_word_id[np.argmax(end_logits_context)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "k3PHA84rarse"
   },
   "source": [
    "\"Advanced\" - making sure that the start of the answer is before the end:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "id": "xiFZ2fUiRU_M"
   },
   "outputs": [],
   "source": [
    "pair_scores = np.ones((len(start_logits_context), len(end_logits_context)))*(-1E10)\n",
    "for i in range(len(start_logits_context-1)):\n",
    "    for j in range(i, len(end_logits_context)):\n",
    "        pair_scores[i, j] = start_logits_context[i] + end_logits_context[j]\n",
    "pair_scores_argmax = np.argmax(pair_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "id": "N9KEiFHPXXeM"
   },
   "outputs": [],
   "source": [
    "start_word_id = context_tok_to_word_id[pair_scores_argmax // len(start_logits_context)]\n",
    "end_word_id = context_tok_to_word_id[pair_scores_argmax % len(end_logits_context)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kJDBL8KKa6NP"
   },
   "source": [
    "Final answer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "id": "j0Y3WDz0XwAw"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The answer to:\n",
      "What are examples of economic actors?\n",
      "is:\n",
      "(worker, capitalist/business owner, landlord).\n"
     ]
    }
   ],
   "source": [
    "predicted_answer = ' '.join(my_context_words[start_word_id:end_word_id+1])\n",
    "print(\"The answer to:\\n\" + my_question + \"\\nis:\\n\" + predicted_answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "id": "bYGSk_5OSYUk"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h2>WHAT ARE EXAMPLES OF ECONOMIC ACTORS?</h2>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<blockquote> Neoclassical economics views inequalities in the distribution of income as arising from differences in value added by labor, capital and land. Within labor income distribution is due to differences in value added by different classifications of workers. In this perspective, wages and profits are determined by the marginal value added of each economic actor <mark>(worker, capitalist/business owner, landlord).</mark> Thus, in a market economy, inequality is a reflection of the productivity gap between highly-paid professions and lower-paid professions. </blockquote>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import HTML\n",
    "display(HTML(f'<h2>{my_question.upper()}</h2>'))\n",
    "marked_text = str(my_context.replace(predicted_answer, f\"<mark>{predicted_answer}</mark>\"))\n",
    "display(HTML(f\"\"\"<blockquote> {marked_text} </blockquote>\"\"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Brg0AXkA0r74"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Bert_QA",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}