C4W3_Assignment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f368f78e",
   "metadata": {
    "colab_type": "text",
    "id": "7yuytuIllsv1"
   },
   "source": [
    "# Assignment 3: Question Answering\n",
    "\n",
    "Welcome to the third assignment of course 4. In this assignment you will explore question answering. You will implement the \"Text to Text Transfer from Transformers\" (better known as T5). Since you implemented transformers from scratch last week you will now be able to use them. \n",
    "\n",
    "<img src = \"images/qa.png\"> \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf3ec561",
   "metadata": {
    "colab_type": "text",
    "id": "Db6LQW5cMSgx"
   },
   "source": [
    "## Table of Contents\n",
    "\n",
    "- [Overview](#0-1)\n",
    "- [Importing the Packages](#0-2)\n",
    "- [1 - Prepare the data for pretraining T5](#1)\n",
    "    - [1.1 - Pre-Training Objective](#1-1)\n",
    "    - [1.2 - C4 Dataset](#1-2)\n",
    "    - [1.3 - Process C4](#1-3)\n",
    "    - [1.4 - Decode to Natural Language](#1-4)\n",
    "    - [1.5 - Tokenizing and Masking](#1-5)\n",
    "        - [Exercise 1 - tokenize_and_mask](#ex-1)\n",
    "    - [1.6 - Creating the Pairs](#1-6)\n",
    "- [2 - Pretrain a T5 model using C4](#2)\n",
    "    - [2.1 - Instantiate a new transformer model](#2-1)\n",
    "    - [2.2 - C4 pretraining](#2-2)\n",
    "- [3 - Fine tune the T5 model for Question Answering](#3)\n",
    "    - [3.1 - Creating a list of paired question and answers](#3-1)\n",
    "        - [Exercise 2 - Parse the SQuaD 2.0 dataset](#ex-2)\n",
    "    - [3.2 - Fine tune the T5 model](#3-2)    \n",
    "    - [3.3 - Implement your Question Answering model](#3-3)\n",
    "        - [Exercise 3 - Implement the question answering function](#ex-3)    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0595e9c4",
   "metadata": {
    "colab_type": "text",
    "id": "ysxogfC1M158"
   },
   "source": [
    "<a name='0-1'></a>\n",
    "## Overview\n",
    "\n",
    "This assignment will be different from the two previous ones. Due to memory constraints of this environment and for the sake of time, your model will be trained with small datasets, so you won't get models that you could use in production but you will gain the necessary knowledge about how the Generative Language models are trained and used. Also you won't spend too much time with the architecture of the models but you will instead take a model that is pre-trained on a larger dataset and fine tune it to get better results.\n",
    "\n",
    "After completing this labs you will:\n",
    "* Understand how the C4 dataset is structured.\n",
    "* Pretrain a transformer model using a Masked Language Model.\n",
    "* Understand how the \"Text to Text Transfer from Transformers\" or T5 model works. \n",
    "* Fine tune the T5 model for Question answering\n",
    "\n",
    "Before getting started take some time to read the following tips:\n",
    "#### TIPS FOR SUCCESSFUL GRADING OF YOUR ASSIGNMENT:\n",
    "- All cells are frozen except for the ones where you need to submit your solutions.\n",
    "- You can add new cells to experiment but these will be omitted by the grader, so don't rely on newly created cells to host your solution code, use the provided places for this.\n",
    "- You can add the comment # grade-up-to-here in any graded cell to signal the grader that it must only evaluate up to that point. This is helpful if you want to check if you are on the right track even if you are not done with the whole assignment. Be sure to remember to delete the comment afterwards!\n",
    "- To submit your notebook, save it and then click on the blue submit button at the beginning of the page."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2156cf78",
   "metadata": {},
   "source": [
    "<a name='0-2'></a>\n",
    "## Importing the Packages\n",
    "\n",
    "Let's start by importing all the required libraries. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3a532381",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "deletable": false,
    "editable": false,
    "id": "uDhi6qLQMHzs",
    "outputId": "64947d91-eef3-425b-9b4b-7ca7cefcc823",
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'\n",
    "\n",
    "import traceback\n",
    "import time\n",
    "import json\n",
    "from termcolor import colored\n",
    "import string\n",
    "import textwrap\n",
    "import itertools\n",
    "import numpy as np\n",
    "import tensorflow_text as tf_text\n",
    "import tensorflow as tf\n",
    "\n",
    "import transformer_utils \n",
    "import utils\n",
    "\n",
    "# Will come in handy later\n",
    "wrapper = textwrap.TextWrapper(width=70)\n",
    "\n",
    "# Set random seed\n",
    "np.random.seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bf711eba",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import w3_unittest"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47bea693",
   "metadata": {
    "colab_type": "text",
    "id": "t7A-LAxsYpDd",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "<a name='1'></a>\n",
    "## 1 -  Prepare the data for pretraining T5 \n",
    "\n",
    "<a name='1-1'></a>\n",
    "### 1.1 - Pre-Training Objective\n",
    "\n",
    "In the initial phase of training a T5 model for a Question Answering task, the pre-training process involves leveraging a masked language model (MLM) on a very large dataset, such as the C4 dataset. The objective is to allow the model to learn contextualized representations of words and phrases, fostering a deeper understanding of language semantics. To initiate pre-training, it is essential to employ the Transformer architecture, which forms the backbone of T5. The Transformer's self-attention mechanism enables the model to weigh different parts of the input sequence dynamically, capturing long-range dependencies effectively.\n",
    "\n",
    "Before delving into pre-training, thorough data preprocessing is crucial. The C4 dataset, a diverse and extensive collection of web pages, provides a rich source for language understanding tasks. The dataset needs to be tokenized into smaller units, such as subwords or words, to facilitate model input. Additionally, the text is often segmented into fixed-length sequences or batches, optimizing computational efficiency during training.\n",
    "\n",
    "For the masked language modeling objective, a percentage of the tokenized input is randomly masked, and the model is trained to predict the original content of these masked tokens. This process encourages the T5 model to grasp contextual relationships between words and phrases, enhancing its ability to generate coherent and contextually appropriate responses during downstream tasks like question answering.\n",
    "\n",
    "In summary, the pre-training of the T5 model involves utilizing the Transformer architecture on a sizable dataset like C4, coupled with meticulous data preprocessing to convert raw text into a format suitable for training. The incorporation of a masked language modeling objective ensures that the model learns robust contextual representations, laying a solid foundation for subsequent fine-tuning on specific tasks such as question answering.\n",
    "\n",
    "**Note:** The word \"mask\" will be used throughout this assignment in context of hiding/removing word(s)\n",
    "\n",
    "You will be implementing the Masked language model (MLM) as shown in the following image. \n",
    "\n",
    "<img src = \"images/loss.png\" width=\"600\" height = \"400\">\n",
    "\n",
    "Assume you have the following text: <span style = \"color:blue\"> **Thank you <span style = \"color:red\">for inviting </span> me to your party <span style = \"color:red\">last</span>  week** </span> \n",
    "\n",
    "\n",
    "Now as input you will mask the words in red in the text: \n",
    "\n",
    "<span style = \"color:blue\"> **Input:**</span> Thank you  **X** me to your party **Y** week.\n",
    "\n",
    "<span style = \"color:blue\">**Output:**</span> The model should predict the words(s) for **X** and **Y**. \n",
    "\n",
    "**[EOS]** will be used to mark the end of the target sequence."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dc25302",
   "metadata": {
    "colab_type": "text",
    "id": "Cwr7LoXwQUW5",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "<a name='1-2'></a>\n",
    "### 1.2 - C4 Dataset\n",
    "\n",
    "The [C4 dataset](https://www.tensorflow.org/datasets/catalog/c4), also known as the Common Crawl C4 (Common Crawl Corpus C4), is a large-scale dataset of web pages collected by the [Common Crawl organization](https://commoncrawl.org/). It is commonly used for various natural language processing tasks and machine learning research. Each sample in the C4 dataset follows a consistent format, making it suitable for pretraining models like BERT. Here's a short explanation and description of the C4 dataset:\n",
    "\n",
    "- Format: Each sample in the C4 dataset is represented as a JSON object, containing several key-value pairs.\n",
    "\n",
    "- Content: The 'text' field in each sample contains the actual text content extracted from web pages. This text often includes a wide range of topics and writing styles, making it diverse and suitable for training language models.\n",
    "\n",
    "- Metadata: The dataset includes metadata such as 'content-length,' 'content-type,' 'timestamp,' and 'url,' providing additional information about each web page. 'Content-length' specifies the length of the content, 'content-type' describes the type of content (e.g., 'text/plain'), 'timestamp' indicates when the web page was crawled, and 'url' provides the source URL of the web page.\n",
    "\n",
    "- Applications: The C4 dataset is commonly used for training and fine-tuning large-scale language models, such as BERT. It serves as a valuable resource for tasks like text classification, named entity recognition, question answering, and more.\n",
    "\n",
    "- Size: The C4 dataset is containing more than 800 GiB of text data, making it suitable for training models with billions of parameters.\n",
    "\n",
    "Run the cell below to see how the C4 dataset looks like. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "aa56acc9",
   "metadata": {
    "deletable": false,
    "editable": false,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "example number 1: \n",
      "\n",
      "{'text': 'Beginners BBQ Class Taking Place in Missoula!\\nDo you want to get better at making delicious BBQ? You will have the opportunity, put this on your calendar now. Thursday, September 22nd join World Class BBQ Champion, Tony Balay from Lonestar Smoke Rangers. He will be teaching a beginner level class for everyone who wants to get better with their culinary skills.\\nHe will teach you everything you need to know to compete in a KCBS BBQ competition, including techniques, recipes, timelines, meat selection and trimming, plus smoker and fire information.\\nThe cost to be in the class is $35 per person, and for spectators it is free. Included in the cost will be either a t-shirt or apron and you will be tasting samples of each meat that is prepared.'} \n",
      "\n",
      "example number 2: \n",
      "\n",
      "{'text': 'Discussion in \\'Mac OS X Lion (10.7)\\' started by axboi87, Jan 20, 2012.\\nI\\'ve got a 500gb internal drive and a 240gb SSD.\\nWhen trying to restore using disk utility i\\'m given the error \"Not enough space on disk ____ to restore\"\\nBut I shouldn\\'t have to do that!!!\\nAny ideas or workarounds before resorting to the above?\\nUse Carbon Copy Cloner to copy one drive to the other. I\\'ve done this several times going from larger HDD to smaller SSD and I wound up with a bootable SSD drive. One step you have to remember not to skip is to use Disk Utility to partition the SSD as GUID partition scheme HFS+ before doing the clone. If it came Apple Partition Scheme, even if you let CCC do the clone, the resulting drive won\\'t be bootable. CCC usually works in \"file mode\" and it can easily copy a larger drive (that\\'s mostly empty) onto a smaller drive. If you tell CCC to clone a drive you did NOT boot from, it can work in block copy mode where the destination drive must be the same size or larger than the drive you are cloning from (if I recall).\\nI\\'ve actually done this somehow on Disk Utility several times (booting from a different drive (or even the dvd) so not running disk utility from the drive your cloning) and had it work just fine from larger to smaller bootable clone. Definitely format the drive cloning to first, as bootable Apple etc..\\nThanks for pointing this out. My only experience using DU to go larger to smaller was when I was trying to make a Lion install stick and I was unable to restore InstallESD.dmg to a 4 GB USB stick but of course the reason that wouldn\\'t fit is there was slightly more than 4 GB of data.'} \n",
      "\n",
      "example number 3: \n",
      "\n",
      "{'text': 'Foil plaid lycra and spandex shortall with metallic slinky insets. Attached metallic elastic belt with O-ring. Headband included. Great hip hop or jazz dance costume. Made in the USA.'} \n",
      "\n",
      "example number 4: \n",
      "\n",
      "{'text': \"How many backlinks per day for new site?\\nDiscussion in 'Black Hat SEO' started by Omoplata, Dec 3, 2010.\\n1) for a newly created site, what's the max # backlinks per day I should do to be safe?\\n2) how long do I have to let my site age before I can start making more blinks?\\nI did about 6000 forum profiles every 24 hours for 10 days for one of my sites which had a brand new domain.\\nThere is three backlinks for every of these forum profile so thats 18 000 backlinks every 24 hours and nothing happened in terms of being penalized or sandboxed. This is now maybe 3 months ago and the site is ranking on first page for a lot of my targeted keywords.\\nbuild more you can in starting but do manual submission and not spammy type means manual + relevant to the post.. then after 1 month you can make a big blast..\\nWow, dude, you built 18k backlinks a day on a brand new site? How quickly did you rank up? What kind of competition/searches did those keywords have?\"} \n",
      "\n",
      "example number 5: \n",
      "\n",
      "{'text': 'The Denver Board of Education opened the 2017-18 school year with an update on projects that include new construction, upgrades, heat mitigation and quality learning environments.\\nWe are excited that Denver students will be the beneficiaries of a four year, $572 million General Obligation Bond. Since the passage of the bond, our construction team has worked to schedule the projects over the four-year term of the bond.\\nDenver voters on Tuesday approved bond and mill funding measures for students in Denver Public Schools, agreeing to invest $572 million in bond funding to build and improve schools and $56.6 million in operating dollars to support proven initiatives, such as early literacy.\\nDenver voters say yes to bond and mill levy funding support for DPS students and schools. Click to learn more about the details of the voter-approved bond measure.\\nDenver voters on Nov. 8 approved bond and mill funding measures for DPS students and schools. Learn more about what’s included in the mill levy measure.'} \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Load example jsons\n",
    "with open('data/c4-en-10k.jsonl', 'r') as file:\n",
    "    example_jsons = [json.loads(line.strip()) for line in file]\n",
    "\n",
    "# Printing the examples to see how the data looks like\n",
    "for i in range(5):\n",
    "    print(f'example number {i+1}: \\n\\n{example_jsons[i]} \\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48901d97",
   "metadata": {
    "colab_type": "text",
    "id": "eeihIgtiaSfh",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "<a name='1-3'></a>\n",
    "### 1.3 - Process C4\n",
    "\n",
    "For the purpose of pretaining the T5 model, you will only use the `content` of each entry. In the following code, you filter only the field `text` from all the entries in the dataset. This is the data that you will use to create the `inputs` and `targets` of your language model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "af728cb2",
   "metadata": {
    "deletable": false,
    "editable": false,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Beginners BBQ Class Taking Place in Missoula!\n",
      "Do you want to get better at making delicious BBQ? You will have the opportunity, put this on your calendar now. Thursday, September 22nd join World Class BBQ Champion, Tony Balay from Lonestar Smoke Rangers. He will be teaching a beginner level class for everyone who wants to get better with their culinary skills.\n",
      "He will teach you everything you need to know to compete in a KCBS BBQ competition, including techniques, recipes, timelines, meat selection and trimming, plus smoker and fire information.\n",
      "The cost to be in the class is $35 per person, and for spectators it is free. Included in the cost will be either a t-shirt or apron and you will be tasting samples of each meat that is prepared.\n"
     ]
    }
   ],
   "source": [
    "# Grab text field from dictionary\n",
    "natural_language_texts = [example_json['text'] for example_json in example_jsons]\n",
    "\n",
    "# Print the first text example\n",
    "print(natural_language_texts[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee4a25a2",
   "metadata": {
    "colab_type": "text",
    "id": "1rMrONRqcCYi",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "<a name='1-4'></a>\n",
    "### 1.4 - Decode to Natural Language\n",
    "\n",
    "The [SentencePieceTokenizer](https://www.tensorflow.org/text/api_docs/python/text/SentencepieceTokenizer), used in the code snippet, tokenizes text into subword units, enhancing handling of complex word structures, out-of-vocabulary words, and multilingual support. It simplifies preprocessing, ensures consistent tokenization, and seamlessly integrates with machine learning frameworks.\n",
    "\n",
    "In this task, a SentencePiece model is loaded from a file, which is used to tokenize text into subwords represented by integer IDs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2ac53d57",
   "metadata": {
    "deletable": false,
    "editable": false,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Special tokens\n",
    "# PAD, EOS = 0, 1\n",
    "\n",
    "with open(\"./models/sentencepiece.model\", \"rb\") as f:\n",
    "    pre_trained_tokenizer = f.read()\n",
    "    \n",
    "tokenizer = tf_text.SentencepieceTokenizer(pre_trained_tokenizer, out_type=tf.int32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "658b0e86",
   "metadata": {},
   "source": [
    "In this tokenizer the string `</s>` is used as `EOS` token. By default, the tokenizer does not add the `EOS` to the end of each sentence, so you need to add it manually when required. Let's verify what id correspond to this token:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7d2fec4b",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EOS: 1\n"
     ]
    }
   ],
   "source": [
    "eos = tokenizer.string_to_id(\"</s>\").numpy()\n",
    "\n",
    "print(\"EOS: \" + str(eos))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e87756f",
   "metadata": {},
   "source": [
    "This code shows the process of tokenizing individual words from a given text, in this case, the first entry of the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "83c48352",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 54
    },
    "colab_type": "code",
    "deletable": false,
    "id": "iCCjgiVZgTSK",
    "outputId": "023a227c-d895-4fd9-ae83-9394fe48cebd",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Word\t\t-->\tTokenization\n",
      "----------------------------------------\n",
      "Foil    \t-->\t[4452, 173]\n",
      "plaid   \t-->\t[30772]\n",
      "lycra   \t-->\t[3, 120, 2935]\n",
      "and     \t-->\t[11]\n",
      "spandex \t-->\t[8438, 26, 994]\n",
      "shortall\t-->\t[710, 1748]\n",
      "with    \t-->\t[28]\n",
      "metallic\t-->\t[18813]\n",
      "slinky  \t-->\t[3, 7, 4907, 63]\n",
      "insets. \t-->\t[16, 2244, 7, 5]\n",
      "Attached\t-->\t[28416, 15, 26]\n",
      "metallic\t-->\t[18813]\n",
      "elastic \t-->\t[15855]\n",
      "belt    \t-->\t[6782]\n",
      "with    \t-->\t[28]\n",
      "O-ring. \t-->\t[411, 18, 1007, 5]\n",
      "Headband\t-->\t[3642, 3348]\n",
      "included.\t-->\t[1285, 5]\n",
      "Great   \t-->\t[1651]\n",
      "hip     \t-->\t[5436]\n",
      "hop     \t-->\t[13652]\n",
      "or      \t-->\t[42]\n",
      "jazz    \t-->\t[9948]\n",
      "dance   \t-->\t[2595]\n",
      "costume.\t-->\t[11594, 5]\n",
      "Made    \t-->\t[6465]\n",
      "in      \t-->\t[16]\n",
      "the     \t-->\t[8]\n",
      "USA.    \t-->\t[2312, 5]\n"
     ]
    }
   ],
   "source": [
    "# printing the encoding of each word to see how subwords are tokenized\n",
    "tokenized_text = [(list(tokenizer.tokenize(word).numpy()), word) for word in natural_language_texts[2].split()]\n",
    "\n",
    "print(\"Word\\t\\t-->\\tTokenization\")\n",
    "print(\"-\"*40)\n",
    "for element in tokenized_text:\n",
    "    print(f\"{element[1]:<8}\\t-->\\t{element[0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4616cf3",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "And as usual, the library provides a function to turn numeric tokens into human readable text. Look how it works. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "92d7037b",
   "metadata": {
    "deletable": false,
    "editable": false,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tokenized: [12847   277]\n",
      "detokenized: b'Beginners'\n"
     ]
    }
   ],
   "source": [
    "# We can see that detokenize successfully undoes the tokenization\n",
    "print(f\"tokenized: {tokenizer.tokenize('Beginners')}\\ndetokenized: {tokenizer.detokenize(tokenizer.tokenize('Beginners'))}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52f63624",
   "metadata": {
    "colab_type": "text",
    "id": "vPKgGOeOxv3w",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "As you can see above, you were able to take a piece of string and tokenize it. \n",
    "\n",
    "Now you will create `input` and `target` pairs that will allow you to train your model. T5 uses the ids at the end of the vocab file as sentinels. For example, it will replace: \n",
    "   - `vocab_size - 1` by `<Z>`\n",
    "   - `vocab_size - 2` by `<Y>`\n",
    "   - and so forth. \n",
    "   \n",
    "It assigns every word a `chr`.\n",
    "\n",
    "The `pretty_decode` function below, which you will use in a bit, helps in handling the type when decoding. Take a look and try to understand what the function is doing.\n",
    "\n",
    "\n",
    "Notice that:\n",
    "```python\n",
    "string.ascii_letters = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'\n",
    "```\n",
    "\n",
    "**NOTE:** Targets may have more than the 52 sentinels we replace, but this is just to give you an idea of things."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b25bb46d",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "deletable": false,
    "editable": false,
    "id": "fCPQL5FTxv3w",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_sentinels(tokenizer, display=False):\n",
    "    sentinels = {}\n",
    "    vocab_size = tokenizer.vocab_size(name=None)\n",
    "    for i, char in enumerate(reversed(string.ascii_letters), 1):\n",
    "        decoded_text = tokenizer.detokenize([vocab_size - i]).numpy().decode(\"utf-8\")\n",
    "        \n",
    "        # Sentinels, ex: <Z> - <a>\n",
    "        sentinels[decoded_text] = f'<{char}>'    \n",
    "    \n",
    "        if display:\n",
    "            print(f'The sentinel is <{char}> and the decoded token is:', decoded_text)\n",
    "\n",
    "    return sentinels\n",
    "\n",
    "def pretty_decode(encoded_str_list, sentinels, tokenizer):\n",
    "    # If already a string, just do the replacements.\n",
    "    if tf.is_tensor(encoded_str_list) and encoded_str_list.dtype == tf.string:\n",
    "        for token, char in sentinels.items():\n",
    "            encoded_str_list = tf.strings.regex_replace(encoded_str_list, token, char)\n",
    "        return encoded_str_list\n",
    "  \n",
    "    # We need to decode and then prettyfy it.\n",
    "    return pretty_decode(tokenizer.detokenize(encoded_str_list), sentinels, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "56d75b6c",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "deletable": false,
    "editable": false,
    "id": "fCPQL5FTxv3w",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The sentinel is <Z> and the decoded token is: Internațional\n",
      "The sentinel is <Y> and the decoded token is: erwachsene\n",
      "The sentinel is <X> and the decoded token is: Cushion\n",
      "The sentinel is <W> and the decoded token is: imunitar\n",
      "The sentinel is <V> and the decoded token is: Intellectual\n",
      "The sentinel is <U> and the decoded token is: traditi\n",
      "The sentinel is <T> and the decoded token is: disguise\n",
      "The sentinel is <S> and the decoded token is: exerce\n",
      "The sentinel is <R> and the decoded token is: nourishe\n",
      "The sentinel is <Q> and the decoded token is: predominant\n",
      "The sentinel is <P> and the decoded token is: amitié\n",
      "The sentinel is <O> and the decoded token is: erkennt\n",
      "The sentinel is <N> and the decoded token is: dimension\n",
      "The sentinel is <M> and the decoded token is: inférieur\n",
      "The sentinel is <L> and the decoded token is: refugi\n",
      "The sentinel is <K> and the decoded token is: cheddar\n",
      "The sentinel is <J> and the decoded token is: unterlieg\n",
      "The sentinel is <I> and the decoded token is: garanteaz\n",
      "The sentinel is <H> and the decoded token is: făcute\n",
      "The sentinel is <G> and the decoded token is: réglage\n",
      "The sentinel is <F> and the decoded token is: pedepse\n",
      "The sentinel is <E> and the decoded token is: Germain\n",
      "The sentinel is <D> and the decoded token is: distinctly\n",
      "The sentinel is <C> and the decoded token is: Schraub\n",
      "The sentinel is <B> and the decoded token is: emanat\n",
      "The sentinel is <A> and the decoded token is: trimestre\n",
      "The sentinel is <z> and the decoded token is: disrespect\n",
      "The sentinel is <y> and the decoded token is: Erasmus\n",
      "The sentinel is <x> and the decoded token is: Australia\n",
      "The sentinel is <w> and the decoded token is: permeabil\n",
      "The sentinel is <v> and the decoded token is: deseori\n",
      "The sentinel is <u> and the decoded token is: manipulated\n",
      "The sentinel is <t> and the decoded token is: suggér\n",
      "The sentinel is <s> and the decoded token is: corespund\n",
      "The sentinel is <r> and the decoded token is: nitro\n",
      "The sentinel is <q> and the decoded token is: oyons\n",
      "The sentinel is <p> and the decoded token is: Account\n",
      "The sentinel is <o> and the decoded token is: échéan\n",
      "The sentinel is <n> and the decoded token is: laundering\n",
      "The sentinel is <m> and the decoded token is: genealogy\n",
      "The sentinel is <l> and the decoded token is: QuickBooks\n",
      "The sentinel is <k> and the decoded token is: constituted\n",
      "The sentinel is <j> and the decoded token is: Fertigung\n",
      "The sentinel is <i> and the decoded token is: goutte\n",
      "The sentinel is <h> and the decoded token is: regulă\n",
      "The sentinel is <g> and the decoded token is: overwhelmingly\n",
      "The sentinel is <f> and the decoded token is: émerg\n",
      "The sentinel is <e> and the decoded token is: broyeur\n",
      "The sentinel is <d> and the decoded token is: povești\n",
      "The sentinel is <c> and the decoded token is: emulator\n",
      "The sentinel is <b> and the decoded token is: halloween\n",
      "The sentinel is <a> and the decoded token is: combustibil\n"
     ]
    }
   ],
   "source": [
    "sentinels = get_sentinels(tokenizer, display=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be73a35d",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fCPQL5FTxv3w",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Now, let's use the `pretty_decode` function in the following sentence. Note that all the words listed as sentinels, will be replaced by the function with the corresponding sentinel. It could be a drawback of this method, but don't worry about it now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1fe92253",
   "metadata": {
    "deletable": false,
    "editable": false,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=string, numpy=b'I want to dress up as an <V> this <b>.'>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pretty_decode(tf.constant(\"I want to dress up as an Intellectual this halloween.\"), sentinels, tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "559b04b7",
   "metadata": {
    "colab_type": "text",
    "id": "Y64F--Nzxv30",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "The functions above make your `inputs` and `targets` more readable. For example, you might see something like this once you implement the masking function below. \n",
    "\n",
    "- <span style=\"color:red\"> Input sentence: </span> Younes and Lukasz were working together in the lab yesterday after lunch. \n",
    "- <span style=\"color:red\">Input: </span> Younes and Lukasz  **Z** together in the **Y** yesterday after lunch.\n",
    "- <span style=\"color:red\">Target: </span> **Z** were working **Y** lab.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "244cd7a8",
   "metadata": {
    "colab_type": "text",
    "id": "NvvNd7n6xv30"
   },
   "source": [
    "<a name='1-5'></a>\n",
    "### 1.5 - Tokenizing and Masking\n",
    "\n",
    "In this task, you will implement the `tokenize_and_mask` function, which tokenizes and masks input words based on a given probability. The probability is controlled by the `noise` parameter, typically set to mask around `15%` of the words in the input text. The function will generate two lists of tokenized sequences following the algorithm outlined below:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7050f25c",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "<a name='ex-1'></a>\n",
    "### Exercise 1 - tokenize_and_mask\n",
    "\n",
    "- Start with two empty lists: `inps` and `targs`\n",
    "- Tokenize the input text using the given tokenizer.\n",
    "- For each `token` in the tokenized sequence:\n",
    "  - Generate a random number(simulating a weighted coin toss)\n",
    "  - If the random value is greater than the given threshold(noise):\n",
    "    - Add the current token to the `inps` list\n",
    "  - Else:\n",
    "    - If a new sentinel must be included(read note **):\n",
    "      - Compute the next sentinel ID using a progression.\n",
    "      - Add a sentinel into the `inps` and `targs` to mark the position of the masked element.\n",
    "    - Add the current token to the `targs` list.\n",
    "\n",
    "** There's a special case to consider. If two or more consecutive tokens get masked during the process, you don't need to add a new sentinel to the sequences. To account for this, use the `prev_no_mask` flag, which starts as `True` but is turned to `False` each time you mask a new element. The code that adds sentinels will only be executed if, before masking the token, the flag was in the `True` state.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c660bf97",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "deletable": false,
    "id": "Bi33WKgRxv31",
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# GRADED FUNCTION: tokenize_and_mask\n",
    "def tokenize_and_mask(text, \n",
    "                      noise=0.15, \n",
    "                      randomizer=np.random.uniform, \n",
    "                      tokenizer=None):\n",
    "    \"\"\"Tokenizes and masks a given input.\n",
    "\n",
    "    Args:\n",
    "        text (str or bytes): Text input.\n",
    "        noise (float, optional): Probability of masking a token. Defaults to 0.15.\n",
    "        randomizer (function, optional): Function that generates random values. Defaults to np.random.uniform.\n",
    "        tokenizer (function, optional): Tokenizer function. Defaults to tokenize.\n",
    "\n",
    "    Returns:\n",
    "        inps, targs: Lists of integers associated to inputs and targets.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Current sentinel number (starts at 0)\n",
    "    cur_sentinel_num = 0\n",
    "    \n",
    "    # Inputs and targets\n",
    "    inps, targs = [], []\n",
    "\n",
    "    # Vocab_size\n",
    "    vocab_size = int(tokenizer.vocab_size())\n",
    "    \n",
    "    # EOS token id \n",
    "    # Must be at the end of each target!\n",
    "    eos = tokenizer.string_to_id(\"</s>\").numpy()\n",
    "    \n",
    "    ### START CODE HERE ###\n",
    "    \n",
    "    # prev_no_mask is True if the previous token was NOT masked, False otherwise\n",
    "    # set prev_no_mask to True\n",
    "    prev_no_mask = True\n",
    "    \n",
    "    # Loop over the tokenized text\n",
    "    for token in tokenizer.tokenize(text).numpy():\n",
    "        \n",
    "        # Generate a random value between 0 and 1\n",
    "        rnd_val = randomizer() \n",
    "        \n",
    "        # Check if the noise is greater than a random value (weighted coin flip)\n",
    "        if noise > rnd_val:\n",
    "            \n",
    "            # Check if previous token was NOT masked\n",
    "            if prev_no_mask:\n",
    "                \n",
    "                # Current sentinel increases by 1\n",
    "                cur_sentinel_num += 1\n",
    "                \n",
    "                # Compute end_id by subtracting current sentinel value out of the total vocabulary size\n",
    "                end_id = vocab_size - cur_sentinel_num\n",
    "                \n",
    "                # Append end_id at the end of the targets\n",
    "                targs.append(end_id)\n",
    "                \n",
    "                # Append end_id at the end of the inputs\n",
    "                inps.append(end_id)\n",
    "                \n",
    "            # Append token at the end of the targets\n",
    "            targs.append(token)\n",
    "            \n",
    "            # set prev_no_mask accordingly\n",
    "            prev_no_mask = False\n",
    "\n",
    "        else:\n",
    "            \n",
    "            # Append token at the end of the inputs\n",
    "            inps.append(token)\n",
    "            \n",
    "            # Set prev_no_mask accordingly\n",
    "            prev_no_mask = True\n",
    "    \n",
    "    \n",
    "    # Add EOS token to the end of the targets\n",
    "    targs.append(eos)\n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "    \n",
    "    return inps, targs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e92edca1",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 122
    },
    "colab_type": "code",
    "deletable": false,
    "editable": false,
    "id": "OlPySQo9xv34",
    "outputId": "2b0dc5e4-8d58-4eb0-a146-0c9f158264ac",
    "scrolled": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tokenized inputs - shape=53:\n",
      "\n",
      "[31999, 15068, 4501, 3, 12297, 3399, 16, 5964, 7115, 31998, 531, 25, 241, 12, 129, 394, 44, 492, 31997, 58, 148, 56, 43, 8, 1004, 6, 474, 31996, 39, 4793, 230, 5, 2721, 6, 1600, 1630, 31995, 1150, 4501, 15068, 16127, 6, 9137, 2659, 5595, 31994, 782, 3624, 14627, 15, 12612, 277, 5]\n",
      "\n",
      "targets - shape=19:\n",
      "\n",
      "[31999, 12847, 277, 31998, 9, 55, 31997, 3326, 15068, 31996, 48, 30, 31995, 727, 1715, 31994, 45, 301, 1]\n"
     ]
    }
   ],
   "source": [
    "# Some logic to mock a np.random value generator\n",
    "# Needs to be in the same cell for it to always generate same output\n",
    "def testing_rnd():\n",
    "    def dummy_generator():\n",
    "        vals = np.linspace(0, 1, 10)\n",
    "        cyclic_vals = itertools.cycle(vals)\n",
    "        for _ in range(100):\n",
    "            yield next(cyclic_vals)\n",
    "\n",
    "    dumr = itertools.cycle(dummy_generator())\n",
    "\n",
    "    def dummy_randomizer():\n",
    "        return next(dumr)\n",
    "    \n",
    "    return dummy_randomizer\n",
    "\n",
    "input_str = 'Beginners BBQ Class Taking Place in Missoula!\\nDo you want to get better at making delicious BBQ? You will have the opportunity, put this on your calendar now. Thursday, September 22nd join World Class BBQ Champion, Tony Balay from Lonestar Smoke Rangers.'\n",
    "\n",
    "inps, targs = tokenize_and_mask(input_str, randomizer=testing_rnd(), tokenizer=tokenizer)\n",
    "print(f\"tokenized inputs - shape={len(inps)}:\\n\\n{inps}\\n\\ntargets - shape={len(targs)}:\\n\\n{targs}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07996252",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "#### **Expected Output:**\n",
    "```\n",
    "tokenized inputs - shape=53:\n",
    "\n",
    "[31999 15068  4501     3 12297  3399    16  5964  7115 31998   531    25\n",
    "   241    12   129   394    44   492 31997    58   148    56    43     8\n",
    "  1004     6   474 31996    39  4793   230     5  2721     6  1600  1630\n",
    " 31995  1150  4501 15068 16127     6  9137  2659  5595 31994   782  3624\n",
    " 14627    15 12612   277     5]\n",
    "\n",
    "targets - shape=19:\n",
    "\n",
    "[31999 12847   277 31998     9    55 31997  3326 15068 31996    48    30\n",
    " 31995   727  1715 31994    45   301     1]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "76daaa5b",
   "metadata": {
    "deletable": false,
    "editable": false,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[92m All tests passed\n"
     ]
    }
   ],
   "source": [
    "# Test your implementation!\n",
    "w3_unittest.test_tokenize_and_mask(tokenize_and_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c87bea8",
   "metadata": {
    "colab_type": "text",
    "id": "_omCqbkLxv36",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "You will now use the inputs and the targets from the `tokenize_and_mask` function you implemented above. Take a look at the decoded version of your masked sentence using your `inps` and `targs` from the sentence above. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "054d51bf",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 105
    },
    "colab_type": "code",
    "deletable": false,
    "editable": false,
    "id": "y6xwo6lGxv37",
    "outputId": "4330ae1e-1805-40c9-daf3-c6bbe92d957b",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inputs: \n",
      "\n",
      " b'<Z> BBQ Class Taking Place in Missoul <Y> Do you want to get better at making <X>? You will have the opportunity, put <W> your calendar now. Thursday, September 22 <V> World Class BBQ Champion, Tony Balay <U>onestar Smoke Rangers.'\n",
      "\n",
      "Targets: \n",
      "\n",
      " b'<Z> Beginners <Y>a! <X> delicious BBQ <W> this on <V>nd join <U> from L'\n"
     ]
    }
   ],
   "source": [
    "print('Inputs: \\n\\n', pretty_decode(inps, sentinels, tokenizer).numpy())\n",
    "print('\\nTargets: \\n\\n', pretty_decode(targs, sentinels, tokenizer).numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0707c320",
   "metadata": {
    "colab_type": "text",
    "id": "24HZiIBLxv3-",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "<a name='1-6'></a>\n",
    "### 1.6 - Creating the Pairs\n",
    "\n",
    "You will now create pairs using your dataset. You will iterate over your data and create (inp, targ) pairs using the functions that we have given you. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ae83fff0",
   "metadata": {
    "deletable": false,
    "editable": false,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Apply tokenize_and_mask\n",
    "inputs_targets_pairs = [tokenize_and_mask(text.encode('utf-8', errors='ignore').decode('utf-8'), tokenizer=tokenizer) \n",
    "                        for text in natural_language_texts[0:2000]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3f157ad1",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "deletable": false,
    "editable": false,
    "id": "c1HiKreWokhs",
    "outputId": "fc194524-41de-4d3b-87d9-ae35c29c9f79",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1]\n",
      "\n",
      "inputs:\n",
      "<Z>il plaid <Y>lycra <X> spandex shortall with metallic slinky\n",
      "<W>sets. Attache <V> metallic elastic belt with O <U>ring. Head <T>\n",
      "included. Great hip hop<S> jazz dance costume.<R> in the USA.\n",
      "\n",
      "targets:\n",
      "<Z> Fo <Y>  <X> and <W> in <V>d <U>- <T>band<S> or<R> Made\n",
      "\n",
      "\n",
      "\n",
      "[2]\n",
      "\n",
      "inputs:\n",
      "I thought I was going to <Z> 3rd season <Y> Wire tonight. <X> there\n",
      "was a commentary <W> 11, so I had to re <V>watch <U> Ground with <T>\n",
      "commentary. Hopefully<S> can finish<R> season <Q>.\n",
      "\n",
      "targets:\n",
      "<Z> finish the <Y> of the <X> But <W> on episode <V>- <U> Middle <T>\n",
      "the<S> I<R> the <Q> next weekend\n",
      "\n",
      "\n",
      "\n",
      "[3]\n",
      "\n",
      "inputs:\n",
      "Pencarian <Z>FILM Untuk \" <Y>eace <X>er 2017 <W> yuk mampir ke channel\n",
      "say <V>. Edges <U> provides the l.. A corrupt cop makes one w.. <T>er\n",
      "2017  ⁇ <S> ⁇  .. Náo Lo ⁇ n - Peace Break.. Please subscribe and hit\n",
      "..<R> in HD at http://.. <Q> cannot believe I manage..\n",
      "\n",
      "targets:\n",
      "<Z>  <Y>P <X> Break <W>\" <V>. <U> East <T> Peace Break<S> <R> uploaded\n",
      "<Q> I\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def display_input_target_pairs(inputs_targets_pairs, sentinels, wrapper=textwrap.TextWrapper(width=70), tokenizer=tokenizer):\n",
    "    for i, inp_tgt_pair in enumerate(inputs_targets_pairs, 1):\n",
    "        inps, tgts = inp_tgt_pair\n",
    "        inps = str(pretty_decode(inps, sentinels, tokenizer).numpy(), encoding='utf-8')\n",
    "        tgts = str(pretty_decode(tgts, sentinels, tokenizer).numpy(), encoding='utf-8')\n",
    "        print(f'[{i}]\\n\\n'\n",
    "              f'inputs:\\n{wrapper.fill(text=inps)}\\n\\n'\n",
    "              f'targets:\\n{wrapper.fill(text=tgts)}\\n\\n\\n')\n",
    "\n",
    "# Print 3 samples. We print inputs with less than 100 tokens. It is just to give you and idea of the process\n",
    "display_input_target_pairs(filter(lambda x: len(x[0]) < 100, inputs_targets_pairs[0:12]), sentinels, wrapper, tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7d5e6d9",
   "metadata": {
    "colab_type": "text",
    "id": "hQI5Jgov5X-d",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "<a name='2'></a>\n",
    "## 2 - Pretrain a T5 model using C4\n",
    "\n",
    "Now you are going to use the Transformer's architecture that you coded in the previous assignment to summarize text, but this time to answer questions. Instead of training the question answering model from scratch, you will first \"pre-train\" the model using the C4 data set you just processed. This will help the model to learn the general structure of language from a large dataset. This is much easier to do, as you don't need to label any data, but just use the masking, which is done automatically. You will then use the data from the SQuAD set to teach the model to answer questions given a context. To start let's review the Transformer's architecture. \n",
    "\n",
    "<img src = \"images/fulltransformer.png\" width=\"300\" height=\"600\">\n",
    "\n",
    "<a name='2-1'></a>\n",
    "### 2.1 - Instantiate a new transformer model\n",
    "\n",
    "We have packaged the code implemented in the previous week into the `Transformer.py` file. You can import it here, and setup with the same configuration used there. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "58ce75dc",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "deletable": false,
    "editable": false,
    "id": "eScMhEG7xv4H",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define the model parameters\n",
    "num_layers = 2\n",
    "embedding_dim = 128\n",
    "fully_connected_dim = 128\n",
    "num_heads = 2\n",
    "positional_encoding_length = 256\n",
    "\n",
    "encoder_vocab_size = int(tokenizer.vocab_size())\n",
    "decoder_vocab_size = encoder_vocab_size\n",
    "\n",
    "# Initialize the model\n",
    "transformer = transformer_utils.Transformer(\n",
    "    num_layers, \n",
    "    embedding_dim, \n",
    "    num_heads, \n",
    "    fully_connected_dim,\n",
    "    encoder_vocab_size, \n",
    "    decoder_vocab_size, \n",
    "    positional_encoding_length, \n",
    "    positional_encoding_length,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "618697cf",
   "metadata": {},
   "source": [
    "Now, you will define the optimizer and the loss function. For this task the model will try to predict the masked words, so, as in the previous lab, the loss function will be the `SparseCategoricalCrossEntropy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7df2d1d1",
   "metadata": {
    "deletable": false,
    "editable": false,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "learning_rate = transformer_utils.CustomSchedule(embedding_dim)\n",
    "optimizer = tf.keras.optimizers.Adam(0.0001, beta_1=0.9, beta_2=0.98, epsilon=1e-9)\n",
    "\n",
    "loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')\n",
    "train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
    "\n",
    "# Here you will store the losses, so you can later plot them\n",
    "losses = []"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03d54376",
   "metadata": {},
   "source": [
    "<a name='2-2'></a>\n",
    "### 2.2 - C4 pretraining\n",
    "\n",
    "For training a Tensorflow model you need to arrange the data into datasets. Now, you will get the `inputs` and the `targets` for the transformer model from the `inputs_targets_pairs`. Before creating the dataset, you need to be sure that all `inputs` have the same length by truncating the longer sequences and padding the shorter ones with `0`. The same must be done for the targets. The function `tf.keras.preprocessing.sequence.pad_sequences` will help you here, as in the previous week assignment.\n",
    "\n",
    "You will use a `BATCH_SIZE = 64`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b03eb998",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Limit the size of the input and output data so this can run in this environment\n",
    "encoder_maxlen = 150\n",
    "decoder_maxlen = 50\n",
    "\n",
    "inputs = tf.keras.preprocessing.sequence.pad_sequences([x[0] for x in inputs_targets_pairs], maxlen=encoder_maxlen, padding='post', truncating='post')\n",
    "targets = tf.keras.preprocessing.sequence.pad_sequences([x[1] for x in inputs_targets_pairs], maxlen=decoder_maxlen, padding='post', truncating='post')\n",
    "\n",
    "inputs = tf.cast(inputs, dtype=tf.int32)\n",
    "targets = tf.cast(targets, dtype=tf.int32)\n",
    "\n",
    "# Create the final training dataset.\n",
    "BUFFER_SIZE = 10000\n",
    "BATCH_SIZE = 64\n",
    "\n",
    "dataset = tf.data.Dataset.from_tensor_slices((inputs, targets)).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e32ae0c",
   "metadata": {},
   "source": [
    "Now, you can run the training loop for 10 epochs. Running it with a big dataset such as C4 on a good computer with enough memory and a good GPU could take more than 24 hours. Here, you will run few epochs using a small portion of the C4 dataset for illustration. It will only take a few minutes, but the model won't be very powerful. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "44fc5f76",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss 10.1148\n",
      "Time taken for one epoch: 29.184126138687134 sec\n",
      "Epoch 2, Loss 9.5108\n",
      "Time taken for one epoch: 9.29906964302063 sec\n",
      "Epoch 3, Loss 8.9211\n",
      "Time taken for one epoch: 8.967788457870483 sec\n",
      "Epoch 4, Loss 8.3773\n",
      "Time taken for one epoch: 6.845104455947876 sec\n",
      "Epoch 5, Loss 7.8819\n",
      "Time taken for one epoch: 7.064520359039307 sec\n",
      "Epoch 6, Loss 7.4382\n",
      "Time taken for one epoch: 5.305743932723999 sec\n",
      "Epoch 7, Loss 7.0450\n",
      "Time taken for one epoch: 4.864109754562378 sec\n",
      "Epoch 8, Loss 6.7119\n",
      "Time taken for one epoch: 4.549200534820557 sec\n",
      "Epoch 9, Loss 6.4377\n",
      "Time taken for one epoch: 5.117376327514648 sec\n",
      "Epoch 10, Loss 6.2390\n",
      "Time taken for one epoch: 3.324408531188965 sec\n"
     ]
    }
   ],
   "source": [
    "# Define the number of epochs\n",
    "epochs = 10\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(epochs):\n",
    "    \n",
    "    start = time.time()\n",
    "    train_loss.reset_states()\n",
    "    number_of_batches=len(list(enumerate(dataset)))\n",
    "\n",
    "    for (batch, (inp, tar)) in enumerate(dataset):\n",
    "        print(f'Epoch {epoch+1}, Batch {batch+1}/{number_of_batches}', end='\\r')\n",
    "        transformer_utils.train_step(inp, tar, transformer, loss_object, optimizer, train_loss)\n",
    "    \n",
    "    print (f'Epoch {epoch+1}, Loss {train_loss.result():.4f}')\n",
    "    losses.append(train_loss.result())\n",
    "    \n",
    "    print (f'Time taken for one epoch: {time.time() - start} sec')\n",
    "\n",
    "# Save the pretrained model\n",
    "# transformer.save_weights('./model_c4_temp')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e8135b5",
   "metadata": {},
   "source": [
    "**Load a pretrained model**\n",
    "\n",
    "To show how powerful this model actually is, we trained it for several epochs with the full dataset in Colab and saved the weights for you. You can load them using the cell below. For the rest of the notebook, you will see the power of the transfer learning in action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "55360633",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.checkpoint.checkpoint.CheckpointLoadStatus at 0x7f406c2a8730>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transformer.load_weights('./pretrained_models/model_c4')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8822756",
   "metadata": {},
   "source": [
    "<a name='3'></a>\n",
    "## 3. Fine tune the T5 model for Question Answering\n",
    "\n",
    "Now,  you are going to fine tune the pretrained model for Question Answering using the [SQUad 2.0 dataset](https://rajpurkar.github.io/SQuAD-explorer/).\n",
    "\n",
    "SQuAD, short for Stanford Question Answering Dataset, is a dataset designed for training and evaluating question answering systems. It consists of real questions posed by humans on a set of Wikipedia articles, where the answer to each question is a specific span of text within the corresponding article.\n",
    "\n",
    "SQuAD 1.1, the previous version of the SQuAD dataset, contains 100,000+ question-answer pairs on about 500 articles.\n",
    "SQuAD 2.0, contains 50.000 additional questions that are not meant to be answered. This extra set of questions can help to train models to detect unanswerable questions.\n",
    "\n",
    "Let's load the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "987571df",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of articles: 442\n"
     ]
    }
   ],
   "source": [
    "with open('data/train-v2.0.json', 'r') as f:\n",
    "    example_jsons = json.load(f)\n",
    "\n",
    "example_jsons = example_jsons['data']\n",
    "\n",
    "print('Number of articles: ' + str(len(example_jsons)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c941761f",
   "metadata": {},
   "source": [
    "The structure of each article is as follows:\n",
    "- `title`: The article title\n",
    "- `paragraphs`: A list of paragraphs and questions related to them\n",
    "    - `context`: The actual paragraph text\n",
    "    - `qas`: A set of question related to the paragraph\n",
    "        - `question`: A question\n",
    "        - `id`: The question unique identifier\n",
    "        - `is_imposible`: Boolean, specifies if the question can be answered or not\n",
    "        - `answers`: A set of possible answers for the question\n",
    "            - `text`: The answer\n",
    "            - `answer_start`: The index of the character that starts the sentence containing the explicit answer to the question\n",
    "            \n",
    "Take a look at an article by running the next cell. Notice that the `context` is usually the last element for every paragraph:           "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7c4c4cfa",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Title: Beyoncé\n",
      "{'qas': [{'question': 'When did Beyonce start becoming popular?', 'id': '56be85543aeaaa14008c9063', 'answers': [{'text': 'in the late 1990s', 'answer_start': 269}], 'is_impossible': False}, {'question': 'What areas did Beyonce compete in when she was growing up?', 'id': '56be85543aeaaa14008c9065', 'answers': [{'text': 'singing and dancing', 'answer_start': 207}], 'is_impossible': False}, {'question': \"When did Beyonce leave Destiny's Child and become a solo singer?\", 'id': '56be85543aeaaa14008c9066', 'answers': [{'text': '2003', 'answer_start': 526}], 'is_impossible': False}, {'question': 'In what city and state did Beyonce  grow up? ', 'id': '56bf6b0f3aeaaa14008c9601', 'answers': [{'text': 'Houston, Texas', 'answer_start': 166}], 'is_impossible': False}, {'question': 'In which decade did Beyonce become famous?', 'id': '56bf6b0f3aeaaa14008c9602', 'answers': [{'text': 'late 1990s', 'answer_start': 276}], 'is_impossible': False}, {'question': 'In what R&B group was she the lead singer?', 'id': '56bf6b0f3aeaaa14008c9603', 'answers': [{'text': \"Destiny's Child\", 'answer_start': 320}], 'is_impossible': False}, {'question': 'What album made her a worldwide known artist?', 'id': '56bf6b0f3aeaaa14008c9604', 'answers': [{'text': 'Dangerously in Love', 'answer_start': 505}], 'is_impossible': False}, {'question': \"Who managed the Destiny's Child group?\", 'id': '56bf6b0f3aeaaa14008c9605', 'answers': [{'text': 'Mathew Knowles', 'answer_start': 360}], 'is_impossible': False}, {'question': 'When did Beyoncé rise to fame?', 'id': '56d43c5f2ccc5a1400d830a9', 'answers': [{'text': 'late 1990s', 'answer_start': 276}], 'is_impossible': False}, {'question': \"What role did Beyoncé have in Destiny's Child?\", 'id': '56d43c5f2ccc5a1400d830aa', 'answers': [{'text': 'lead singer', 'answer_start': 290}], 'is_impossible': False}, {'question': 'What was the first album Beyoncé released as a solo artist?', 'id': '56d43c5f2ccc5a1400d830ab', 'answers': [{'text': 'Dangerously in Love', 'answer_start': 505}], 'is_impossible': False}, {'question': 'When did Beyoncé release Dangerously in Love?', 'id': '56d43c5f2ccc5a1400d830ac', 'answers': [{'text': '2003', 'answer_start': 526}], 'is_impossible': False}, {'question': 'How many Grammy awards did Beyoncé win for her first solo album?', 'id': '56d43c5f2ccc5a1400d830ad', 'answers': [{'text': 'five', 'answer_start': 590}], 'is_impossible': False}, {'question': \"What was Beyoncé's role in Destiny's Child?\", 'id': '56d43ce42ccc5a1400d830b4', 'answers': [{'text': 'lead singer', 'answer_start': 290}], 'is_impossible': False}, {'question': \"What was the name of Beyoncé's first solo album?\", 'id': '56d43ce42ccc5a1400d830b5', 'answers': [{'text': 'Dangerously in Love', 'answer_start': 505}], 'is_impossible': False}], 'context': 'Beyoncé Giselle Knowles-Carter (/biːˈjɒnseɪ/ bee-YON-say) (born September 4, 1981) is an American singer, songwriter, record producer and actress. Born and raised in Houston, Texas, she performed in various singing and dancing competitions as a child, and rose to fame in the late 1990s as lead singer of R&B girl-group Destiny\\'s Child. Managed by her father, Mathew Knowles, the group became one of the world\\'s best-selling girl groups of all time. Their hiatus saw the release of Beyoncé\\'s debut album, Dangerously in Love (2003), which established her as a solo artist worldwide, earned five Grammy Awards and featured the Billboard Hot 100 number-one singles \"Crazy in Love\" and \"Baby Boy\".'}\n"
     ]
    }
   ],
   "source": [
    "example_article = example_jsons[0]\n",
    "example_article\n",
    "\n",
    "print(\"Title: \" + example_article[\"title\"])\n",
    "print(example_article[\"paragraphs\"][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2982be57",
   "metadata": {},
   "source": [
    "The previous article might be difficult to navigate so here is a nicely formatted example paragraph:\n",
    "```python\n",
    "{\n",
    "  \"context\": \"Beyoncé Giselle Knowles-Carter (/biːˈjɒnseɪ/ bee-YON-say) (born September 4, 1981) is an American singer, songwriter, record producer and actress. Born and raised in Houston, Texas, she performed in various singing and dancing competitions as a child, and rose to fame in the late 1990s as lead singer of R&B girl-group Destiny's Child. Managed by her father, Mathew Knowles, the group became one of the world's best-selling girl groups of all time. Their hiatus saw the release of Beyoncé's debut album, Dangerously in Love (2003), which established her as a solo artist worldwide, earned five Grammy Awards and featured the Billboard Hot 100 number-one singles 'Crazy in Love' and 'Baby Boy'\",\n",
    "  \"qas\": [\n",
    "    {\n",
    "      \"question\": \"When did Beyonce start becoming popular?\",\n",
    "      \"id\": \"56be85543aeaaa14008c9063\",\n",
    "      \"answers\": [\n",
    "        {\n",
    "          \"text\": \"in the late 1990s\",\n",
    "          \"answer_start\": 269\n",
    "        }\n",
    "      ],\n",
    "      \"is_impossible\": false\n",
    "    },\n",
    "    {\n",
    "      \"question\": \"What areas did Beyonce compete in when she was growing up?\",\n",
    "      \"id\": \"56be85543aeaaa14008c9065\",\n",
    "      \"answers\": [\n",
    "        {\n",
    "          \"text\": \"singing and dancing\",\n",
    "          \"answer_start\": 207\n",
    "        }\n",
    "      ],\n",
    "      \"is_impossible\": false\n",
    "    }\n",
    "  ]\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3345571",
   "metadata": {},
   "source": [
    "<a name='3-1'></a>\n",
    "### 3.1 - Creating a list of paired question and answers \n",
    "\n",
    "You are tasked with generating input/output pairs for a Question Answering (QA) model using the SQuAD 2.0 dataset. Each pair follows the structure:\n",
    "\n",
    "- inputs: `question: <Q> context: <P>`\n",
    "- targets: `answer: <A>`\n",
    "    \n",
    "Here, `<Q>` represents the question in the context of the given paragraph `<P>`, and `<A>` is a possible answer.\n",
    "\n",
    "In this notebook, we will focus on a single answer per question. However, it's essential to note that the dataset contains questions with multiple answers. When training a model in real-life scenarios, consider including all available information.\n",
    "\n",
    "<a name='ex-2'></a>\n",
    "### Exercise 2 - Parse the SQuaD 2.0 Dataset\n",
    "\n",
    "Your task is to implement the parse_squad function, which iterates over all the articles, paragraphs, and questions in the SQuAD dataset. Extract pairs of inputs and targets for the QA model using the provided code template.\n",
    "- Start with two empty lists: `inputs` and `targets`.\n",
    "- Loop over all the articles in the dataset.\n",
    "- For each article, loop over each paragraph.\n",
    "- Extract the context from the paragraph.\n",
    "- Loop over each question in the given paragraph.\n",
    "- Check if the question is not impossible and has at least one answer.\n",
    "- If the above condition is met, create the `question_context` sequence as described in the input structure.\n",
    "- Create the `answer` sequence using the first answer from the available answers.\n",
    "- Append the `question_context` to the `inputs` list.\n",
    "- Append the `answer` to the `targets` list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "b5344f35",
   "metadata": {
    "deletable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# GRADED FUNCTION: parse_squad\n",
    "def parse_squad(dataset):\n",
    "    \"\"\"Extract all the answers/questions pairs from the SQuAD dataset\n",
    "\n",
    "    Args:\n",
    "        dataset (dict): The imported JSON dataset\n",
    "\n",
    "    Returns:\n",
    "        inputs, targets: Two lists containing the inputs and the targets for the QA model\n",
    "    \"\"\"\n",
    "\n",
    "    inputs, targets = [], []\n",
    "\n",
    "    ### START CODE HERE ###\n",
    "    \n",
    "    # Loop over all the articles\n",
    "    for article in dataset:\n",
    "        \n",
    "        # Loop over each paragraph of each article\n",
    "        for paragraph in article['paragraphs']:\n",
    "            \n",
    "            # Extract context from the paragraph\n",
    "            context = paragraph['context']\n",
    "            \n",
    "            #Loop over each question of the given paragraph\n",
    "            for qa in paragraph['qas']:\n",
    "                \n",
    "                # If this question is not impossible and there is at least one answer\n",
    "                if len(qa['answers']) > 0 and not(qa['is_impossible']):\n",
    "                    \n",
    "                    # Create the question/context sequence\n",
    "                    question_context = 'question: ' + qa['question'] + ' context: ' + paragraph['context']\n",
    "                    \n",
    "                    # Create the answer sequence. Use the text field of the first answer\n",
    "                    answer = 'answer: ' + qa['answers'][0]['text']\n",
    "                    \n",
    "                    # Add the question_context to the inputs list\n",
    "                    inputs.append(question_context)\n",
    "                    \n",
    "                    # Add the answer to the targets list\n",
    "                    targets.append(answer)\n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "    \n",
    "    return inputs, targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "6744c424",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of question/answer pairs: 86821\n",
      "\n",
      "First Q/A pair:\n",
      "\n",
      "inputs: \u001b[34mquestion: When did Beyonce start becoming popular? context: Beyoncé Giselle Knowles-Carter (/biːˈjɒnseɪ/ bee-YON-say) (born September 4, 1981) is an American singer, songwriter, record producer and actress. Born and raised in Houston, Texas, she performed in various singing and dancing competitions as a child, and rose to fame in the late 1990s as lead singer of R&B girl-group Destiny's Child. Managed by her father, Mathew Knowles, the group became one of the world's best-selling girl groups of all time. Their hiatus saw the release of Beyoncé's debut album, Dangerously in Love (2003), which established her as a solo artist worldwide, earned five Grammy Awards and featured the Billboard Hot 100 number-one singles \"Crazy in Love\" and \"Baby Boy\".\u001b[0m\n",
      "\n",
      "targets: \u001b[32manswer: in the late 1990s\u001b[0m\n",
      "\n",
      "Last Q/A pair:\n",
      "\n",
      "inputs: \u001b[34mquestion: What is KMC an initialism of? context: Kathmandu Metropolitan City (KMC), in order to promote international relations has established an International Relations Secretariat (IRC). KMC's first international relationship was established in 1975 with the city of Eugene, Oregon, United States. This activity has been further enhanced by establishing formal relationships with 8 other cities: Motsumoto City of Japan, Rochester of the USA, Yangon (formerly Rangoon) of Myanmar, Xi'an of the People's Republic of China, Minsk of Belarus, and Pyongyang of the Democratic Republic of Korea. KMC's constant endeavor is to enhance its interaction with SAARC countries, other International agencies and many other major cities of the world to achieve better urban management and developmental programs for Kathmandu.\u001b[0m\n",
      "\n",
      "targets: \u001b[32manswer: Kathmandu Metropolitan City\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "inputs, targets =  parse_squad(example_jsons)          \n",
    "print(\"Number of question/answer pairs: \" + str(len(inputs)))\n",
    "\n",
    "print('\\nFirst Q/A pair:\\n\\ninputs: ' + colored(inputs[0], 'blue'))\n",
    "print('\\ntargets: ' + colored(targets[0], 'green'))\n",
    "print('\\nLast Q/A pair:\\n\\ninputs: ' + colored(inputs[-1], 'blue'))\n",
    "print('\\ntargets: ' + colored(targets[-1], 'green'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2b164c2",
   "metadata": {},
   "source": [
    "#### **Expected Output:**\n",
    "```\n",
    "Number of question/answer pairs: 86821\n",
    "\n",
    "First Q/A pair:\n",
    "\n",
    "inputs: question: When did Beyonce start becoming popular? context: Beyoncé Giselle Knowles-Carter (/biːˈjɒnseɪ/ bee-YON-say) (born September 4, 1981) is an American singer, songwriter, record producer and actress. Born and raised in Houston, Texas, she performed in various singing and dancing competitions as a child, and rose to fame in the late 1990s as lead singer of R&B girl-group Destiny's Child. Managed by her father, Mathew Knowles, the group became one of the world's best-selling girl groups of all time. Their hiatus saw the release of Beyoncé's debut album, Dangerously in Love (2003), which established her as a solo artist worldwide, earned five Grammy Awards and featured the Billboard Hot 100 number-one singles \"Crazy in Love\" and \"Baby Boy\".\n",
    "\n",
    "targets: answer: in the late 1990s\n",
    "\n",
    "Last Q/A pair:\n",
    "\n",
    "inputs: question: What is KMC an initialism of? context: Kathmandu Metropolitan City (KMC), in order to promote international relations has established an International Relations Secretariat (IRC). KMC's first international relationship was established in 1975 with the city of Eugene, Oregon, United States. This activity has been further enhanced by establishing formal relationships with 8 other cities: Motsumoto City of Japan, Rochester of the USA, Yangon (formerly Rangoon) of Myanmar, Xi'an of the People's Republic of China, Minsk of Belarus, and Pyongyang of the Democratic Republic of Korea. KMC's constant endeavor is to enhance its interaction with SAARC countries, other International agencies and many other major cities of the world to achieve better urban management and developmental programs for Kathmandu.\n",
    "\n",
    "targets: answer: Kathmandu Metropolitan City\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "f197bb69",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[92m All tests passed\n"
     ]
    }
   ],
   "source": [
    "# UNIT TEST\n",
    "w3_unittest.test_parse_squad(parse_squad)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1f69b24",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "You will use 40000 samples for training and 5000 samples for testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "947354ad",
   "metadata": {
    "deletable": false,
    "editable": false,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 40K pairs for training\n",
    "inputs_train = inputs[0:40000] \n",
    "targets_train = targets[0:40000]  \n",
    "\n",
    "# 5K pairs for testing\n",
    "inputs_test = inputs[40000:45000] \n",
    "targets_test =  targets[40000:45000] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c21fd31",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Now, you can create the batch dataset of padded sequences. You will first tokenize the inputs and the targets. Then, using the function `tf.keras.preprocessing.sequence.pad_sequences`, you will ensure that the inputs and the outputs have the required lengths. Remember that the sequences longer than the required size will be truncated and the shorter ones will be padded with `0`. This setup is very similar to the other one used in this and the previous notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "83393c74",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Limit the size of the input and output data so this can run in this environment\n",
    "encoder_maxlen = 150\n",
    "decoder_maxlen = 50\n",
    "\n",
    "inputs_str = [tokenizer.tokenize(s) for s in inputs_train]\n",
    "targets_str = [tf.concat([tokenizer.tokenize(s), [1]], 0) for s in targets_train]\n",
    "\n",
    "inputs = tf.keras.preprocessing.sequence.pad_sequences(inputs_str, maxlen=encoder_maxlen, padding='post', truncating='post')\n",
    "targets = tf.keras.preprocessing.sequence.pad_sequences(targets_str, maxlen=decoder_maxlen, padding='post', truncating='post')\n",
    "\n",
    "inputs = tf.cast(inputs, dtype=tf.int32)\n",
    "targets = tf.cast(targets, dtype=tf.int32)\n",
    "\n",
    "# Create the final training dataset.\n",
    "BUFFER_SIZE = 10000\n",
    "BATCH_SIZE = 64\n",
    "dataset = tf.data.Dataset.from_tensor_slices((inputs, targets)).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df82c8a9",
   "metadata": {},
   "source": [
    "<a name='3-2'></a>\n",
    "### 3.2 Fine tune the T5 model\n",
    "\n",
    "Now, you will train the model for 2 epochs. In the T5 model, all the weights are adjusted during the fine tuning. As usual, fine tuning this model to get state of the art results would require more time and resources than there are available in this environment, but you are welcome to train the model for more epochs and with more data using Colab GPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "aaaba558",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss 5.986125\n",
      "Time taken for one epoch: 74.45405197143555 sec\n",
      "Epoch 2, Loss 5.356925\n",
      "Time taken for one epoch: 35.4884033203125 sec\n"
     ]
    }
   ],
   "source": [
    "# Define the number of epochs\n",
    "epochs = 2\n",
    "losses = []\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(epochs):\n",
    "    \n",
    "    start = time.time()\n",
    "    train_loss.reset_states()\n",
    "    number_of_batches=len(list(enumerate(dataset)))\n",
    "\n",
    "    for (batch, (inp, tar)) in enumerate(dataset):\n",
    "        print(f'Epoch {epoch+1}, Batch {batch+1}/{number_of_batches}', end='\\r')\n",
    "        transformer_utils.train_step(inp, tar, transformer, loss_object, optimizer, train_loss)\n",
    "    \n",
    "    print (f'Epoch {epoch+1}, Loss {train_loss.result():.4f}')\n",
    "    losses.append(train_loss.result())\n",
    "    \n",
    "    print (f'Time taken for one epoch: {time.time() - start} sec')\n",
    "    #if epoch % 15 == 0:\n",
    "        #transformer.save_weights('./pretrained_models/model_qa_temp')\n",
    "# Save the final model\n",
    "#transformer.save_weights('./pretrained_models/model_qa_temp')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23b8dc0c",
   "metadata": {},
   "source": [
    "To get a model that works properly, you would need to train for about 100 epochs. So, we have pretrained a model for you. Just load the weights in the current model and let's use it for answering questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "144e769b",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.checkpoint.checkpoint.CheckpointLoadStatus at 0x7f40b006f880>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Restore the weights\n",
    "transformer.load_weights('./pretrained_models/model_qa3')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "360e09fd",
   "metadata": {},
   "source": [
    "<a name='3-3'></a>\n",
    "### 3.3 - Implement your Question Answering model\n",
    "In this final step, you will implement the answer_question function, utilizing a pre-trained transformer model for question answering.\n",
    "\n",
    "To help you out the `transformer_utils.next_word` function is provided. This function receives the question and beginning of the answer (both in tensor format) alongside the model to predict the next token in the answer. The next cell shows how to use this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "92b40de0",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted next word is: 'blue'\n",
      "Answer so far: 'answer: blue'\n"
     ]
    }
   ],
   "source": [
    "# Define an example question\n",
    "example_question = \"question: What color is the sky? context: Sky is blue\"\n",
    "\n",
    "# Question is tokenized and padded\n",
    "# Note that this is hardcoded here but you must implement this in the upcoming exercise\n",
    "tokenized_padded_question = tf.constant([[822, 10, 363, 945, 19, 8, 5796, 58, 2625, 10, 5643, 19, 1692, 0, 0]])\n",
    "\n",
    "# All answers begin with the string \"answer: \"\n",
    "# Feel free to check that this is indeed the tokenized version of that string\n",
    "tokenized_answer = tf.constant([[1525,   10]])\n",
    "\n",
    "# Predict the next word using the transformer_utils.next_word function\n",
    "# Notice that it expects the question, answer and model (in that order)\n",
    "next_word = transformer_utils.next_word(tokenized_padded_question, tokenized_answer, transformer)\n",
    "\n",
    "print(f\"Predicted next word is: '{tokenizer.detokenize(next_word).numpy()[0].decode('utf-8')}'\")\n",
    "\n",
    "# Concatenate predicted word with answer so far\n",
    "answer_so_far = tf.concat([tokenized_answer, next_word], axis=-1)\n",
    "\n",
    "print(f\"Answer so far: '{tokenizer.detokenize(answer_so_far).numpy()[0].decode('utf-8')}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e23a6be",
   "metadata": {},
   "source": [
    "<a name='ex-3'></a>\n",
    "### Exercise 3 - Implement the question answering function\n",
    "\n",
    "Implement the `answer_question` function. Here are the steps:\n",
    "- **Question Setup:**\n",
    "\n",
    "  - Tokenize the given question using the provided tokenizer.\n",
    "  - Add an extra dimension to the tensor for compatibility.\n",
    "  - Pad the question tensor using `pad_sequences` to ensure the sequence has the specified max length. This function will truncate the sequence if it is larger or pad with zeros if it is shorter.\n",
    "- **Answer Setup:**\n",
    "  - Tokenize the initial answer, noting that all answers begin with the string \"answer: \".\n",
    "  - Add an extra dimension to the tensor for compatibility.\n",
    "  - Get the id of the `EOS` token, typically represented by 1.\n",
    "- **Generate Answer:**\n",
    "  - Loop for `decoder_maxlen` iterations.\n",
    "  - Use the `transformer_utils.next_word` function, which predicts the next token in the answer using the model, input document, and the current state of the output.\n",
    "  - Concatenate the predicted next word to the output tensor.\n",
    "- **Stop Condition:**\n",
    "  - The text generation stops if the model predicts the `EOS` token.\n",
    "  - If the `EOS` token is predicted, break out of the loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "91def253",
   "metadata": {
    "deletable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# GRADED FUNCTION: answer_question\n",
    "def answer_question(question, model, tokenizer, encoder_maxlen=150, decoder_maxlen=50):\n",
    "    \"\"\"\n",
    "    A function for question answering using the transformer model\n",
    "    Arguments:\n",
    "        question (tf.Tensor): Input data with question and context\n",
    "        model (tf.keras.model): The transformer model\n",
    "        tokenizer (function): The SentencePiece tokenizer\n",
    "        encoder_maxlen (number): Max length of the encoded sequence\n",
    "        decoder_maxlen (number): Max length of the decoded sequence\n",
    "    Returns:\n",
    "        _ (str): The answer to the question\n",
    "    \"\"\"\n",
    "    \n",
    "    ### START CODE HERE ###\n",
    "    \n",
    "    # QUESTION SETUP\n",
    "    \n",
    "    # Tokenize the question\n",
    "    tokenized_question = tokenizer.tokenize(question)\n",
    "    \n",
    "    # Add an extra dimension to the tensor\n",
    "    tokenized_question = tf.expand_dims(tokenized_question, 0) \n",
    "    \n",
    "    # Pad the question tensor\n",
    "    padded_question = tf.keras.preprocessing.sequence.pad_sequences(tokenized_question,\n",
    "                                                                    maxlen=encoder_maxlen,\n",
    "                                                                    padding='post', \n",
    "                                                                    truncating='post') \n",
    "    # ANSWER SETUP\n",
    "    \n",
    "    # Tokenize the answer\n",
    "    # Hint: All answers begin with the string \"answer: \"\n",
    "    tokenized_answer = tokenizer.tokenize(\"answer: \")\n",
    "    \n",
    "    # Add an extra dimension to the tensor\n",
    "    tokenized_answer = tf.expand_dims(tokenized_answer, 0)\n",
    "    \n",
    "    # Get the id of the EOS token\n",
    "    eos = tokenizer.string_to_id(\"</s>\") \n",
    "    \n",
    "    # Loop for decoder_maxlen iterations\n",
    "    for i in range(decoder_maxlen):\n",
    "        \n",
    "        # Predict the next word using the model, the input document and the current state of output\n",
    "        next_word = transformer_utils.next_word(padded_question, tokenized_answer, model)\n",
    "        \n",
    "        # Concat the predicted next word to the output \n",
    "        tokenized_answer = tf.concat([tokenized_answer, next_word], axis=1)\n",
    "        \n",
    "        # The text generation stops if the model predicts the EOS token\n",
    "        if next_word.numpy()[0][0] == eos:\n",
    "            break \n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "\n",
    "    return tokenized_answer "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de501d8c",
   "metadata": {},
   "source": [
    "Let's test the model with some question from the training dataset. Check if the answers match the correct one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "163e79eb",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[34mb'answer: January 9, 1957'\u001b[0m\n",
      "\n",
      "question: When was the Chechen-Ingush Autonomous Soviet Socialist Republic transferred from the Georgian SSR? context: On January 9, 1957, Karachay Autonomous Oblast and Chechen-Ingush Autonomous Soviet Socialist Republic were restored by Khrushchev and they were transferred from the Georgian SSR back to the Russian SFSR.\n",
      "\u001b[32manswer: January 9, 1957\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "idx = 10408\n",
    "\n",
    "result = answer_question(inputs_train[idx], transformer, tokenizer)\n",
    "print(colored(pretty_decode(result, sentinels, tokenizer).numpy()[0], 'blue'))\n",
    "print()\n",
    "print(inputs_train[idx])\n",
    "print(colored(targets_train[idx], 'green'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fae9e8d",
   "metadata": {},
   "source": [
    "#### **Expected Output:**\n",
    "```\n",
    "b'answer: January 9, 1957'\n",
    "\n",
    "question: When was the Chechen-Ingush Autonomous Soviet Socialist Republic transferred from the Georgian SSR? context: On January 9, 1957, Karachay Autonomous Oblast and Chechen-Ingush Autonomous Soviet Socialist Republic were restored by Khrushchev and they were transferred from the Georgian SSR back to the Russian SFSR.\n",
    "answer: January 9, 1957\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "19ac8067",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[92m All tests passed\n"
     ]
    }
   ],
   "source": [
    "# UNIT TEST\n",
    "w3_unittest.test_answer_question(answer_question)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06588341",
   "metadata": {},
   "source": [
    "Test the model with question 110"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "6c381df3",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[34mb'answer: 50'\u001b[0m\n",
      "\n",
      "question:  What percentage of the vote was recorded as approving Napoleon's constitution? context: Napoleon established a political system that historian Martyn Lyons called \"dictatorship by plebiscite.\" Worried by the democratic forces unleashed by the Revolution, but unwilling to ignore them entirely, Napoleon resorted to regular electoral consultations with the French people on his road to imperial power. He drafted the Constitution of the Year VIII and secured his own election as First Consul, taking up residence at the Tuileries. The constitution was approved in a rigged plebiscite held the following January, with 99.94 percent officially listed as voting \"yes.\" Napoleon's brother, Lucien, had falsified the returns to show that 3 million people had participated in the plebiscite; the real number was 1.5 million. Political observers at the time assumed the eligible French voting public numbered about 5 million people, so the regime artificially doubled the participation rate to indicate popular enthusiasm for the Consulate. In the first few months of the Consulate, with war in Europe still raging and internal instability still plaguing the country, Napoleon's grip on power remained very tenuous.\n",
      "\u001b[32manswer: 99.94\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "idx = 110\n",
    "result = answer_question(inputs_test[idx], transformer, tokenizer)\n",
    "print(colored(pretty_decode(result, sentinels, tokenizer).numpy()[0], 'blue'))\n",
    "print()\n",
    "print(inputs_test[idx])\n",
    "print(colored(targets_test[idx], 'green'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd09ec41",
   "metadata": {},
   "source": [
    "Test the model with question 301. Use this cell to play with the model by selecting other test questions. Look if the model has learnt something or if it is just generating random text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "fc9c898f",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[34mb'answer: June 1840'\u001b[0m\n",
      "\n",
      "question:  On what date was a state funeral held for Napoleon? context: In 1840, Louis Philippe I obtained permission from the British to return Napoleon's remains to France. On 15 December 1840, a state funeral was held. The hearse proceeded from the Arc de Triomphe down the Champs-Élysées, across the Place de la Concorde to the Esplanade des Invalides and then to the cupola in St Jérôme's Chapel, where it remained until the tomb designed by Louis Visconti was completed. In 1861, Napoleon's remains were entombed in a porphyry sarcophagus in the crypt under the dome at Les Invalides.\n",
      "\u001b[32manswer: 15 December 1840\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "idx = 311\n",
    "result = answer_question(inputs_test[idx], transformer, tokenizer)\n",
    "print(colored(pretty_decode(result, sentinels, tokenizer).numpy()[0], 'blue'))\n",
    "print()\n",
    "print(inputs_test[idx])\n",
    "print(colored(targets_test[idx], 'green'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e3d00de",
   "metadata": {},
   "source": [
    "Congratulations, you have finished the last assignment of this specialization. Now, you know what is behind the powerful models like ChatGPT. Now it is time for you to find and solve the huge amount of problems that can be approached with NLP."
   ]
  }
 ],
 "metadata": {
  "grader_version": "1",
  "jupytext": {
   "encoding": "# -*- coding: utf-8 -*-"
  },
  "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": 5
}