Copy_of_t5_trivia.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "name": "Copy of t5-trivia",
      "provenance": [],
      "collapsed_sections": [
        "zrtR2urJV3ST",
        "-pFvyrHmm6Mx"
      ],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jonbaer/googlecolab/blob/master/Copy_of_t5_trivia.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "YONnGjpAYUdU"
      },
      "source": [
        "\n",
        "<a href=\"https://colab.research.google.com/github/google-research/text-to-text-transfer-transformer/blob/master/notebooks/t5-trivia.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zrtR2urJV3ST"
      },
      "source": [
        "##### Copyright 2019 The T5 Authors\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "DWdCSqJ6WHBh",
        "colab": {}
      },
      "source": [
        "# Copyright 2019 The T5 Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zSeyoqE7WMwu"
      },
      "source": [
        "# Fine-Tuning the Text-To-Text Transfer Transformer (T5) for Closed-Book Question Answering\n",
        "## _Or: What does T5 know?_\n",
        "\n",
        "*The following tutorial guides you through the process of fine-tuning a pre-trained T5 model, evaluating its accuracy, and using it for prediction,\n",
        "all on a free Google Cloud TPU <a href=\"https://colab.research.google.com/github/google-research/text-to-text-transfer-transformer/blob/master/notebooks/t5-trivia.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>.*\n",
        "\n",
        "### Background\n",
        "\n",
        "T5 was introduced in the paper [_Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer_](https://arxiv.org/abs/1910.10683). In that paper, we provided a comprehensive picture of how we pre-trained a standard text-to-text Transformer model on a large text corpus, achieving state-of-the-art results on many NLP tasks after fine-tuning.\n",
        "\n",
        "We pre-trained T5 on a mixture of supervised and unsupervised tasks with the majoriy of data coming from an unlabeled dataset we developed called [C4](https://www.tensorflow.org/datasets/catalog/c4). C4 is based on a massive scrape of the web produced by [Common Crawl](https://commoncrawl.org). Loosely speaking, pre-training on C4 ideally gives T5 an understanding of natural language in addition to general world knowledge.\n",
        "\n",
        "### How can we assess what T5 knows?\n",
        "\n",
        "As the name implies, T5 is a text-to-text model, which enables us to train it on arbitrary tasks involving a textual input and output. As we showed in our paper, a huge variety of NLP tasks can be cast in this format, including translation, summarization, and even classification and regression tasks.\n",
        "\n",
        "One way to use this text-to-text framework is on reading comprehension problems, where the model is fed some context along with a question and is trained to predict the question's answer. For example, we might feed the model the text from the Wikipedia article about [Hurrican Connie](https://en.wikipedia.org/wiki/Hurricane_Connie) along with the question \"On what date did Hurricane Connie occur?\" and train the model to predict the answer \"August 3rd, 1955\".\n",
        "A related task is open-domain question answering (QA) where the model is not provided with this oracle context. Typically, open-domain QA systems include a mechanism to look up information in an external knowledge source. This setting is similar to an \"open-book\" exam.\n",
        "\n",
        "In this notebook, we'll be training T5 on a variant of this task which we call **closed-book question answering**. In closed-book QA, we feed the model a question *without any context or access to external knowledge* and train it to predict the answer. Since the model doesn't receive any context, the primary way it can learn to answer these questions is based on the \"knowledge\" it obtained during pre-training. We don't expect T5 to contain super specific information, so we will be focusing on two question-answering datasets which largely include trivia questions (i.e. facts about well-known subjects). [Similar](https://arxiv.org/abs/1909.01066) [investigations](https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf) have recently been done to test the knowledge stored by BERT and GPT-2.\n",
        "\n",
        "T5 was not pre-trained on closed-book QA, so in this notebook we'll first create two new tasks and then use the [`t5`](https://github.com/google-research/text-to-text-transfer-transformer) library to fine-tune, evaluate, and obtain predictions from T5. In the end, T5's performance on closed-book QA can give us a sense of what kind (and how much) information T5 managed to learn during pre-training.\n",
        "\n",
        "We published a [more in-depth investigation](https://tiny.cc/t5-qa) of closed-book QA with T5 where we achieved suprisingly strong performance on open-domain variants of Natural Questions, WebQuestions, and TriviaQA. The code in this notebook is a simplified version of those experiments but still produces good results.\n",
        "\n",
        "\n",
        "### Caveats\n",
        "\n",
        "* While we provide instructions for running on a [Cloud TPU](https://cloud.google.com/tpu/) via Colab for free, a [Google Cloud Storage (GCS)](http://console.cloud.google.com/storage) bucket is required for storing model parameters and data. The [GCS free tier](https://cloud.google.com/free/) provides 5 GB of storage, which should be enough to train the `large` model and smaller but not the `3B` or `11B` parameter models. You can use part of your initial $300 credit to get more space.\n",
        "* The Cloud TPU provided by Colab (a `v2-8`) does not have enough memory to fine-tune the `11B` parameter model. For this model, you will need to fine-tune inside of a GCP instance (see [README](https://github.com/google-research/text-to-text-transfer-transformer/)).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yAb_APDrefs6"
      },
      "source": [
        "# Set Up"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "eDeE_yVuHMYg"
      },
      "source": [
        "<h3><a href=\"https://cloud.google.com/tpu/\"><img valign=\"middle\" src=\"https://raw.githubusercontent.com/GoogleCloudPlatform/tensorflow-without-a-phd/master/tensorflow-rl-pong/images/tpu-hexagon.png\" width=\"50\"></a>  &nbsp;&nbsp;Train on TPU</h3>\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "   1. Create a Cloud Storage bucket for your data and model checkpoints at http://console.cloud.google.com/storage, and fill in the `BASE_DIR` parameter in the following form. There is a [free tier](https://cloud.google.com/free/) if you do not yet have an account.\n",
        " \n",
        "   1. On the main menu, click Runtime and select **Change runtime type**. Set \"TPU\" as the hardware accelerator.\n",
        "   1. Run the following cell and follow instructions to:\n",
        "    *  Set up a Colab TPU running environment\n",
        "    *   Verify that you are connected to a TPU device\n",
        "    *   Upload your credentials to TPU to access your GCS bucket\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "xYh-IaN4C7Z1",
        "colab": {}
      },
      "source": [
        "# TODO(adarob): Add support for 2.x.\n",
        "%tensorflow_version 1.x\n",
        "\n",
        "import datetime\n",
        "import functools\n",
        "import json\n",
        "import os\n",
        "import pprint\n",
        "import random\n",
        "import string\n",
        "import sys\n",
        "import tensorflow as tf\n",
        "\n",
        "BASE_DIR = \"\" #@param { type: \"string\" }\n",
        "if not BASE_DIR:\n",
        "  raise ValueError(\"You must enter a BASE_DIR.\")\n",
        "DATA_DIR = os.path.join(BASE_DIR, \"data\")\n",
        "MODELS_DIR = os.path.join(BASE_DIR, \"models\")\n",
        "ON_CLOUD = True\n",
        "\n",
        "\n",
        "if ON_CLOUD:\n",
        "  assert \"COLAB_TPU_ADDR\" in os.environ, \"ERROR: Not connected to a TPU runtime; please see the first cell in this notebook for instructions!\"\n",
        "  TPU_ADDRESS = \"grpc://\" + os.environ[\"COLAB_TPU_ADDR\"] \n",
        "  TPU_TOPOLOGY = \"2x2\"\n",
        "  print(\"TPU address is\", TPU_ADDRESS)\n",
        "\n",
        "  from google.colab import auth\n",
        "  auth.authenticate_user()\n",
        "  with tf.Session(TPU_ADDRESS) as session:\n",
        "    print('TPU devices:')\n",
        "    pprint.pprint(session.list_devices())\n",
        "\n",
        "    # Upload credentials to TPU.\n",
        "    with open('/content/adc.json', 'r') as f:\n",
        "      auth_info = json.load(f)\n",
        "    tf.contrib.cloud.configure_gcs(session, credentials=auth_info)\n",
        "    # Now credentials are set for all future sessions on this TPU."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "zzZU6QZ1bHyC",
        "colab": {}
      },
      "source": [
        "#@title Install and import required packages\n",
        "if ON_CLOUD:\n",
        "  !pip install -qU t5\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\", category=DeprecationWarning)\n",
        "\n",
        "import t5\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "import time\n",
        "\n",
        "# Improve logging.\n",
        "from contextlib import contextmanager\n",
        "import logging as py_logging\n",
        "\n",
        "if ON_CLOUD:\n",
        "  tf.get_logger().propagate = False\n",
        "  py_logging.root.setLevel('INFO')\n",
        "\n",
        "@contextmanager\n",
        "def tf_verbosity_level(level):\n",
        "  og_level = tf.logging.get_verbosity()\n",
        "  tf.logging.set_verbosity(level)\n",
        "  yield\n",
        "  tf.logging.set_verbosity(og_level)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "dMoJ-G9mqDqa"
      },
      "source": [
        "# Creating new Tasks and Mixture"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zwoLPQhE6bef"
      },
      "source": [
        "Two core components of the T5 library are `Task` and `Mixture` objects.\n",
        "\n",
        "A `Task` is a dataset along with preprocessing functions and evaluation metrics. A `Mixture` is a collection of `Task` objects along with a mixing rate or a function defining how to compute a mixing rate based on the properties of the constituent `Tasks`.\n",
        "\n",
        "For this example, we will fine-tune the model to do closed-book question answering."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "152zECujzPMk"
      },
      "source": [
        "### Natural Questions\n",
        "\n",
        "[Natural Questions (NQ)](https://ai.google.com/research/NaturalQuestions) is a challenging corpus for open-domain QA. Each example includes a question along with an entire Wikipedia article that may or may not contain its answer. The goal is to produce the correct answer given this context. In our case, we will be ignoring the provided context in hopes that the model will learn to find the answers from the world knowledge it has acquired during pre-training.\n",
        "\n",
        "Since the raw data splits are stored as JSONL files, we will first need to convert them to TSV format to make them parseable in TensorFlow. We will also take the opportunity to drop information we will not be using, remove questions with multiple answers, and to do a bit of cleaning of the text."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "OjEonhK3zNRu",
        "outputId": "feb1e52f-0b86-4108-89b8-51d516499f11",
        "colab": {
          "height": 34
        }
      },
      "source": [
        "import gzip\n",
        "import json\n",
        "\n",
        "# Public directory of Natural Questions data on GCS.\n",
        "NQ_JSONL_DIR = \"gs://natural_questions/v1.0-simplified/\"\n",
        "NQ_SPLIT_FNAMES = {\n",
        "    \"train\": \"simplified-nq-train.jsonl.gz\",\n",
        "    \"validation\": \"nq-dev-all.jsonl.gz\"\n",
        "}\n",
        "nq_counts_path = os.path.join(DATA_DIR, \"nq-counts.json\")\n",
        "nq_tsv_path = {\n",
        "    \"train\": os.path.join(DATA_DIR, \"nq-train.tsv\"),\n",
        "    \"validation\": os.path.join(DATA_DIR, \"nq-validation.tsv\")\n",
        "}\n",
        "\n",
        "def nq_jsonl_to_tsv(in_fname, out_fname):\n",
        "\n",
        "  def extract_answer(tokens, span):\n",
        "    \"\"\"Reconstruct answer from token span and remove extra spaces.\"\"\"\n",
        "    start, end = span[\"start_token\"], span[\"end_token\"]  \n",
        "    ans = \" \".join(tokens[start:end])\n",
        "    # Remove incorrect spacing around punctuation.\n",
        "    ans = ans.replace(\" ,\", \",\").replace(\" .\", \".\").replace(\" %\", \"%\")\n",
        "    ans = ans.replace(\" - \", \"-\").replace(\" : \", \":\").replace(\" / \", \"/\")\n",
        "    ans = ans.replace(\"( \", \"(\").replace(\" )\", \")\")\n",
        "    ans = ans.replace(\"`` \", \"\\\"\").replace(\" ''\", \"\\\"\")\n",
        "    ans = ans.replace(\" 's\", \"'s\").replace(\"s ' \", \"s' \")\n",
        "    return ans\n",
        "\n",
        "  count = 0\n",
        "  with tf.io.gfile.GFile(in_fname, \"rb\") as infile,\\\n",
        "       tf.io.gfile.GFile(out_fname, \"w\") as outfile:\n",
        "    for line in gzip.open(infile):\n",
        "      ex = json.loads(line)\n",
        "      # Remove any examples with more than one answer.\n",
        "      if len(ex['annotations'][0]['short_answers']) != 1:\n",
        "        continue\n",
        "      # Questions in NQ do not include a question mark.\n",
        "      question = ex[\"question_text\"] + \"?\"\n",
        "      answer_span = ex['annotations'][0]['short_answers'][0]\n",
        "      # Handle the two document formats in NQ (tokens or text).\n",
        "      if \"document_tokens\" in ex:\n",
        "        tokens = [t[\"token\"] for t in ex[\"document_tokens\"]]\n",
        "      elif \"document_text\" in ex:\n",
        "        tokens = ex[\"document_text\"].split(\" \")\n",
        "      answer = extract_answer(tokens, answer_span)\n",
        "      # Write this line as <question>\\t<answer>\n",
        "      outfile.write(\"%s\\t%s\\n\" % (question, answer))\n",
        "      count += 1\n",
        "      tf.logging.log_every_n(\n",
        "          tf.logging.INFO,\n",
        "          \"Wrote %d examples to %s.\" % (count, out_fname),\n",
        "          1000)\n",
        "    return count\n",
        "\n",
        "if tf.io.gfile.exists(nq_counts_path):\n",
        "  # Used cached data and counts.\n",
        "  tf.logging.info(\"Loading NQ from cache.\")\n",
        "  num_nq_examples = json.load(tf.io.gfile.GFile(nq_counts_path))\n",
        "else:\n",
        "  # Create TSVs and get counts.\n",
        "  tf.logging.info(\"Generating NQ TSVs.\")\n",
        "  num_nq_examples = {}\n",
        "  for split, fname in NQ_SPLIT_FNAMES.items():\n",
        "    num_nq_examples[split] = nq_jsonl_to_tsv(\n",
        "        os.path.join(NQ_JSONL_DIR, fname), nq_tsv_path[split])\n",
        "  json.dump(num_nq_examples, tf.io.gfile.GFile(nq_counts_path, \"w\"))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "I1206 00:11:00.169766 248738 <ipython-input-3-45e03d923fbf>:51] Loading NQ from cache.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "R-Ja8akCX1dR"
      },
      "source": [
        "Next, we define a function to load the TSV data as a `tf.data.Dataset` in TensorFlow."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "KPOteeqctpzw",
        "outputId": "ef61274d-170d-4c56-dadc-427a679521e9",
        "colab": {
          "height": 139
        }
      },
      "source": [
        "def nq_dataset_fn(split, shuffle_files=False):\n",
        "  # We only have one file for each split.\n",
        "  del shuffle_files\n",
        "\n",
        "  # Load lines from the text file as examples.\n",
        "  ds = tf.data.TextLineDataset(nq_tsv_path[split])\n",
        "  # Split each \"<question>\\t<answer>\" example into (question, answer) tuple.\n",
        "  ds = ds.map(\n",
        "      functools.partial(tf.io.decode_csv, record_defaults=[\"\", \"\"],\n",
        "                        field_delim=\"\\t\", use_quote_delim=False),\n",
        "      num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "  # Map each tuple to a {\"question\": ... \"answer\": ...} dict.\n",
        "  ds = ds.map(lambda *ex: dict(zip([\"question\", \"answer\"], ex)))\n",
        "  return ds\n",
        "\n",
        "print(\"A few raw validation examples...\")\n",
        "for ex in tfds.as_numpy(nq_dataset_fn(\"validation\").take(5)):\n",
        "  print(ex)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "A few raw validation examples...\n",
            "{'question': b'what do the 3 dots mean in math?', 'answer': b'the therefore sign (\\xe2\\x88\\xb4) is generally used before a logical consequence, such as the conclusion of a syllogism'}\n",
            "{'question': b'who is playing the halftime show at super bowl 2016?', 'answer': b'Coldplay with special guest performers Beyonc\\xc3\\xa9 and Bruno Mars'}\n",
            "{'question': b'who won the 2017 sports personality of the year?', 'answer': b'Mo Farah'}\n",
            "{'question': b'where was the world economic forum held this year?', 'answer': b'Davos, a mountain resort in Graub\\xc3\\xbcnden, in the eastern Alps region of Switzerland'}\n",
            "{'question': b'who has made the most premier league appearances?', 'answer': b'Gareth Barry'}\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MCUYT7JmX9Tj"
      },
      "source": [
        "Now, we write a preprocess function to convert the examples in the `tf.data.Dataset` into a text-to-text format, with both `inputs` and `targets` fields. The preprocessor also normalizes the text by lowercasing it and removing quotes since the answers are sometimes formatted in odd ways. Finally, we prepend 'trivia question:' to the inputs so that the model knows what task it's trying to solve."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "x8tNn6HMYLMb",
        "colab": {}
      },
      "source": [
        "def trivia_preprocessor(ds):\n",
        "  def normalize_text(text):\n",
        "    \"\"\"Lowercase and remove quotes from a TensorFlow string.\"\"\"\n",
        "    text = tf.strings.lower(text)\n",
        "    text = tf.strings.regex_replace(text,\"'(.*)'\", r\"\\1\")\n",
        "    return text\n",
        "\n",
        "  def to_inputs_and_targets(ex):\n",
        "    \"\"\"Map {\"question\": ..., \"answer\": ...}->{\"inputs\": ..., \"targets\": ...}.\"\"\"\n",
        "    return {\n",
        "        \"inputs\":\n",
        "             tf.strings.join(\n",
        "                 [\"trivia question: \", normalize_text(ex[\"question\"])]),\n",
        "        \"targets\": normalize_text(ex[\"answer\"])\n",
        "    }\n",
        "  return ds.map(to_inputs_and_targets, \n",
        "                num_parallel_calls=tf.data.experimental.AUTOTUNE)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gm1Pm2aRZ9Ow"
      },
      "source": [
        "Finally, we put everything together to create a `Task`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "yJyRavOpZ7UW",
        "colab": {}
      },
      "source": [
        "t5.data.TaskRegistry.add(\n",
        "    \"nq_context_free\",\n",
        "    # Supply a function which returns a tf.data.Dataset.\n",
        "    dataset_fn=nq_dataset_fn,\n",
        "    splits=[\"train\", \"validation\"],\n",
        "    # Supply a function which preprocesses text from the tf.data.Dataset.\n",
        "    text_preprocessor=[trivia_preprocessor],\n",
        "    # Use the same vocabulary that we used for pre-training.\n",
        "    sentencepiece_model_path=t5.data.DEFAULT_SPM_PATH,\n",
        "    # Lowercase targets before computing metrics.\n",
        "    postprocess_fn=t5.data.postprocessors.lower_text, \n",
        "    # We'll use accuracy as our evaluation metric.\n",
        "    metric_fns=[t5.evaluation.metrics.accuracy],\n",
        "    # Not required, but helps for mixing and auto-caching.\n",
        "    num_input_examples=num_nq_examples\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "qe4o_0jFbP-p"
      },
      "source": [
        "Let's look at a few pre-processed examples from the validation set. Note they contain both the tokenized (integer) and plain-text inputs and targets.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "I64TqHGxbOJ2",
        "outputId": "68df8ac5-9d34-4dd2-de40-1c6b65feca3c",
        "colab": {
          "height": 275
        }
      },
      "source": [
        "nq_task = t5.data.TaskRegistry.get(\"nq_context_free\")\n",
        "ds = nq_task.get_dataset(split=\"validation\", sequence_length={\"inputs\": 128, \"targets\": 32})\n",
        "print(\"A few preprocessed validation examples...\")\n",
        "for ex in tfds.as_numpy(ds.take(5)):\n",
        "  print(ex)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "A few preprocessed validation examples...\n",
            "{'inputs_plaintext': b'trivia question: what is the average height of a chinese man?', 'inputs': array([22377,   822,    10,   125,    19,     8,  1348,  3902,    13,\n",
            "           3,     9,     3,  1436,  1496,    15,   388,    58,     1]), 'targets_plaintext': b'167.1 cm (5 ft 6 in)', 'targets': array([  898, 25059,  2446,  9209,     3,    89,    17,   431,    16,\n",
            "          61,     1])}\n",
            "{'inputs_plaintext': b'trivia question: what is the population of fayetteville north carolina?', 'inputs': array([22377,   822,    10,   125,    19,     8,  2074,    13,     3,\n",
            "          89,     9,    63,  1954,  1420,  3457,   443, 12057,     9,\n",
            "          58,     1]), 'targets_plaintext': b'204,408 in 2013', 'targets': array([    3, 26363,     6,  2445,   927,    16,  2038,     1])}\n",
            "{'inputs_plaintext': b'trivia question: capital of georgia the former soviet republic 7 letters?', 'inputs': array([22377,   822,    10,  1784,    13,   873,  1677,    23,     9,\n",
            "           8,  1798,    78,  5914,    17, 20237,   489,  5487,    58,\n",
            "           1]), 'targets_plaintext': b'tbilisi', 'targets': array([   3,   17, 3727,  159,   23,    1])}\n",
            "{'inputs_plaintext': b'trivia question: who plays jill bigelow in line of duty?', 'inputs': array([22377,   822,    10,   113,  4805,     3,   354,  1092,   600,\n",
            "          15,  3216,    16,   689,    13,  5461,    58,     1]), 'targets_plaintext': b'polly walker', 'targets': array([ 5492,    63,     3, 24063,     1])}\n",
            "{'inputs_plaintext': b'trivia question: when did we first put a rover on mars?', 'inputs': array([22377,   822,    10,   116,   410,    62,   166,   474,     3,\n",
            "           9,     3,    52,  1890,    30,  8113,    58,     1]), 'targets_plaintext': b'january 2004', 'targets': array([   3, 7066,   76, 1208, 4406,    1])}\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "G1ktIEGePdBr"
      },
      "source": [
        "**Note**: Instead of defining `nq_dataset_fn` and above, we also could have used the `TextLineTask` class with the `parse_tsv` preprocessor for equivalent results as follows:\n",
        "\n",
        "```py\n",
        "t5.data.TaskRegistry.add(\n",
        "    \"nq_context_free\",\n",
        "    t5.data.TextLineTask,\n",
        "    split_to_filepattern=nq_tsv_path,\n",
        "    text_preprocessor=[\n",
        "      functools.partial(\n",
        "          t5.data.preprocessors.parse_tsv,\n",
        "          field_names=[\"question\", \"answer\"]),\n",
        "      trivia_preprocessor\n",
        "    ],\n",
        "    postprocess_fn=t5.data.postprocessors.lower_text, \n",
        "    metric_fns=[t5.evaluation.metrics.accuracy],\n",
        "    num_input_examples=num_nq_examples\n",
        ")\n",
        "```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "y4_1gpcK9i3W"
      },
      "source": [
        "## TriviaQA\n",
        "\n",
        "A second dataset we will use is related to [TriviaQA](https://nlp.cs.washington.edu/triviaqa/). It is also intended for reading comprehension, but, once again, we will modify the task here by ignoring the provided context.\n",
        "\n",
        "Since the dataset has been imported into [TensorFlow Datasets (TFDS)](https://www.tensorflow.org/datasets/catalog/trivia_qa), we can let it handle the data parsing for us. It will take a few minutes to download and preprocess the first time, but we'll be able to access it instantly from our data directory afterward."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "mQTQHS94Z0Tq",
        "outputId": "b0c0c90b-e4c3-418d-971f-bdfaa52d6b40",
        "colab": {
          "height": 156
        }
      },
      "source": [
        "ds = tfds.load(\n",
        "    \"trivia_qa/unfiltered.nocontext\",\n",
        "    data_dir=DATA_DIR,\n",
        "    # Download data locally for preprocessing to avoid using GCS space.\n",
        "    download_and_prepare_kwargs={\"download_dir\": \"./downloads\"})\n",
        "print(\"A few raw validation examples...\")\n",
        "for ex in tfds.as_numpy(ds[\"validation\"].take(2)):\n",
        "  print(ex)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "A few raw validation examples...\n",
            "{'answer': {'aliases': array([b'Torquemada (disambiguation)', b'Torquemada'], dtype=object), 'matched_wiki_entity_name': b'', 'normalized_aliases': array([b'torquemada', b'torquemada disambiguation'], dtype=object), 'normalized_matched_wiki_entity_name': b'', 'normalized_value': b'torquemada', 'type': b'WikipediaEntity', 'value': b'Torquemada'}, 'entity_pages': {'doc_source': array([], dtype=object), 'filename': array([], dtype=object), 'title': array([], dtype=object), 'wiki_context': array([], dtype=object)}, 'question': b'In 1483, who was appointed the first grand inquisitor of the Spanish Inquisition?', 'question_id': b'qw_16011', 'question_source': b'http://www.quizwise.com/', 'search_results': {'description': array([], dtype=object), 'filename': array([], dtype=object), 'rank': array([], dtype=int32), 'search_context': array([], dtype=object), 'title': array([], dtype=object), 'url': array([], dtype=object)}}\n",
            "{'answer': {'aliases': array([b'Austerlitz (disambiguation)', b'Austerlitz', b'AUSTERLITZ'],\n",
            "      dtype=object), 'matched_wiki_entity_name': b'', 'normalized_aliases': array([b'austerlitz', b'austerlitz disambiguation'], dtype=object), 'normalized_matched_wiki_entity_name': b'', 'normalized_value': b'austerlitz', 'type': b'WikipediaEntity', 'value': b'AUSTERLITZ'}, 'entity_pages': {'doc_source': array([], dtype=object), 'filename': array([], dtype=object), 'title': array([], dtype=object), 'wiki_context': array([], dtype=object)}, 'question': b'Which celebrated battle was fought near Brno on 2nd December 1805?', 'question_id': b'dpql_4053', 'question_source': b'https://derbyshirepubquizleague.wordpress.com/', 'search_results': {'description': array([], dtype=object), 'filename': array([], dtype=object), 'rank': array([], dtype=int32), 'search_context': array([], dtype=object), 'title': array([], dtype=object), 'url': array([], dtype=object)}}\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gq5U_rjDb1bn"
      },
      "source": [
        "As with Natural Questions, we need to preprocess the raw examples into `inputs` and `targets` features. We can reuse the `trivia_preprocessor` above, but first we need to convert the TriviaQA examples into the correct format, ignoring the fields we don't need for our task.\n",
        "\n",
        "We'll then define our `Task` and print out a few preprocessed examples from the validation set.\n",
        "\n",
        "Note that we do not need to specify the splits or number of examples since that information is provided by TFDS."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "6rU32DjyeLuL",
        "outputId": "5f643cf8-c997-4201-d1da-a6c0beda473d",
        "colab": {
          "height": 258
        }
      },
      "source": [
        "def tiviaqa_extract_qa(ds):\n",
        "  def exract_qa(ex):\n",
        "    return {\n",
        "        \"question\": ex[\"question\"],\n",
        "        \"answer\": ex[\"answer\"][\"value\"]\n",
        "    }\n",
        "  return ds.map(exract_qa, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "t5.data.TaskRegistry.add(\n",
        "    \"triviaqa_context_free\",\n",
        "    # A TfdsTask takes in a TFDS name instead of a tf.data.Dataset function.\n",
        "    t5.data.TfdsTask,\n",
        "    tfds_name=\"trivia_qa/unfiltered.nocontext:1.1.0\",\n",
        "    tfds_data_dir=DATA_DIR,\n",
        "    sentencepiece_model_path=t5.data.DEFAULT_SPM_PATH,\n",
        "    text_preprocessor=[tiviaqa_extract_qa, trivia_preprocessor],\n",
        "    postprocess_fn=t5.data.postprocessors.lower_text,\n",
        "    metric_fns=[t5.evaluation.metrics.accuracy]\n",
        ")\n",
        "\n",
        "# Load and print a few examples.\n",
        "triviaqa_task = t5.data.TaskRegistry.get(\"triviaqa_context_free\")\n",
        "ds = triviaqa_task.get_dataset(split=\"validation\", sequence_length={\"inputs\": 128, \"targets\": 32})\n",
        "print(\"A few preprocessed validation examples...\")\n",
        "for ex in tfds.as_numpy(ds.take(3)):\n",
        "  print(ex)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "A few preprocessed validation examples...\n",
            "{'inputs_plaintext': b'trivia question: what does a farrier do?', 'inputs': array([22377,   822,    10,   125,   405,     3,     9,   623,  6711,\n",
            "         103,    58,     1]), 'targets_plaintext': b'he shoes horses', 'targets': array([    3,    88,  4439, 10235,     1])}\n",
            "{'inputs_plaintext': b'trivia question: what is the name of the wooden panelled lining applied to a room', 'inputs': array([22377,   822,    10,   125,    19,     8,   564,    13,     8,\n",
            "        5726,  2952,  1361,     3,  9424,  2930,    12,     3,     9,\n",
            "         562,     1]), 'targets_plaintext': b'wainscotting', 'targets': array([    3,   210, 13676, 10405,    53,     1])}\n",
            "{'inputs_plaintext': b'trivia question: how did gus grissom, ed white and roger b. chaffee die in 1967?', 'inputs': array([22377,   822,    10,   149,   410,     3,  1744,     7, 19116,\n",
            "       10348,     6,     3,    15,    26,   872,    11,     3,  3822,\n",
            "          49,     3,   115,     5,     3,  3441,  7398,    15,    67,\n",
            "          16, 18148,    58,     1]), 'targets_plaintext': b'burned to death', 'targets': array([16644,    12,  1687,     1])}\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "wlghm_3rAd-M"
      },
      "source": [
        "## Dataset Mixture\n",
        "\n",
        "We now create a `Mixture` from the above `Tasks`, which we will fine-tune on.\n",
        "\n",
        "There are different ways to automatically set the rate (for example, based on the number of examples using `rate_num_examples`), but we will just hardcode an equal mixture for simplicity."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "Zgs-s3eDAU37",
        "colab": {}
      },
      "source": [
        "t5.data.MixtureRegistry.remove(\"trivia_all\")\n",
        "t5.data.MixtureRegistry.add(\n",
        "    \"trivia_all\",\n",
        "    [\"nq_context_free\", \"triviaqa_context_free\"],\n",
        "     default_rate=1.0\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "CUkorodCENGw"
      },
      "source": [
        "# Transferring to new Tasks\n",
        "\n",
        "We are now ready to fine-tune one of the pre-trained T5 models on our new mixture of closed-book QA tasks.\n",
        "\n",
        "First, we'll instantiate a `Model` object using the model size of your choice. Note that larger models are slower to train and use but will likely achieve higher accuracy. You also may be able to increase accuracy by training longer with more `FINETUNE_STEPS` below.\n",
        "\n",
        "\n",
        "## Caveats\n",
        "\n",
        "* Due to its memory requirements, you will not be able to train the `11B` parameter model on the TPU provided by Colab. Instead, you will need to fine-tune inside of a GCP instance (see [README](https://github.com/google-research/text-to-text-transfer-transformer/)).\n",
        "* Due to the checkpoint size, you will not be able use the 5GB GCS free tier for the `3B` parameter models. You will need at least 25GB of space, which you can purchase with your $300 of initial credit on GCP.\n",
        "* While `large` can achieve decent results, it is recommended that you fine-tune at least the `3B` parameter model.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "syte5n0nnMOC"
      },
      "source": [
        "## Define Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "colab_type": "code",
        "id": "yGQ-zpgy3raf",
        "colab": {}
      },
      "source": [
        "MODEL_SIZE = \"3B\" #@param[\"small\", \"base\", \"large\", \"3B\", \"11B\"]\n",
        "# Public GCS path for T5 pre-trained model checkpoints\n",
        "BASE_PRETRAINED_DIR = \"gs://t5-data/pretrained_models\"\n",
        "PRETRAINED_DIR = os.path.join(BASE_PRETRAINED_DIR, MODEL_SIZE)\n",
        "MODEL_DIR = os.path.join(MODELS_DIR, MODEL_SIZE)\n",
        "\n",
        "if ON_CLOUD and MODEL_SIZE == \"3B\":\n",
        "  tf.logging.warn(\n",
        "      \"The `3B` model is too large to use with the 5GB GCS free tier. \"\n",
        "      \"Make sure you have at least 25GB on GCS before continuing.\"\n",
        "  )\n",
        "elif ON_CLOUD and MODEL_SIZE == \"11B\":\n",
        "  raise ValueError(\n",
        "      \"The `11B` parameter is too large to fine-tune on the `v2-8` TPU \"\n",
        "      \"provided by Colab. Please comment out this Error if you're running \"\n",
        "      \"on a larger TPU.\"\n",
        "  )\n",
        "\n",
        "# Set parallelism and batch size to fit on v2-8 TPU (if possible).\n",
        "# Limit number of checkpoints to fit within 5GB (if possible).\n",
        "model_parallelism, train_batch_size, keep_checkpoint_max = {\n",
        "    \"small\": (1, 256, 16),\n",
        "    \"base\": (2, 128, 8),\n",
        "    \"large\": (8, 64, 4),\n",
        "    \"3B\": (8, 16, 1),\n",
        "    \"11B\": (8, 16, 1)}[MODEL_SIZE]\n",
        "\n",
        "tf.io.gfile.makedirs(MODEL_DIR)\n",
        "# The models from our paper are based on the Mesh Tensorflow Transformer.\n",
        "model = t5.models.MtfModel(\n",
        "    model_dir=MODEL_DIR,\n",
        "    tpu=TPU_ADDRESS,\n",
        "    tpu_topology=TPU_TOPOLOGY,\n",
        "    model_parallelism=model_parallelism,\n",
        "    batch_size=train_batch_size,\n",
        "    sequence_length={\"inputs\": 128, \"targets\": 32},\n",
        "    learning_rate_schedule=0.003,\n",
        "    save_checkpoints_steps=5000,\n",
        "    keep_checkpoint_max=keep_checkpoint_max if ON_CLOUD else None,\n",
        "    iterations_per_loop=100,\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "dInuo63ZQrFi"
      },
      "source": [
        "Before we continue, let's load a [TensorBoard](https://www.tensorflow.org/tensorboard) visualizer so that we can keep monitor our progress. The page should automatically update as fine-tuning and evaluation proceed."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "M5mPyYATNsVT",
        "colab": {}
      },
      "source": [
        "if ON_CLOUD:\n",
        "  %reload_ext tensorboard\n",
        "  import tensorboard as tb\n",
        "tb.notebook.start(\"--logdir \" + MODELS_DIR)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DZhAd0U_4B_o"
      },
      "source": [
        "## Fine-tune\n",
        "\n",
        "We are now ready to fine-tune our model. This will take a while (~2 hours with default settings), so please be patient! The larger the model and more `FINETUNE_STEPS` you use, the longer it will take.\n",
        "\n",
        "Don't worry, you can always come back later and increase the number of steps, and it will automatically pick up where you left off."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "V7t7a25LBTj9",
        "colab": {}
      },
      "source": [
        "FINETUNE_STEPS = 25000 #@param {type: \"integer\"}\n",
        "\n",
        "model.finetune(\n",
        "    mixture_or_task_name=\"trivia_all\",\n",
        "    pretrained_model_dir=PRETRAINED_DIR,\n",
        "    finetune_steps=FINETUNE_STEPS\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-pFvyrHmm6Mx"
      },
      "source": [
        "## Expected Results [SPOILER ALERT]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "i_-7qYemnEHl"
      },
      "source": [
        "Below are the expected accuracies on the Natural Question (NQ) and TriviQA validation sets for various model sizes. The full 11B model produces the exact text of the answer 34.5% and 25.1% of the time on TriviaQA and NQ, respectively. The 3B parameter model, which is the largest that can be trained with a free Cloud TPU in Colab, achieves 29.7% and 23.7%, respectively.\n",
        "\n",
        "In reality, the model performs better than this since requiring exact match is too strict of a metric, as you’ll see in the examples below. This helps to explain why the model appears to perform better on TriviaQA than NQ, as the latter tends to include more long-form answers extracted from the context.\n",
        "\n",
        "Please see our [paper on closed-book QA](https://tiny.cc/t5-qa) where achieved even better results.\n",
        "\n",
        "<img src=\"https://storage.googleapis.com/t5-data/assets/t5_trivia_expected.png\">"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "eYeciUZ_D7T2"
      },
      "source": [
        "## Evaluate\n",
        "\n",
        "We now evaluate on the validation sets of the tasks in our mixture. Accuracy results will be logged and added to the TensorBoard above."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "bz6CJRHzNfd3",
        "colab": {}
      },
      "source": [
        "# Use a larger batch size for evaluation, which requires less memory.\n",
        "model.batch_size = train_batch_size * 4\n",
        "model.eval(\n",
        "    mixture_or_task_name=\"trivia_all\",\n",
        "    checkpoint_steps=\"all\"\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "92dClA1SWwIx"
      },
      "source": [
        "Let's look at a few random predictions from the validation sets. Note that we measure accuracy based on an *exact match* of the predicted answer and the ground-truth answer. As a result, some of the answers are semantically correct but are counted wrong by the exact match score."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "-FuqHRuvxOct",
        "outputId": "e36475d1-a8e7-4f2f-b4fa-6a2a2dd0c607",
        "colab": {
          "height": 1000
        }
      },
      "source": [
        "def print_random_predictions(task_name, n=10):\n",
        "  \"\"\"Print n predictions from the validation split of a task.\"\"\"\n",
        "  # Grab the dataset for this task.\n",
        "  ds = t5.data.TaskRegistry.get(task_name).get_dataset(\n",
        "      split=\"validation\",\n",
        "      sequence_length={\"inputs\": 128, \"targets\": 32},\n",
        "      shuffle=False)\n",
        "\n",
        "  def _prediction_file_to_ckpt(path):\n",
        "    \"\"\"Extract the global step from a prediction filename.\"\"\"\n",
        "    return int(path.split(\"_\")[-2])\n",
        "\n",
        "  # Grab the paths of all logged predictions.\n",
        "  prediction_files = tf.io.gfile.glob(\n",
        "      os.path.join(\n",
        "          MODEL_DIR,\n",
        "          \"validation_eval/%s_*_predictions\" % task_name))\n",
        "  # Get most recent prediction file by sorting by their step.\n",
        "  latest_prediction_file = sorted(\n",
        "      prediction_files, key=_prediction_file_to_ckpt)[-1]\n",
        "\n",
        "  # Collect (inputs, targets, prediction) from the dataset and predictions file\n",
        "  results = []\n",
        "  with tf.io.gfile.GFile(latest_prediction_file) as preds:\n",
        "    for ex, pred in zip(tfds.as_numpy(ds), preds):\n",
        "      results.append((tf.compat.as_text(ex[\"inputs_plaintext\"]),\n",
        "                      tf.compat.as_text(ex[\"targets_plaintext\"]),\n",
        "                      pred.strip()))\n",
        "\n",
        "  print(\"<== Random predictions for %s using checkpoint %s ==>\\n\" %\n",
        "        (task_name, \n",
        "         _prediction_file_to_ckpt(latest_prediction_file)))\n",
        "\n",
        "  for inp, tgt, pred in random.choices(results, k=10):\n",
        "    print(\"Input:\", inp)\n",
        "    print(\"Target:\", tgt)\n",
        "    print(\"Prediction:\", pred)\n",
        "    print(\"Counted as Correct?\", tgt == pred)\n",
        "    print()\n",
        "\n",
        "print_random_predictions(\"triviaqa_context_free\")\n",
        "print_random_predictions(\"nq_context_free\")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "<== Random predictions for triviaqa_context_free using checkpoint 1100000 ==>\n",
            "\n",
            "Input: trivia question: jackpot counter, ghost drop and drop zone are all terms used in which uk television game show?\n",
            "Target: tipping point\n",
            "Prediction: countdown\n",
            "Counted as Correct? False\n",
            "\n",
            "Input: trivia question: cursed to sail around the cape of good hope, which ghost ship is the theme of an 1841 opera by richard wagner?\n",
            "Target: the flying dutchman\n",
            "Prediction: baron von munchhausen\n",
            "Counted as Correct? False\n",
            "\n",
            "Input: trivia question: at what fret are found the same notes as the open strings, but an octave higher, on a standard guitar?\n",
            "Target: 12th\n",
            "Prediction: 12th\n",
            "Counted as Correct? True\n",
            "\n",
            "Input: trivia question: how many legs does a ladybird have?\n",
            "Target: six\n",
            "Prediction: six\n",
            "Counted as Correct? True\n",
            "\n",
            "Input: trivia question: in which city’s harbour was the ship queen elizabeth ravaged by fire in 1972?\n",
            "Target: hong kong\n",
            "Prediction: hong kong\n",
            "Counted as Correct? True\n",
            "\n",
            "Input: trivia question: what are the three largest islands in the world beginning with the letter n\n",
            "Target: new guinea, north island\n",
            "Prediction: new zealand; namibia and nova scotia\n",
            "Counted as Correct? False\n",
            "\n",
            "Input: trivia question: lenny bruce was in what field of entertainment in the 1960s?\n",
            "Target: standup comedy\n",
            "Prediction: comedy\n",
            "Counted as Correct? False\n",
            "\n",
            "Input: trivia question: in which sea are the cayman islands?\n",
            "Target: caribbean\n",
            "Prediction: caribbean\n",
            "Counted as Correct? True\n",
            "\n",
            "Input: trivia question: what is an astronomical event that occurs when one celestial object moves into the shadow of another?\n",
            "Target: eclipse\n",
            "Prediction: lunar eclipse\n",
            "Counted as Correct? False\n",
            "\n",
            "Input: trivia question: which tv cartoon series was about a meek janitor who led a double life as an unfortunate super-detective?\n",
            "Target: hong kong fuey\n",
            "Prediction: scooby-doo\n",
            "Counted as Correct? False\n",
            "\n",
            "<== Random predictions for nq_context_free using checkpoint 1100000 ==>\n",
            "\n",
            "Input: trivia question: who is known as the super fast boy in the series the icredible?\n",
            "Target: dashiell robert parr/dash\n",
            "Prediction: dash\n",
            "Counted as Correct? False\n",
            "\n",
            "Input: trivia question: who played santa in the santa clause movies?\n",
            "Target: tim allen\n",
            "Prediction: tim allen\n",
            "Counted as Correct? True\n",
            "\n",
            "Input: trivia question: who has sold more albums kelly or carrie?\n",
            "Target: carrie underwood\n",
            "Prediction: carrie underwood\n",
            "Counted as Correct? True\n",
            "\n",
            "Input: trivia question: when did sweet caroline start at red sox games?\n",
            "Target: at least 1997\n",
            "Prediction: 2004\n",
            "Counted as Correct? False\n",
            "\n",
            "Input: trivia question: who plays mr wilson in dennis the menace?\n",
            "Target: joseph sherrard kearns\n",
            "Prediction: joseph sherrard kearns\n",
            "Counted as Correct? True\n",
            "\n",
            "Input: trivia question: who had a baby at 100 in the bible?\n",
            "Target: abraham\n",
            "Prediction: sarah\n",
            "Counted as Correct? False\n",
            "\n",
            "Input: trivia question: who is doing 2018 super bowl half time show?\n",
            "Target: justin timberlake\n",
            "Prediction: justin timberlake\n",
            "Counted as Correct? True\n",
            "\n",
            "Input: trivia question: what is the official slogan for the 2018 winter olympics?\n",
            "Target: passion. connected.\n",
            "Prediction: every step counts\n",
            "Counted as Correct? False\n",
            "\n",
            "Input: trivia question: ray charles hit the road jack album name?\n",
            "Target: ray charles greatest hits\n",
            "Prediction: the road jack album\n",
            "Counted as Correct? False\n",
            "\n",
            "Input: trivia question: who sang the theme song to step by step?\n",
            "Target: jesse frederick james conaway\n",
            "Prediction: frederick and teresa james\n",
            "Counted as Correct? False\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "vbqiq2Ab4PJk"
      },
      "source": [
        "## Predict\n",
        "\n",
        "Now that we have fine-tuned the model, we can feed T5 arbitrary questions and have it predict the answers!\n",
        "\n",
        "There is a significant amount of overhead in initializing the model so this may take a few minutes to run each time even though the prediction itself is quite fast.\n",
        "\n",
        "\n",
        "To avoid this overhead, you might consider exporting a `SavedModel` and running it on [Cloud ML Engine](https://cloud.google.com/ml-engine/).\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "colab_type": "code",
        "id": "xatHPuCJsPns",
        "outputId": "18e6b677-cf57-4d45-b16d-ca015937dcf9",
        "colab": {
          "height": 340
        }
      },
      "source": [
        "question_1 = \"Where is the Google headquarters located?\" #@param {type:\"string\"}\n",
        "question_2 = \"What is the most populous country in the world?\" #@param {type:\"string\"}\n",
        "question_3 = \"Who are the 4 members of The Beatles?\" #@param {type:\"string\"}\n",
        "question_4 = \"How many teeth do humans have?\" #@param {type:\"string\"}\n",
        "\n",
        "questions = [question_1, question_2, question_3, question_4]\n",
        "\n",
        "now = time.time()\n",
        "# Write out the supplied questions to text files.\n",
        "predict_inputs_path = os.path.join(MODEL_DIR, \"predict_inputs_%d.txt\" % now)\n",
        "predict_outputs_path = os.path.join(MODEL_DIR, \"predict_outputs_%d.txt\" % now)\n",
        "# Manually apply preprocessing by prepending \"triviaqa question:\".\n",
        "with tf.io.gfile.GFile(predict_inputs_path, \"w\") as f:\n",
        "  for q in questions:\n",
        "    f.write(\"trivia question: %s\\n\" % q.lower())\n",
        "\n",
        "# Ignore any logging so that we only see the model's answers to the questions.\n",
        "with tf_verbosity_level('ERROR'):\n",
        "  model.batch_size = len(questions)\n",
        "  model.predict(\n",
        "      input_file=predict_inputs_path,\n",
        "      output_file=predict_outputs_path,\n",
        "      # Select the most probable output token at each step.\n",
        "      temperature=0,\n",
        "  )\n",
        "\n",
        "# The output filename will have the checkpoint appended so we glob to get \n",
        "# the latest.\n",
        "prediction_files = sorted(tf.io.gfile.glob(predict_outputs_path + \"*\"))\n",
        "print(\"\\nPredictions using checkpoint %s:\\n\" % prediction_files[-1].split(\"-\")[-1])\n",
        "with tf.io.gfile.GFile(prediction_files[-1]) as f:\n",
        "  for q, a in zip(questions, f):\n",
        "    if q:\n",
        "      print(\"Q: \" + q)\n",
        "      print(\"A: \" + a)\n",
        "      print()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Predictions using checkpoint 1100000:\n",
            "\n",
            "Q: Where is the Google headquarters located?\n",
            "A: mountain view, california\n",
            "\n",
            "\n",
            "Q: What is the most populous country in the world?\n",
            "A: china\n",
            "\n",
            "\n",
            "Q: Who are the 4 members of The Beatles?\n",
            "A: john lennon, paul mccartney, george harrison and ringo starr\n",
            "\n",
            "\n",
            "Q: How many teeth do humans have?\n",
            "A: 30\n",
            "\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xxElhphBZMD5"
      },
      "source": [
        "# Export Model for Serving\n",
        "\n",
        "As mentioned in the previous section, exporting a [`SavedModel`](https://www.tensorflow.org/guide/saved_model) can be useful for improving performance during inference or allowing your model to be deployed on a variety of platforms (e.g., TFLite, TensorFlow.js, TensorFlow Serving, or TensorFlow Hub)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "l_YuEL9FZ-UR"
      },
      "source": [
        "## Export SavedModel\n",
        "\n",
        "We first export the SavedModel. We set a batch size of 1 for simplicity, but it may be more efficient to use a larger batch size if you want to handle multiple requests per call.\n",
        "\n",
        "For 3B and 11B models the export will take approximately 30-45 minutes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "eWu8lbh3aHjc",
        "colab": {}
      },
      "source": [
        "export_dir = os.path.join(MODEL_DIR, \"export\")\n",
        "\n",
        "model.batch_size = 1 # make one prediction per call\n",
        "model.export(\n",
        "    export_dir,\n",
        "    checkpoint_step=-1,  # use most recent\n",
        "    beam_size=1,  # no beam search\n",
        "    temperature=1.0,  # sample according to predicted distribution\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "h82x5mSQaKjk",
        "colab": {}
      },
      "source": [
        "# Find the most-recently exported model.\n",
        "saved_model_path = os.path.dirname(\n",
        "    max(tf.io.gfile.glob(os.path.join(MODEL_DIR, \"export\", \"*\", \"saved_model.pb\")))\n",
        ")\n",
        "print(\"Model saved to:\", saved_model_path)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "PZ8WXpFkaNTP"
      },
      "source": [
        "## Load SavedModel\n",
        "\n",
        "One way to test our model is to load it in a TF 1.x `session` so that we can repeatedly predict from the model without the overhead of loading the graph and weights on the TPU each time.\n",
        "\n",
        "We pay the overhead once here, but it shouldn't take more than a few minutes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "1TpeMFGhaN7r",
        "colab": {}
      },
      "source": [
        "import tensorflow.compat.v1 as tf\n",
        "tf.reset_default_graph()\n",
        "sess = tf.Session(TPU_ADDRESS)\n",
        "meta_graph_def = tf.saved_model.loader.load(sess, [\"serve\"], saved_model_path)\n",
        "signature_def = meta_graph_def.signature_def[\"serving_default\"]\n",
        "\n",
        "def answer(question):\n",
        "  return sess.run(\n",
        "      fetches=signature_def.outputs[\"outputs\"].name, \n",
        "      feed_dict={signature_def.inputs[\"input\"].name: [question]}\n",
        "  )[0].decode('utf-8')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "YbGC8xefaYtV"
      },
      "source": [
        "## Predict\n",
        "\n",
        "We can now call the predict method with different inputs each time and relatively quickly get results."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "3WA0BYI9abgv",
        "colab": {}
      },
      "source": [
        "for question in [\"trivia question: where is the google headquarters?\",\n",
        "                 \"trivia question: what is the most populous country in the world?\",\n",
        "                 \"trivia question: who are the 4 members of the beatles?\",\n",
        "                 \"trivia question: how many teeth do humans have?\"]:\n",
        "    print(answer(question))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ePXPEQhDafmV"
      },
      "source": [
        "## Deploy SavedModel\n",
        "\n",
        "You can now deploy your SavedModel for serving (e.g., with [TensorFlow Serving](https://www.tensorflow.org/tfx/tutorials/serving/rest_simple))."
      ]
    }
  ]
}