T5_SQuAD_Model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "C4W3_Colab_T5_SQuAD_Model.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/anshudaur/DeepLearning/blob/master/T5_SQuAD_Model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7yuytuIllsv1"
      },
      "source": [
        "# Assignment 3, Part 2: T5 SQuAD Model \n",
        "\n",
        "Welcome to the part 2 of testing the models for this week's assignment. This time we will perform decoding using the T5 SQuAD model. In this notebook we'll perform Question Answering by providing a \"Question\", its \"Context\" and see how well we get the \"Target\" answer. \n",
        "\n",
        "## IMPORTANT\n",
        "\n",
        "- As you cannot save the changes you make to this colab, you have to make a copy of this notebook in your own drive and run that. You can do so by going to `File -> Save a copy in Drive`. Close this colab and open the copy which you have made in your own drive.\n",
        "\n",
        "- Go to this [google drive folder](https://drive.google.com/drive/folders/1rOZsbEzcpMRVvgrRULRh1JPFpkIG_JOz?usp=sharing) named `NLP C4 W3 Data`. In the folder, next to its name use the drop down menu to select `\"Add shortcut to Drive\" -> \"My Drive\" and then press ADD SHORTCUT`. This should add a shortcut to the folder `NLP C4 W3 Data` within your own google drive. Please make sure this happens, as you'll be reading the data for this notebook from this folder.\n",
        "\n",
        "- Make sure your runtime is GPU (_not_ CPU or TPU). And if it is an option, make sure you are using _Python 3_. You can select these settings by going to `Runtime -> Change runtime type -> Select the above mentioned settings and then press SAVE`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Db6LQW5cMSgx"
      },
      "source": [
        "**Note: Restarting the runtime maybe required**.\n",
        "\n",
        "Colab will tell you if the restarting is necessary -- you can do this from the:\n",
        "\n",
        "Runtime > Restart Runtime\n",
        "\n",
        "option in the dropdown."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OTeAcYQo_4OW"
      },
      "source": [
        "## Outline\n",
        "\n",
        "- [Part 0: Downloading and loading dependencies](#0)\n",
        "- [Part 1: Mounting your drive for data accessibility](#1)\n",
        "- [Part 2: Getting things ready](#2)\n",
        "- [Part 3: Fine-tuning on SQuAD](#3)\n",
        "    - [3.1 Loading in the data and preprocessing](#3.1)\n",
        "    - [3.2 Decoding from a fine-tuned model](#3.2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ysxogfC1M158"
      },
      "source": [
        "<a name='0'></a>\n",
        "# Part 0: Downloading and loading dependencies\n",
        "\n",
        "Uncomment the code cell below and run it to download some dependencies that you will need. You need to download them once every time you open the colab. You can ignore the `kfac` error."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1BNZzCg0xv3R"
      },
      "source": [
        "#!pip -q install trax"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uDhi6qLQMHzs"
      },
      "source": [
        "import string\n",
        "import t5\n",
        "import numpy as np\n",
        "import trax \n",
        "from trax.supervised import decoding\n",
        "import textwrap \n",
        "# Will come handy later.\n",
        "wrapper = textwrap.TextWrapper(width=70)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cwr7LoXwQUW5"
      },
      "source": [
        "<a name='1'></a>\n",
        "# Part 1: Mounting your drive for data accessibility\n",
        "\n",
        "Run the code cell below and follow the instructions to mount your drive. The data is the same as used in the coursera version of the assignment."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P7ZF7KiXzQEg"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive/', force_remount=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HuTyft5EBQK6"
      },
      "source": [
        "<a name='2'></a>\n",
        "# Part 2: Getting things ready \n",
        "\n",
        "Run the code cell below to ready some functions which will later help us in decoding. The code and the functions are the same as the ones you previsouly ran on the coursera version of the assignment."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LJ8s_xZ1QtkI"
      },
      "source": [
        "PAD, EOS, UNK = 0, 1, 2\n",
        "\n",
        "def detokenize(np_array):\n",
        "  return trax.data.detokenize(\n",
        "      np_array,\n",
        "      vocab_type='sentencepiece',\n",
        "      vocab_file='sentencepiece.model',\n",
        "      vocab_dir='/content/drive/My Drive/NLP C4 W3 Data/')\n",
        "\n",
        "def tokenize(s):\n",
        "  # The trax.data.tokenize function operates on streams,\n",
        "  # that's why we have to create 1-element stream with iter\n",
        "  # and later retrieve the result with next.\n",
        "  return next(trax.data.tokenize(\n",
        "      iter([s]),\n",
        "      vocab_type='sentencepiece',\n",
        "      vocab_file='sentencepiece.model',\n",
        "      vocab_dir='/content/drive/My Drive/NLP C4 W3 Data/'))\n",
        " \n",
        "vocab_size = trax.data.vocab_size(\n",
        "    vocab_type='sentencepiece',\n",
        "    vocab_file='sentencepiece.model',\n",
        "    vocab_dir='/content/drive/My Drive/NLP C4 W3 Data/')\n",
        "\n",
        "def get_sentinels(vocab_size):\n",
        "    sentinels = {}\n",
        "\n",
        "    for i, char in enumerate(reversed(string.ascii_letters), 1):\n",
        "\n",
        "        decoded_text = detokenize([vocab_size - i]) \n",
        "        \n",
        "        # Sentinels, ex: <Z> - <a>\n",
        "        sentinels[decoded_text] = f'<{char}>'\n",
        "        \n",
        "    return sentinels\n",
        "\n",
        "sentinels = get_sentinels(vocab_size)    \n",
        "\n",
        "\n",
        "def pretty_decode(encoded_str_list, sentinels=sentinels):\n",
        "    # If already a string, just do the replacements.\n",
        "    if isinstance(encoded_str_list, (str, bytes)):\n",
        "        for token, char in sentinels.items():\n",
        "            encoded_str_list = encoded_str_list.replace(token, char)\n",
        "        return encoded_str_list\n",
        "  \n",
        "    # We need to decode and then prettyfy it.\n",
        "    return pretty_decode(detokenize(encoded_str_list))    "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HEoSSKNwgDVA"
      },
      "source": [
        "<a name='3'></a>\n",
        "# Part 3: Fine-tuning on SQuAD\n",
        "\n",
        "Now let's try to fine tune on SQuAD and see what becomes of the model. For this, we need to write a function that will create and process the SQuAD `tf.data.Dataset`. Below is how T5 pre-processes SQuAD dataset as a text2text example. Before we jump in, we will have to first load in the data. \n",
        "\n",
        "<a name='3.1'></a>\n",
        "### 3.1 Loading in the data and preprocessing\n",
        "\n",
        "You first start by loading in the dataset. The text2text example for a SQuAD example looks like:\n",
        "\n",
        "```\n",
        "{\n",
        "  'inputs': 'question: <question> context: <article>',\n",
        "  'targets': '<answer_0>',\n",
        "}\n",
        "```\n",
        "\n",
        "The squad pre-processing function takes in the dataset and processes it using the sentencePiece vocabulary you have seen above. It generates the features from the vocab and encodes the string features. It takes on question, context, and answer, and returns \"question: Q context: C\" as input and \"A\" as target."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RcdR5Dh9UVEw"
      },
      "source": [
        "# Retrieve Question, C, A and return \"question: Q context: C\" as input and \"A\" as target.\n",
        "def squad_preprocess_fn(dataset, mode='train'):\n",
        "  return t5.data.preprocessors.squad(dataset)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0NCIejwexv4x"
      },
      "source": [
        "# train generator, this takes about 1 minute\n",
        "train_generator_fn, eval_generator_fn = trax.data.tf_inputs.data_streams(\n",
        "  'squad/v1.1:2.0.0',\n",
        "  data_dir='/content/drive/My Drive/NLP C4 W3 Data/data/',\n",
        "  bare_preprocess_fn=squad_preprocess_fn,\n",
        "  input_name='inputs',\n",
        "  target_name='targets'\n",
        ")\n",
        "train_generator = train_generator_fn()\n",
        "next(train_generator)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QGQsExH8xv40"
      },
      "source": [
        "#print example from train_generator\n",
        "(inp, out) = next(train_generator)\n",
        "print(inp.decode('utf8').split('context:')[0])\n",
        "print()\n",
        "print('context:', inp.decode('utf8').split('context:')[1])\n",
        "print()\n",
        "print('target:', out.decode('utf8'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F_eGJHpPh-rz"
      },
      "source": [
        "<a name='3.2'></a>\n",
        "### 3.2 Decoding from a fine-tuned model\n",
        "\n",
        "You will now use an existing model that we trained for you. You will initialize, then load in your model, and then try with your own input. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dZztX5XLa01p"
      },
      "source": [
        "# Initialize the model \n",
        "model = trax.models.Transformer(\n",
        "    d_ff = 4096,\n",
        "    d_model = 1024,\n",
        "    max_len = 2048,\n",
        "    n_heads = 16,\n",
        "    dropout = 0.1,\n",
        "    input_vocab_size = 32000,\n",
        "    n_encoder_layers = 24,\n",
        "    n_decoder_layers = 24,\n",
        "    mode='predict')  # Change to 'eval' for slow decoding."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OQRAGG2fnSMe"
      },
      "source": [
        "# load in the model\n",
        "# this will take a minute\n",
        "shape11 = trax.shapes.ShapeDtype((1, 1), dtype=np.int32)\n",
        "model.init_from_file('/content/drive/My Drive/NLP C4 W3 Data/models/model_squad.pkl.gz',\n",
        "                     weights_only=True, input_signature=(shape11, shape11))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l33umWrRyYNm"
      },
      "source": [
        "# Uncomment to see the transformer's structure.\n",
        "# print(model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zzrxtGGDuUkc"
      },
      "source": [
        "# create inputs\n",
        "# a simple example \n",
        "# inputs = 'question: She asked him where is john? context: John was at the game'\n",
        "\n",
        "# an extensive example\n",
        "inputs = 'question: What are some of the colours of a rose? context: A rose is a woody perennial flowering plant of the genus Rosa, in the family Rosaceae, or the flower it bears.There are over three hundred species and tens of thousands of cultivars. They form a group of plants that can be erect shrubs, climbing, or trailing, with stems that are often armed with sharp prickles. Flowers vary in size and shape and are usually large and showy, in colours ranging from white through yellows and reds. Most species are native to Asia, with smaller numbers native to Europe, North America, and northwestern Africa. Species, cultivars and hybrids are all widely grown for their beauty and often are fragrant.'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nVlu-iwHcZRY"
      },
      "source": [
        "# tokenizing the input so we could feed it for decoding\n",
        "print(tokenize(inputs))\n",
        "test_inputs = tokenize(inputs) "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PJh_Qw9G5jND"
      },
      "source": [
        "# Temperature is a parameter for sampling.\n",
        "#   # * 0.0: same as argmax, always pick the most probable token\n",
        "#   # * 1.0: sampling from the distribution (can sometimes say random things)\n",
        "#   # * values inbetween can trade off diversity and quality, try it out!\n",
        "output = decoding.autoregressive_sample(model, inputs=np.array(test_inputs)[None, :],\n",
        "                                        temperature=0.0, max_length=10)\n",
        "print(wrapper.fill(pretty_decode(output[0])))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BadBty9C-8wY"
      },
      "source": [
        "### Note: As you can see the RAM is almost full, it is because the model and the decoding is memory heavy. You can run decoding just once. Running it the second time with another example might give you the same answer as before, or not run at all (crash). If that happens restart the runtime (see how to at the start of the notebook) and run all the cells again."
      ]
    }
  ]
}