Summarization using Transformer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7yuytuIllsv1"
   },
   "source": [
    "\n",
    "# Assignment 2: Transformer Summarizer\n",
    "\n",
    "Welcome to the second assignment of course 4. In this assignment you will explore summarization using the transformer model. Yes, you will implement the transformer decoder from scratch, but we will slowly walk you through it. There are many hints in this notebook so feel free to use them as needed. \n",
    "\n",
    "<img src = \"transformerNews.png\">\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4-3lxSnXRWPx"
   },
   "source": [
    "## Outline\n",
    "\n",
    "- [Introduction](#0)\n",
    "- [Part 1: Importing the dataset](#1)\n",
    "    - [1.1 Encode & Decode helper functions](#1.1)\n",
    "    - [1.2 Defining parameters](#1.2)\n",
    "    - [1.3 Exploring the data](#1.3)\n",
    "- [Part 2: Summarization with transformer](#2)\n",
    "    - [2.1 Dot product attention](#2.1)\n",
    "        - [Exercise 01](#ex01)\n",
    "    - [2.2 Causal Attention](#2.2)\n",
    "        - [Exercise 02](#ex02)\n",
    "    - [2.3 Transformer decoder block](#2.3)\n",
    "        - [Exercise 03](#ex03)\n",
    "    - [2.4 Transformer Language model](#2.4)\n",
    "        - [Exercise 04](#ex04)\n",
    "- [Part 3: Training](#3)\n",
    "    - [3.1 Training the model](#3.1)\n",
    "        - [Exercise 05](#ex05)\n",
    "- [Part 4: Evaluation](#4)\n",
    "    - [4.1 Loading in a trained model](#4.1)\n",
    "- [Part 5: Testing with your own input](#5) \n",
    "    - [Exercise 6](#ex06)\n",
    "    - [5.1 Greedy decoding](#5.1)\n",
    "        - [Exercise 07](#ex07)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "H4NlfEQhRWPy"
   },
   "source": [
    "<a name='0'></a>\n",
    "### Introduction\n",
    "\n",
    "Summarization is an important task in natural language processing and could be useful for a consumer enterprise. For example, bots can be used to scrape articles, summarize them, and then you can use sentiment analysis to identify the sentiment about certain stocks. Anyways who wants to read an article or a long email today, when you can build a transformer to summarize text for you. Let's get started, by completing this assignment you will learn to:  \n",
    "\n",
    "- Use built-in functions to preprocess your data\n",
    "- Implement DotProductAttention\n",
    "- Implement Causal Attention\n",
    "- Understand how attention works\n",
    "- Build the transformer model\n",
    "- Evaluate your model\n",
    "- Summarize an article\n",
    "\n",
    "As you can tell, this model is slightly different than the ones you have already implemented. This is heavily based on attention and does not rely on sequences, which allows for parallel computing. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "CChWzW-rEHVb",
    "outputId": "a0b3e98b-7fc6-492d-c8ad-3a263b54f670"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:tokens_length=568 inputs_length=512 targets_length=114 noise_density=0.15 mean_noise_span_length=3.0 \n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import textwrap\n",
    "wrapper = textwrap.TextWrapper(width=70)\n",
    "\n",
    "import trax\n",
    "from trax import layers as tl\n",
    "from trax.fastmath import numpy as jnp\n",
    "\n",
    "# to print the entire np array\n",
    "np.set_printoptions(threshold=sys.maxsize)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "kEL2rvaHRWP4"
   },
   "source": [
    "<a name='1'></a>\n",
    "## Part 1: Importing the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Trax makes it easy to work with Tensorflow's datasets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VInmKSkhEhle"
   },
   "outputs": [],
   "source": [
    "# This will download the dataset if no data_dir is specified.\n",
    "# Downloading and processing can take bit of time,\n",
    "# so we have the data already in 'data/' for you\n",
    "\n",
    "# Importing CNN/DailyMail articles dataset\n",
    "train_stream_fn = trax.data.TFDS('cnn_dailymail',\n",
    "                                 data_dir='data/',\n",
    "                                 keys=('article', 'highlights'),\n",
    "                                 train=True)\n",
    "\n",
    "# This should be much faster as the data is downloaded already.\n",
    "eval_stream_fn = trax.data.TFDS('cnn_dailymail',\n",
    "                                data_dir='data/',\n",
    "                                keys=('article', 'highlights'),\n",
    "                                train=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='1.1'></a>\n",
    "## 1.1 Tokenize & Detokenize helper functions\n",
    "\n",
    "Just like in the previous assignment, the cell above loads in the encoder for you. Given any data set, you have to be able to map words to their indices, and indices to their words. The inputs and outputs to your [Trax](https://github.com/google/trax) models are usually tensors of numbers where each number corresponds to a word. If you were to process your data manually, you would have to make use of the following: \n",
    "\n",
    "- <span style='color:blue'> word2Ind: </span> a dictionary mapping the word to its index.\n",
    "- <span style='color:blue'> ind2Word:</span> a dictionary mapping the index to its word.\n",
    "- <span style='color:blue'> word2Count:</span> a dictionary mapping the word to the number of times it appears. \n",
    "- <span style='color:blue'> num_words:</span> total number of words that have appeared. \n",
    "\n",
    "Since you have already implemented these in previous assignments of the specialization, we will provide you with helper functions that will do this for you. Run the cell below to get the following functions:\n",
    "\n",
    "- <span style='color:blue'> tokenize: </span> converts a text sentence to its corresponding token list (i.e. list of indices). Also converts words to subwords.\n",
    "- <span style='color:blue'> detokenize: </span> converts a token list to its corresponding sentence (i.e. string)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "djTiSLcaNFGa"
   },
   "outputs": [],
   "source": [
    "def tokenize(input_str, EOS=1):\n",
    "    \"\"\"Input str to features dict, ready for inference\"\"\"\n",
    "  \n",
    "    # Use the trax.data.tokenize method. It takes streams and returns streams,\n",
    "    # we get around it by making a 1-element stream with `iter`.\n",
    "    inputs =  next(trax.data.tokenize(iter([input_str]),\n",
    "                                      vocab_dir='vocab_dir/',\n",
    "                                      vocab_file='summarize32k.subword.subwords'))\n",
    "    \n",
    "    # Mark the end of the sentence with EOS\n",
    "    return list(inputs) + [EOS]\n",
    "\n",
    "def detokenize(integers):\n",
    "    \"\"\"List of ints to str\"\"\"\n",
    "  \n",
    "    s = trax.data.detokenize(integers,\n",
    "                             vocab_dir='vocab_dir/',\n",
    "                             vocab_file='summarize32k.subword.subwords')\n",
    "    \n",
    "    return wrapper.fill(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7WvhaFbCRWQS"
   },
   "source": [
    "<a name='1.2'></a>\n",
    "\n",
    "## 1.2 Preprocessing for Language Models: Concatenate It!\n",
    "\n",
    "This week you will use a language model -- Transformer Decoder -- to solve\n",
    "an input-output problem. As you know, language models only predict the next\n",
    "word, they have no notion of inputs. To create a single input suitable for\n",
    "a language model, we concatenate inputs with targets putting a separator\n",
    "in between. We also need to create a mask -- with 0s at inputs and 1s at targets -- so that the model is not penalized for mis-predicting the article and only focuses on the summary. See the preprocess function below for how this is done."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "c4rgPxYSRWQS"
   },
   "outputs": [],
   "source": [
    "# Special tokens\n",
    "SEP = 0 # Padding or separator token\n",
    "EOS = 1 # End of sentence token\n",
    "\n",
    "# Concatenate tokenized inputs and targets using 0 as separator.\n",
    "def preprocess(stream):\n",
    "    for (article, summary) in stream:\n",
    "        joint = np.array(list(article) + [EOS, SEP] + list(summary) + [EOS])\n",
    "        mask = [0] * (len(list(article)) + 2) + [1] * (len(list(summary)) + 1) # Accounting for EOS and SEP\n",
    "        yield joint, joint, np.array(mask)\n",
    "\n",
    "# You can combine a few data preprocessing steps into a pipeline like this.\n",
    "input_pipeline = trax.data.Serial(\n",
    "    # Tokenizes\n",
    "    trax.data.Tokenize(vocab_dir='vocab_dir/',\n",
    "                       vocab_file='summarize32k.subword.subwords'),\n",
    "    # Uses function defined above\n",
    "    preprocess,\n",
    "    # Filters out examples longer than 2048\n",
    "    trax.data.FilterByLength(2048)\n",
    ")\n",
    "\n",
    "# Apply preprocessing to data streams.\n",
    "train_stream = input_pipeline(train_stream_fn())\n",
    "eval_stream = input_pipeline(eval_stream_fn())\n",
    "\n",
    "train_input, train_target, train_mask = next(train_stream)\n",
    "\n",
    "assert sum((train_input - train_target)**2) == 0  # They are the same in Language Model (LM)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 238
    },
    "colab_type": "code",
    "id": "uKFoGsUKSa_I",
    "outputId": "bc4d6634-d716-4311-d49c-1956bca2bc2d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Single example mask:\n",
      "\n",
      " [0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1\n",
      " 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1\n",
      " 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]\n"
     ]
    }
   ],
   "source": [
    "# prints mask, 0s on article, 1s on summary\n",
    "print(f'Single example mask:\\n\\n {train_mask}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 459
    },
    "colab_type": "code",
    "id": "S4uHyCkbSuUo",
    "outputId": "52845be8-f2fc-4803-bf7a-ed9725fe2bac"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Single example:\n",
      "\n",
      " A curry house has been fined £3,000 after two dead mice were\n",
      "discovered next to a sack of onions. Health inspectors also found\n",
      "rodent droppings in the kitchen and food stored inside a dirty shed\n",
      "when they visited Khans Tandoori and Balti Takeaway in Portsmouth,\n",
      "Hampshire. The kitchen sink was also dry and had no soap and the staff\n",
      "toilet was in a poor condition, while layers of dirt, grease and\n",
      "debris had built up in the areas where food was handled and stored.\n",
      "Khans Tandoori and Balti Takeaway in Portsmouth, Hampshire, has been\n",
      "fined £3,000 after these dead mice were discovered next to a sack of\n",
      "onions . Magistrates ordered the restaurant to pay a £1,800 fine,\n",
      "£1,117 in costs and a £36 victim surcharge, after owner of Harold\n",
      "Southsea Ltd, admitted five charges under hygiene legislation on\n",
      "behalf of the takeaway. Health officer Christopher Larkin from\n",
      "Portsmouth City Council had discovered the dead mice on sticky boards\n",
      "that had been put out to catch the vermin when he visited the\n",
      "restaurant last July. Staff were forced to close the restaurant to\n",
      "make emergency repairs, and it was only allowed to reopen once the\n",
      "kitchen had been cleaned and disinfected, and the food moved out of\n",
      "the lean-to shed. Other improvements were also made. Magistrates\n",
      "ordered the restaurant (pictured) to pay a £1,800 fine, £1,117 in\n",
      "costs and a £36 victim surcharge, after owner of Harold Southsea Ltd,\n",
      "admitted five charges under hygiene legislation on behalf of the\n",
      "takeaway . Inspectors found that the kitchen sink was dry and had no\n",
      "soap, while the staff toilet was in a poor condition . 'Food hygiene\n",
      "is of prime importance and we cannot allow poor standards and\n",
      "behaviour like this,' said Alan Cufley, head of environmental health\n",
      "at the Portsmouth City Council. 'When problems are found, we try to\n",
      "work with businesses to help them improve, but if necessary we will\n",
      "take appropriate action, including prosecution, to protect the\n",
      "public.' The restaurant admitted failing to maintain the premises in\n",
      "good repair, failing to maintain them in a clean condition, failing to\n",
      "protect food from risk of contamination, failing to protect the\n",
      "kitchen from pest entry, and failing to control pests. These bags of\n",
      "cleaning products were discovered when inspectors visited the takeaway\n",
      "in July last year . Mr Khan, who owns the takeaway, said: 'It's been\n",
      "cleaned up. Basically we have done everything that's required. 'We\n",
      "have been trying to improve it for the future. We have never had it\n",
      "this bad before. 'We have taken measures to have a clean kitchen.'\n",
      "Council environment chief Robert New said: 'There's just no excuse for\n",
      "mice to be in kitchens. 'It has been a priority since I took over.\n",
      "It's about protecting the public.' He added there has been a change of\n",
      "strategy at the council, with officers now tackling 'medium to bad'\n",
      "offenders as a priority, rather than looking at each restaurant in\n",
      "turn. Health inspectors also found food being stored inside this dirty\n",
      "shed. The takeaway was forced to close, and could only reopen once the\n",
      "restaurant had been cleaned, and the food removed from the lean-to .\n",
      "Mouse droppings were found on the floor of the kitchen, and in a goods\n",
      "store used by staff .<EOS><pad>Rodentdroppings found in the Khans\n",
      "Tandoori and Balti Takeaway kitchen . Food was being stored in a dirty\n",
      "shed, and the kitchen sink had no soap . Dead mice were discovered on\n",
      "sticky boards intended to catch them . Takeaway fined after owner\n",
      "admitted hygiene charges on its behalf .<EOS>\n"
     ]
    }
   ],
   "source": [
    "# prints: [Example][<EOS>][<pad>][Example Summary][<EOS>]\n",
    "print(f'Single example:\\n\\n {detokenize(train_input)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "T4sDS1WIVaYG"
   },
   "source": [
    "<a name='1.3'></a>\n",
    "\n",
    "## 1.3 Batching with bucketing\n",
    "\n",
    "As in the previous week, we use bucketing to create batches of data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "oqj1NsbERWQX"
   },
   "outputs": [],
   "source": [
    "# Bucketing to create batched generators.\n",
    "\n",
    "# Buckets are defined in terms of boundaries and batch sizes.\n",
    "# Batch_sizes[i] determines the batch size for items with length < boundaries[i]\n",
    "# So below, we'll take a batch of 16 sentences of length < 128 , 8 of length < 256,\n",
    "# 4 of length < 512. And so on. \n",
    "boundaries =  [128, 256,  512, 1024]\n",
    "batch_sizes = [16,    8,    4,    2, 1]\n",
    "\n",
    "# Create the streams.\n",
    "train_batch_stream = trax.data.BucketByLength(\n",
    "    boundaries, batch_sizes)(train_stream)\n",
    "\n",
    "eval_batch_stream = trax.data.BucketByLength(\n",
    "    boundaries, batch_sizes)(eval_stream)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "P6M5OA8QRWQb"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 1024)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Every execution will result in generation of a different article\n",
    "# Try running this cell multiple times to see how the length of the examples affects the batch size\n",
    "input_batch, _, mask_batch = next(train_batch_stream)\n",
    "\n",
    "# Shape of the input_batch\n",
    "input_batch.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "SjNOlljxTGuQ",
    "outputId": "9227c68c-6369-4ce8-8137-506c594f6ad2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  567   379  6230  2014 13088  5611   379 24404   186 15121 20195   784\n",
      "  2989  1337  4758  8046    23  1205    64   809  2328  3586    28  7100\n",
      "  6204     7    79  1019    68    82  1940  1248 20988  9372     2 11659\n",
      "    16    68   527 26995    16   320   117   763   101  1095  1099    34\n",
      "   163 12339 12554     4  2215   809   213   225  6735    27 11156    14\n",
      "  1337     2   131  1113  3611  4751  5988  2328  3586    28   717  7100\n",
      " 10952    79 13316    47     7    26  1566  6199    26   213 23375  4864\n",
      " 12032  1387   763   101     8 12293    12   191  4231   364   213   763\n",
      "   157 16885    68     2    41   221    68  5537  6510 27884     4 14561\n",
      " 27872  1099     9  3766  1041   141    28   984   102  3622  7100  6204\n",
      "     7    79  1636  4758  8046    70  1779    23    46  5754   527  2699\n",
      "    68   221  3105   527  3130     6 23375  4864  8927   316  7054  6882\n",
      "  2884    70   132    28  2622  2685  1224   320   589   213  2147   527\n",
      "    68   221  3130     3  9175  6051     4   246  1019   846   379 20205\n",
      "    16  4339    11 24404   186 15121 20195   784  2989  1337  4758  8046\n",
      "    23  1205    64   809  2328  3586    28  7100  6204     7    79  1019\n",
      "    68    82  1940  1248 20988  9372     2 11659    16    68   527 26995\n",
      "    16   320   117   763   101  1095    80 25236   911    11    34    28\n",
      " 12554     4  2215   809   213   225  6735    27 11156    14  1337     2\n",
      "   131  1113  3611  4751  5988  2328  3586    28   717  7100 10952    79\n",
      " 13316    47     7    26  1566  6199    26   213 23375  4864 12032  1387\n",
      "   763   101     8 12293    12   191  4231   364   213   763   157 16885\n",
      "    68     2    41   221    68  5537  6510 27884     4 14561 27872    80\n",
      "     9  2989  1337   229    92 14354     4   320   379  8295    29   132\n",
      "    28   754  2264  1248 14831 10990     2   131    38    35  6122   379\n",
      "   285  7054  6882  2884   229  2334   412    28  3130     6  5029 13626\n",
      "     4     2    19   412    28 18072  1019   379  2228  7692   412   103\n",
      "  2065  1782   312   379    33   245   285  2030   186    33   547    28\n",
      "  2030   527  4758  8046 22846     2   824   229   379  2754    33     7\n",
      "   165  4991   101    70   213  1184   456   578  3898   131   127  1782\n",
      "   244  6847  2754    98  1593   456   483   320  2444   103     3    52\n",
      "     7     5  2754   103   229     3    13   358     7    26   456   662\n",
      "  2002  5466    68  8470  1012   186 10635  1385  5387     2   103     7\n",
      "     5    92  5258   213 15680    20  2974  1353 26681    17  7511  3622\n",
      "  7100  6204     7    79   316    68    64   132    28  2622   809 14409\n",
      "  2687     7     5  1168  1665   132  3501   752    78   754   728     3\n",
      "  8398 21748 15616    11    34    28   772  2622     2  3622  7100  6204\n",
      "     7    79   816    64 14236    28  1541  1838    28  3331  1779   127\n",
      "   131  1410    19   320  2444  4758  8046     7     5  3130     6 20806\n",
      "  1968  8927  7054  6882  2884   166   213  6006    40  3037    68   320\n",
      "   602    68   221  3130   379   252     6  2423    11  3622  7100  6204\n",
      "     7    79   511   163  7890  1019  1877 14080    16  7008  4555  1019\n",
      "    68   453   132   429   327   225  6735    27 11156    14   379    34\n",
      "    68  1566     2   213   286     6   104     6   292  1475    61  2685\n",
      "  7270  6000  1571   527  4415   640  2486    68   741     6 18196     4\n",
      "     2  1510  5190  6117    28  1541   753   320    68   691    28   506\n",
      "  2635  1779  4924    68   412    28   453   682  1782 13860  2328  3586\n",
      "    28  3898  9831   213  1541  1782    13   362    33     7   165   456\n",
      "  9564   320  1151   824   767    35   764   824  1123   132  3501 13094\n",
      "     3    13  1353   141  2685   320  2444  4758  8046     7     5  7054\n",
      "  6882  2884  8927   320  5029   279   130  3130  7511    33  1472    78\n",
      "   213   175  1517   186  6671   156 10220   861  1348 15204    28   445\n",
      "  7511    13   816   161   911  3898   213  6006   127   123 16948     2\n",
      " 11333  7270   412    28   966     2   131     2   403     2    62 18171\n",
      "   285   131     7    21    36   194  9160    61  1248 23500  3130     3\n",
      " 19338 26052  5274   997    11  4758  8046   229    92 14354     4   320\n",
      "  8295    29   132    28   754  2264  1248 14831 10990     2   131  5446\n",
      "  6122   285  7054  6882  2884   229  2334   412    28  3130     6  5029\n",
      " 13626     4     2    19   412    28 18072  1019  2228  7692   412   103\n",
      "  2065   379  4632  5438 26082   113   213   917  3611   312    33   245\n",
      "   285  2030   186    33   547    28  2030   527  4758  8046 22846     2\n",
      "   824   229  2754    33     7   165  4991   101    70   213  1184   456\n",
      "   578  3898   131   127  1782   244  6847  2754    98  1593   456   483\n",
      "   320  2444   103    80   200   192   213  6006   980   213  1541   412\n",
      "    28  1175   527  6872     2  4758  8046   436   103   412    28   776\n",
      "  1287     3   392    68 12339 12554     4   809  3622  7100  6204     7\n",
      "    79     2  5807 14527    17   132     2  1248    36   537 12554    16\n",
      "  3611   182 19890  6376  4368  3586    28   229   213    82   882   527\n",
      " 20988 15909   156 10163  9773   158 21710    59 10163   244    33     7\n",
      "   165   141  8223 14091   669 27884     4 14561 27872  1099    34    68\n",
      "  1111     2  4758  8046   127   527   213 12770 12778     5   458  3611\n",
      "   200    41  3224 23375  4864  8927 11033     4   669 27884     4 14561\n",
      " 27872  1099     9  2989  1337   229   761  8795   320 20988  9372     7\n",
      "     5 11492     4 19507     4   598   527 12770 12778     5     2  1480\n",
      "  1435   962 12393    21   412   117 22437 10014    80   285   117 11862\n",
      " 15898   132   613   186 18120    47   213 16620 19610   527   213  3130\n",
      "  1099   244   412   525   412  4758  8046     7     5  1534   285 20988\n",
      "  9372   229    28   117   763   101   458   957    28  4319  5699  7507\n",
      "    84   527   213  1438  5539   285   103    23    28   203   527 17249\n",
      "   185   186 10754   132 22846  3130 20264    10     1     0  4758  8046\n",
      "     7     5  3766   229 10635  1385  4689   131  1086  7855    17  8295\n",
      "  1019  2699   117  2228  2688  7760  1816    80  7054  6882  2884     2\n",
      "  1480   229  2449   149   412    28  3130     6 20806  1968  8927  2104\n",
      "     1     0     0     0     0     0     0     0     0     0     0     0\n",
      "     0     0     0     0     0     0     0     0     0     0     0     0\n",
      "     0     0     0     0]\n"
     ]
    }
   ],
   "source": [
    "# print corresponding integer values\n",
    "print(input_batch[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GD-72TENV2Jk"
   },
   "source": [
    "Things to notice:\n",
    " - First we see the corresponding values of the words.\n",
    " - The first 1, which represents the `<EOS>` tag of the article.\n",
    " - Followed by a 0, which represents a `<pad>` tag.\n",
    " - After the first 0 (`<pad>` tag) the corresponding values are of the words that are used for the summary of the article.\n",
    " - The second 1 represents the `<EOS>` tag for the summary.\n",
    " - All the trailing 0s represent `<pad>` tags which are appended to maintain consistent length (If you don't see them then it would mean it is already of max length)\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "Bu05ZwbWTE6P",
    "outputId": "3d455bd7-e343-4c25-a467-572d2abd837f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Article:\n",
      "\n",
      " By . Margot Peppers . Nigerian and Cameroonian pop star Dencia has hit\n",
      "out at Lupita Nyong'o for her new contract with Lancome, accusing her\n",
      "of bowing to 'white people companies'. In an angry tweet directed at\n",
      "the 12 Years A Slave star, she wrote: 'Oh @Lupita_Nyongo cln't talk\n",
      "abt the bleaching creams white people (Companies) make cuz the white\n",
      "man pays her, they own her!! [sic]'. The comment comes just a month\n",
      "after Miss Nyong'o mentioned Dencia - who has been accused of\n",
      "marketing her own brand of skin-bleaching cream called Whitenicious -\n",
      "in a speech about learning to value the color of her own skin. Scroll\n",
      "down for video . Butting heads: Nigerian and Cameroonian pop star\n",
      "Dencia has hit out at Lupita Nyong'o for her new contract with\n",
      "Lancome, accusing her of bowing to 'white people companies' Fighting\n",
      "words: In a tweet directed at the 12 Years A Slave star, she wrote:\n",
      "'Oh @Lupita_Nyongo cln't talk abt the bleaching creams white people\n",
      "(Companies) make cuz the white man pays her, they own her!! [sic]' The\n",
      "pop star is no stranger to . controversy; in a February interview with\n",
      "Ebony, she all but admitted . that Whitenicious is intended as a skin-\n",
      "lightener, not as a cure for . dark spots as it claims. 'When . you\n",
      "take that picture and you put a picture of Dencia darker, this is .\n",
      "what you're telling people - the product really works,' she said. 'And\n",
      "guess what? People really want to buy it. It's what it is. I don't\n",
      "really care.' Given her defiant and hypocritical attitude, it's no\n",
      "surprise the fiery singer was angered when Miss Nyong'o called her out\n",
      "in a speech at Essence's Black Women in Hollywood event on February\n",
      "27. Influential: In a recent speech, Miss Nyong'o read out loud a\n",
      "letter from a fan who said she decided not to buy Dencia's skin-\n",
      "whitening cream Whitenicious because the actress had inspired her to\n",
      "love her own skin . On-screen: Miss Nyong'o won an Oscar for Best\n",
      "Supporting Actress for her role in 2013 film 12 Years A Slave . In her\n",
      "talk, the 30-year-old opened up about how conventional standards of\n",
      "beauty once affected her self-esteem, reading aloud a letter written\n",
      "to her by a young girl who viewed her as a role model. 'Dear Lupita,'\n",
      "reads the letter. 'I think you're really lucky to be this black but\n",
      "yet this successful in Hollywood overnight. I was just about to buy\n",
      "Dencia's Whitenicious cream to lighten my skin when you appeared on\n",
      "the world map and saved me.' 'My heart bled a little when I read those\n",
      "words,' the actress said through tears, explaining how as a child,\n",
      "she, too, would pray that she'd one day wake up with lighter skin.\n",
      "Hypocritical: Dencia is no stranger to controversy; in a February\n",
      "interview with Ebony, she essentially admitted that Whitenicious is\n",
      "intended as a skin-lightener, not as a cure for dark spots as it\n",
      "claims . Perpetuating the problem: 'When you take that picture and you\n",
      "put a picture of Dencia darker, this is what you're telling people -\n",
      "the product really works,' she said. 'And guess what? People really\n",
      "want to buy it' But while the actress saw the letter as a source of\n",
      "inspiration, Dencia took it as a personal attack. After her angry\n",
      "tweet at Miss Nyong'o, criticism poured in, with one person tweeting:\n",
      "'B**** lupita is the new face of Lancôme!! SHE WINS!! And you're just\n",
      "TRASH [sic]'. In her response, Dencia said of the cosmetics company:\n",
      "'But they sell bleaching cream tho [sic]'. The pop star is likely\n",
      "referring to Lancome's Blanc Expert range of cosmetics, which are\n",
      "actually advertised as 'brighteners' that 'regulate melanin production\n",
      "and awaken the luminosity of the skin'. And as far as Dencia's claim\n",
      "that Lancome is a 'white people company', a quick perusal of the\n",
      "website reveals that it has a number of concealers and foundations in\n",
      "darker skin tones.<EOS><pad>Dencia's comment is hypocritical\n",
      "considering she recently courted controversy for marketing 'dark spot\n",
      "remover' Whitenicious, which is frequently used as a skin-whitening\n",
      "cream .<EOS><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pa\n",
      "d><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pa\n",
      "d><pad>\n"
     ]
    }
   ],
   "source": [
    "# print the article and its summary\n",
    "print('Article:\\n\\n', detokenize(input_batch[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aNFVhgHoncGm"
   },
   "source": [
    "You can see that the data has the following structure:\n",
    "- <span style='color:blue'> [Article] </span> -> `<EOS>` -> `<pad>` -> <span style='color:blue'> [Article Summary] </span> -> `<EOS>` -> (possibly) multiple `<pad>`\n",
    "\n",
    "The loss is taken only on the summary using cross_entropy as loss function. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Un8NHIRoj-1W"
   },
   "source": [
    "<a name='2'></a>\n",
    "# Part 2: Summarization with transformer\n",
    "\n",
    "Now that we have given you the data generator and have handled the preprocessing for you, it is time for you to build your own model. We saved you some time because we know you have already preprocessed data before in this specialization, so we would rather you spend your time doing the next steps. \n",
    "\n",
    "You will be implementing the attention from scratch and then using it in your transformer model. Concretely, you will understand how attention works, how you use it to connect the encoder and the decoder.\n",
    "\n",
    "<img src=\"transformer_decoder_zoomin.png\">\n",
    "\n",
    "<a name='2.1'></a>\n",
    "## 2.1 Dot product attention \n",
    "\n",
    "Now you will implement dot product attention which takes in a query, key, value, and a mask. It returns the output. \n",
    "\n",
    "<img src =\"dotproduct.png\">\n",
    "\n",
    "\n",
    "Here are some helper functions that will help you create tensors and display useful information:\n",
    "   - `create_tensor`  creates a `jax numpy array` from a list of lists.\n",
    "   - `display_tensor` prints out the shape and the actual tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_tensor(t):\n",
    "    \"\"\"Create tensor from list of lists\"\"\"\n",
    "    return jnp.array(t)\n",
    "\n",
    "\n",
    "def display_tensor(t, name):\n",
    "    \"\"\"Display shape and tensor\"\"\"\n",
    "    print(f'{name} shape: {t.shape}\\n')\n",
    "    print(f'{t}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before implementing it yourself, you can play around with a toy example of `dot product attention` without the softmax  operation. Technically it would not be `dot product attention` without the softmax but this is done to avoid giving away too much of the answer and the idea is to display these tensors to give you a sense of how they look like.\n",
    "\n",
    "The formula for attention is this one:\n",
    "\n",
    "$$\n",
    "\\text { Attention }(Q, K, V)=\\operatorname{softmax}\\left(\\frac{Q K^{T}}{\\sqrt{d_{k}}}+{M}\\right) V\\tag{1}\\\n",
    "$$\n",
    "\n",
    "$d_{k}$ stands for the dimension of queries and keys.\n",
    "\n",
    "The `query`, `key`, `value` and `mask` vectors are provided for this example.\n",
    "\n",
    "Notice that the masking is done using very negative values that will yield a similar effect to using $-\\infty $. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 221
    },
    "colab_type": "code",
    "id": "_0x0HJXwRWQk",
    "outputId": "d6d78a8e-e3cc-47af-9584-2bdcdfcca0cd"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "query shape: (2, 3)\n",
      "\n",
      "[[1 0 0]\n",
      " [0 1 0]]\n",
      "\n",
      "key shape: (2, 3)\n",
      "\n",
      "[[1 2 3]\n",
      " [4 5 6]]\n",
      "\n",
      "value shape: (2, 3)\n",
      "\n",
      "[[0 1 0]\n",
      " [1 0 1]]\n",
      "\n",
      "mask shape: (2, 2)\n",
      "\n",
      "[[ 0.e+00  0.e+00]\n",
      " [-1.e+09  0.e+00]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "q = create_tensor([[1, 0, 0], [0, 1, 0]])\n",
    "display_tensor(q, 'query')\n",
    "k = create_tensor([[1, 2, 3], [4, 5, 6]])\n",
    "display_tensor(k, 'key')\n",
    "v = create_tensor([[0, 1, 0], [1, 0, 1]])\n",
    "display_tensor(v, 'value')\n",
    "m = create_tensor([[0, 0], [-1e9, 0]])\n",
    "display_tensor(m, 'mask')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "query shape: (2, 3)\n",
    "\n",
    "[[1 0 0]\n",
    " [0 1 0]]\n",
    "\n",
    "key shape: (2, 3)\n",
    "\n",
    "[[1 2 3]\n",
    " [4 5 6]]\n",
    "\n",
    "value shape: (2, 3)\n",
    "\n",
    "[[0 1 0]\n",
    " [1 0 1]]\n",
    "\n",
    "mask shape: (2, 2)\n",
    "\n",
    "[[ 0.e+00  0.e+00]\n",
    " [-1.e+09  0.e+00]]\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 102
    },
    "colab_type": "code",
    "id": "kVR9u4faRWQo",
    "outputId": "f01ea4ca-4152-4b54-b76a-e4b5917ae2b7"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "query dot key shape: (2, 2)\n",
      "\n",
      "[[0.57735026 2.309401  ]\n",
      " [1.1547005  2.8867514 ]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "q_dot_k = q @ k.T / jnp.sqrt(3)\n",
    "display_tensor(q_dot_k, 'query dot key')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "query dot key shape: (2, 2)\n",
    "\n",
    "[[0.57735026 2.309401  ]\n",
    " [1.1547005  2.8867514 ]]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "masked query dot key shape: (2, 2)\n",
      "\n",
      "[[ 5.7735026e-01  2.3094010e+00]\n",
      " [-1.0000000e+09  2.8867514e+00]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "masked = q_dot_k + m\n",
    "display_tensor(masked, 'masked query dot key')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "masked query dot key shape: (2, 2)\n",
    "\n",
    "[[ 5.7735026e-01  2.3094010e+00]\n",
    " [-1.0000000e+09  2.8867514e+00]]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "masked query dot key dot value shape: (2, 3)\n",
      "\n",
      "[[ 2.3094010e+00  5.7735026e-01  2.3094010e+00]\n",
      " [ 2.8867514e+00 -1.0000000e+09  2.8867514e+00]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "display_tensor(masked @ v, 'masked query dot key dot value')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "masked query dot key dot value shape: (2, 3)\n",
    "\n",
    "[[ 2.3094010e+00  5.7735026e-01  2.3094010e+00]\n",
    " [ 2.8867514e+00 -1.0000000e+09  2.8867514e+00]]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to use the previous dummy tensors to test some of the graded functions, a batch dimension should be added to them so they mimic the shape of real-life examples. The mask is also replaced by a version of it that resembles the one that is used by trax:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "query with batch dim shape: (1, 2, 3)\n",
      "\n",
      "[[[1 0 0]\n",
      "  [0 1 0]]]\n",
      "\n",
      "key with batch dim shape: (1, 2, 3)\n",
      "\n",
      "[[[1 2 3]\n",
      "  [4 5 6]]]\n",
      "\n",
      "value with batch dim shape: (1, 2, 3)\n",
      "\n",
      "[[[0 1 0]\n",
      "  [1 0 1]]]\n",
      "\n",
      "boolean mask shape: (2, 2)\n",
      "\n",
      "[[ True  True]\n",
      " [False  True]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "q_with_batch = q[None,:]\n",
    "display_tensor(q_with_batch, 'query with batch dim')\n",
    "k_with_batch = k[None,:]\n",
    "display_tensor(k_with_batch, 'key with batch dim')\n",
    "v_with_batch = v[None,:]\n",
    "display_tensor(v_with_batch, 'value with batch dim')\n",
    "m_bool = create_tensor([[True, True], [False, True]])\n",
    "display_tensor(m_bool, 'boolean mask')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "query with batch dim shape: (1, 2, 3)\n",
    "\n",
    "[[[1 0 0]\n",
    "  [0 1 0]]]\n",
    "\n",
    "key with batch dim shape: (1, 2, 3)\n",
    "\n",
    "[[[1 2 3]\n",
    "  [4 5 6]]]\n",
    "\n",
    "value with batch dim shape: (1, 2, 3)\n",
    "\n",
    "[[[0 1 0]\n",
    "  [1 0 1]]]\n",
    "\n",
    "boolean mask shape: (2, 2)\n",
    "\n",
    "[[ True  True]\n",
    " [False  True]]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='ex01'></a>\n",
    "### Exercise 01\n",
    "\n",
    "**Instructions:** Implement the dot product attention. Concretely, implement the following equation\n",
    "\n",
    "\n",
    "$$\n",
    "\\text { Attention }(Q, K, V)=\\operatorname{softmax}\\left(\\frac{Q K^{T}}{\\sqrt{d_{k}}}+{M}\\right) V\\tag{1}\\\n",
    "$$\n",
    "\n",
    "$Q$ - query, \n",
    "$K$ - key, \n",
    "$V$ - values, \n",
    "$M$ - mask, \n",
    "${d_k}$ - depth/dimension of the queries and keys (used for scaling down)\n",
    "\n",
    "You can implement this formula either by `trax` numpy (trax.math.numpy) or regular `numpy` but it is recommended to use `jnp`.\n",
    "\n",
    "Something to take into consideration is that within trax, the masks are tensors of `True/False` values not 0's and $-\\infty$ as in the previous example. Within the graded function don't think of applying the mask by summing up matrices, instead use `jnp.where()` and treat the **mask as a tensor of boolean values with `False` for values that need to be masked and True for the ones that don't.**\n",
    "\n",
    "Also take into account that the real tensors are far more complex than the toy ones you just played with. Because of this avoid using shortened operations such as `@` for dot product or `.T` for transposing. Use `jnp.matmul()` and `jnp.swapaxes()` instead.\n",
    "\n",
    "This is the self-attention block for the transformer decoder. Good luck!  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kSauPt0NUl_o"
   },
   "outputs": [],
   "source": [
    "# UNQ_C1\n",
    "# GRADED FUNCTION: DotProductAttention\n",
    "def DotProductAttention(query, key, value, mask):\n",
    "    \"\"\"Dot product self-attention.\n",
    "    Args:\n",
    "        query (jax.interpreters.xla.DeviceArray): array of query representations with shape (L_q by d)\n",
    "        key (jax.interpreters.xla.DeviceArray): array of key representations with shape (L_k by d)\n",
    "        value (jax.interpreters.xla.DeviceArray): array of value representations with shape (L_k by d) where L_v = L_k\n",
    "        mask (jax.interpreters.xla.DeviceArray): attention-mask, gates attention with shape (L_q by L_k)\n",
    "\n",
    "    Returns:\n",
    "        jax.interpreters.xla.DeviceArray: Self-attention array for q, k, v arrays. (L_q by L_k)\n",
    "    \"\"\"\n",
    "\n",
    "    assert query.shape[-1] == key.shape[-1] == value.shape[-1], \"Embedding dimensions of q, k, v aren't all the same\"\n",
    "\n",
    "    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###\n",
    "    # Save depth/dimension of the query embedding for scaling down the dot product\n",
    "    depth = query.shape[-1] \n",
    "\n",
    "    # Calculate scaled query key dot product according to formula above\n",
    "    dots = jnp.matmul(query, jnp.swapaxes(key, -1, -2)) / jnp.sqrt(depth)\n",
    "    \n",
    "    # Apply the mask\n",
    "    if mask is not None: # The 'None' in this line does not need to be replaced\n",
    "        dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))\n",
    "    \n",
    "    # Softmax formula implementation\n",
    "    # Use trax.fastmath.logsumexp of dots to avoid underflow by division by large numbers\n",
    "    # Hint: Last axis should be used and keepdims should be True\n",
    "    # Note: softmax = e^(dots - logsumexp(dots)) = E^dots / sumexp(dots)\n",
    "    logsumexp = trax.fastmath.logsumexp(dots, axis=-1, keepdims=True)\n",
    "\n",
    "    # Take exponential of dots minus logsumexp to get softmax\n",
    "    # Use jnp.exp()\n",
    "    dots = jnp.exp(dots - logsumexp)\n",
    "\n",
    "    # Multiply dots by value to get self-attention\n",
    "    # Use jnp.matmul()\n",
    "    attention = jnp.matmul(dots, value)\n",
    "\n",
    "    ## END CODE HERE ###\n",
    "    \n",
    "    return attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "8o0K7VWKRWQw",
    "outputId": "1c51af3a-5f11-480f-b33b-419072d8298c"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[[0.8496746 , 0.15032545, 0.8496746 ],\n",
       "              [1.        , 0.        , 1.        ]]], dtype=float32)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DotProductAttention(q_with_batch, k_with_batch, v_with_batch, m_bool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "DeviceArray([[[0.8496746 , 0.15032545, 0.8496746 ],\n",
    "              [1.        , 0.        , 1.        ]]], dtype=float32)\n",
    "```    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2y2PSiLVRWQ2"
   },
   "source": [
    "<a name='2.2'></a>\n",
    "\n",
    "## 2.2 Causal Attention\n",
    "\n",
    "Now you are going to implement causal attention: multi-headed attention with a mask to attend only to words that occurred before. \n",
    "\n",
    "<img src = \"causal.png\">\n",
    "\n",
    "In the image above, a word can see everything that is before it, but not what is after it. To implement causal attention, you will have to transform vectors and do many reshapes. You will need to implement the functions below.\n",
    "\n",
    "\n",
    "<a name='ex02'></a>\n",
    "### Exercise 02\n",
    "\n",
    "Implement the following functions that will be needed for Causal Attention:\n",
    "\n",
    "- <span style='color:blue'> compute_attention_heads </span>: Gets an input $x$ of dimension (batch_size, seqlen, n_heads $\\times$ d_head) and splits the last (depth) dimension and stacks it to the zeroth dimension to allow matrix multiplication (batch_size $\\times$ n_heads, seqlen, d_head).\n",
    "- <span style='color:blue'> dot_product_self_attention </span>: Creates a mask matrix with `False` values above the diagonal and `True` values below and calls DotProductAttention which implements dot product self attention.\n",
    "- <span style='color:blue'> compute_attention_output </span>: Undoes compute_attention_heads by splitting first (vertical) dimension and stacking in the last (depth) dimension (batch_size, seqlen, n_heads $\\times$ d_head). These operations concatenate (stack/merge) the heads. \n",
    "\n",
    "Next there are some toy tensors which may serve to give you an idea of the data shapes and opperations involved in Causal Attention. They are also useful to test out your functions! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 442
    },
    "colab_type": "code",
    "id": "VRH67YcrRWQ3",
    "outputId": "847a9416-877a-4246-c738-0eacdf46de59"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "query matrix (2D tensor) shape: (2, 3)\n",
      "\n",
      "[[1 0 0]\n",
      " [0 1 0]]\n",
      "\n",
      "batch of two (multi-head) collections of query matrices (4D tensor) shape: (2, 2, 2, 3)\n",
      "\n",
      "[[[[1 0 0]\n",
      "   [0 1 0]]\n",
      "\n",
      "  [[1 0 0]\n",
      "   [0 1 0]]]\n",
      "\n",
      "\n",
      " [[[1 0 0]\n",
      "   [0 1 0]]\n",
      "\n",
      "  [[1 0 0]\n",
      "   [0 1 0]]]]\n",
      "\n",
      "one batch of concatenated heads of query matrices (3d tensor) shape: (1, 2, 6)\n",
      "\n",
      "[[[1 0 0 1 0 0]\n",
      "  [0 1 0 0 1 0]]]\n",
      "\n",
      "three batches of concatenated heads of query matrices (3d tensor) shape: (3, 2, 6)\n",
      "\n",
      "[[[1 0 0 1 0 0]\n",
      "  [0 1 0 0 1 0]]\n",
      "\n",
      " [[1 0 0 1 0 0]\n",
      "  [0 1 0 0 1 0]]\n",
      "\n",
      " [[1 0 0 1 0 0]\n",
      "  [0 1 0 0 1 0]]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "tensor2d = create_tensor(q)\n",
    "display_tensor(tensor2d, 'query matrix (2D tensor)')\n",
    "\n",
    "tensor4d2b = create_tensor([[q, q], [q, q]])\n",
    "display_tensor(tensor4d2b, 'batch of two (multi-head) collections of query matrices (4D tensor)')\n",
    "\n",
    "tensor3dc = create_tensor([jnp.concatenate([q, q], axis = -1)])\n",
    "display_tensor(tensor3dc, 'one batch of concatenated heads of query matrices (3d tensor)')\n",
    "\n",
    "tensor3dc3b = create_tensor([jnp.concatenate([q, q], axis = -1), jnp.concatenate([q, q], axis = -1), jnp.concatenate([q, q], axis = -1)])\n",
    "display_tensor(tensor3dc3b, 'three batches of concatenated heads of query matrices (3d tensor)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is important to know that the following 3 functions would normally be defined within the `CausalAttention` function further below. \n",
    "\n",
    "However this makes these functions harder to test. Because of this, these functions are shown individually using a `closure` (when necessary) that simulates them being inside of the `CausalAttention` function. This is done because they rely on some variables that can be accessed from within `CausalAttention`.\n",
    "\n",
    "### Support Functions\n",
    "\n",
    "<span style='color:blue'> compute_attention_heads </span>: Gets an input $x$ of dimension (batch_size, seqlen, n_heads $\\times$ d_head) and splits the last (depth) dimension and stacks it to the zeroth dimension to allow matrix multiplication (batch_size $\\times$ n_heads, seqlen, d_head).\n",
    "\n",
    "**For the closures you only have to fill the inner function.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# UNQ_C2\n",
    "# GRADED FUNCTION: compute_attention_heads_closure\n",
    "def compute_attention_heads_closure(n_heads, d_head):\n",
    "    \"\"\" Function that simulates environment inside CausalAttention function.\n",
    "    Args:\n",
    "        d_head (int):  dimensionality of heads.\n",
    "        n_heads (int): number of attention heads.\n",
    "    Returns:\n",
    "        function: compute_attention_heads function\n",
    "    \"\"\"\n",
    "\n",
    "    def compute_attention_heads(x):\n",
    "        \"\"\" Compute the attention heads.\n",
    "        Args:\n",
    "            x (jax.interpreters.xla.DeviceArray): tensor with shape (batch_size, seqlen, n_heads X d_head).\n",
    "        Returns:\n",
    "            jax.interpreters.xla.DeviceArray: reshaped tensor with shape (batch_size X n_heads, seqlen, d_head).\n",
    "        \"\"\"\n",
    "        ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###\n",
    "        \n",
    "        # Size of the x's batch dimension\n",
    "        batch_size = x.shape[0]\n",
    "        # Length of the sequence\n",
    "        # Should be size of x's first dimension without counting the batch dim\n",
    "        seqlen = x.shape[1]\n",
    "        # Reshape x using jnp.reshape()\n",
    "        # batch_size, seqlen, n_heads*d_head -> batch_size, seqlen, n_heads, d_head\n",
    "        x = jnp.reshape(x, (batch_size, seqlen, n_heads, d_head))\n",
    "        # Transpose x using jnp.transpose()\n",
    "        # batch_size, seqlen, n_heads, d_head -> batch_size, n_heads, seqlen, d_head\n",
    "        # Note that the values within the tuple are the indexes of the dimensions of x and you must rearrange them\n",
    "        x = jnp.transpose(x, (0, 2, 1, 3))\n",
    "        # Reshape x using jnp.reshape()\n",
    "        # batch_size, n_heads, seqlen, d_head -> batch_size*n_heads, seqlen, d_head\n",
    "        x = jnp.reshape(x, (-1, seqlen, d_head))\n",
    "        \n",
    "        ### END CODE HERE ###\n",
    "        \n",
    "        return x\n",
    "    \n",
    "    return compute_attention_heads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input tensor shape: (3, 2, 6)\n",
      "\n",
      "[[[1 0 0 1 0 0]\n",
      "  [0 1 0 0 1 0]]\n",
      "\n",
      " [[1 0 0 1 0 0]\n",
      "  [0 1 0 0 1 0]]\n",
      "\n",
      " [[1 0 0 1 0 0]\n",
      "  [0 1 0 0 1 0]]]\n",
      "\n",
      "output tensor shape: (6, 2, 3)\n",
      "\n",
      "[[[1 0 0]\n",
      "  [0 1 0]]\n",
      "\n",
      " [[1 0 0]\n",
      "  [0 1 0]]\n",
      "\n",
      " [[1 0 0]\n",
      "  [0 1 0]]\n",
      "\n",
      " [[1 0 0]\n",
      "  [0 1 0]]\n",
      "\n",
      " [[1 0 0]\n",
      "  [0 1 0]]\n",
      "\n",
      " [[1 0 0]\n",
      "  [0 1 0]]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "display_tensor(tensor3dc3b, \"input tensor\")\n",
    "result_cah = compute_attention_heads_closure(2,3)(tensor3dc3b)\n",
    "display_tensor(result_cah, \"output tensor\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "input tensor shape: (3, 2, 6)\n",
    "\n",
    "[[[1 0 0 1 0 0]\n",
    "  [0 1 0 0 1 0]]\n",
    "\n",
    " [[1 0 0 1 0 0]\n",
    "  [0 1 0 0 1 0]]\n",
    "\n",
    " [[1 0 0 1 0 0]\n",
    "  [0 1 0 0 1 0]]]\n",
    "\n",
    "output tensor shape: (6, 2, 3)\n",
    "\n",
    "[[[1 0 0]\n",
    "  [0 1 0]]\n",
    "\n",
    " [[1 0 0]\n",
    "  [0 1 0]]\n",
    "\n",
    " [[1 0 0]\n",
    "  [0 1 0]]\n",
    "\n",
    " [[1 0 0]\n",
    "  [0 1 0]]\n",
    "\n",
    " [[1 0 0]\n",
    "  [0 1 0]]\n",
    "\n",
    " [[1 0 0]\n",
    "  [0 1 0]]]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style='color:blue'> dot_product_self_attention </span>: Creates a mask matrix with `False` values above the diagonal and `True` values below and calls DotProductAttention which implements dot product self attention."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# UNQ_C3\n",
    "# GRADED FUNCTION: dot_product_self_attention\n",
    "def dot_product_self_attention(q, k, v):\n",
    "    \"\"\" Masked dot product self attention.\n",
    "    Args:\n",
    "        q (jax.interpreters.xla.DeviceArray): queries.\n",
    "        k (jax.interpreters.xla.DeviceArray): keys.\n",
    "        v (jax.interpreters.xla.DeviceArray): values.\n",
    "    Returns:\n",
    "        jax.interpreters.xla.DeviceArray: masked dot product self attention tensor.\n",
    "    \"\"\"\n",
    "    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###\n",
    "    \n",
    "    # Hint: mask size should be equal to L_q. Remember that q has shape (batch_size, L_q, d)\n",
    "    mask_size = q.shape[-2]\n",
    "\n",
    "    # Creates a matrix with ones below the diagonal and 0s above. It should have shape (1, mask_size, mask_size)\n",
    "    # Notice that 1's and 0's get casted to True/False by setting dtype to jnp.bool_\n",
    "    # Use jnp.tril() - Lower triangle of an array and jnp.ones()\n",
    "    mask = jnp.tril(jnp.ones((1, mask_size, mask_size), dtype=jnp.bool_), k=0)\n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "    \n",
    "    return DotProductAttention(q, k, v, mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[[0.        , 1.        , 0.        ],\n",
       "              [0.8496746 , 0.15032543, 0.8496746 ]]], dtype=float32)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dot_product_self_attention(q_with_batch, k_with_batch, v_with_batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "DeviceArray([[[0.        , 1.        , 0.        ],\n",
    "              [0.8496746 , 0.15032543, 0.8496746 ]]], dtype=float32)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style='color:blue'> compute_attention_output </span>: Undoes compute_attention_heads by splitting first (vertical) dimension and stacking in the last (depth) dimension (batch_size, seqlen, n_heads $\\times$ d_head). These operations concatenate (stack/merge) the heads. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# UNQ_C4\n",
    "# GRADED FUNCTION: compute_attention_output_closure\n",
    "def compute_attention_output_closure(n_heads, d_head):\n",
    "    \"\"\" Function that simulates environment inside CausalAttention function.\n",
    "    Args:\n",
    "        d_head (int):  dimensionality of heads.\n",
    "        n_heads (int): number of attention heads.\n",
    "    Returns:\n",
    "        function: compute_attention_output function\n",
    "    \"\"\"\n",
    "    \n",
    "    def compute_attention_output(x):\n",
    "        \"\"\" Compute the attention output.\n",
    "        Args:\n",
    "            x (jax.interpreters.xla.DeviceArray): tensor with shape (batch_size X n_heads, seqlen, d_head).\n",
    "        Returns:\n",
    "            jax.interpreters.xla.DeviceArray: reshaped tensor with shape (batch_size, seqlen, n_heads X d_head).\n",
    "        \"\"\"\n",
    "        ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###\n",
    "        \n",
    "        # Length of the sequence\n",
    "        # Should be size of x's first dimension without counting the batch dim\n",
    "        seqlen = x.shape[1]\n",
    "        # Reshape x using jnp.reshape() to shape (batch_size, n_heads, seqlen, d_head)\n",
    "        x = jnp.reshape(x, ( -1, n_heads, seqlen, d_head))\n",
    "        # Transpose x using jnp.transpose() to shape (batch_size, seqlen, n_heads, d_head)\n",
    "        x = jnp.transpose(x, ( 0, 2, 1 , 3))\n",
    "        \n",
    "        ### END CODE HERE ###\n",
    "        \n",
    "        # Reshape to allow to concatenate the heads\n",
    "        return jnp.reshape(x, (-1, seqlen, n_heads * d_head))\n",
    "    \n",
    "    return compute_attention_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input tensor shape: (6, 2, 3)\n",
      "\n",
      "[[[1 0 0]\n",
      "  [0 1 0]]\n",
      "\n",
      " [[1 0 0]\n",
      "  [0 1 0]]\n",
      "\n",
      " [[1 0 0]\n",
      "  [0 1 0]]\n",
      "\n",
      " [[1 0 0]\n",
      "  [0 1 0]]\n",
      "\n",
      " [[1 0 0]\n",
      "  [0 1 0]]\n",
      "\n",
      " [[1 0 0]\n",
      "  [0 1 0]]]\n",
      "\n",
      "output tensor shape: (3, 2, 6)\n",
      "\n",
      "[[[1 0 0 1 0 0]\n",
      "  [0 1 0 0 1 0]]\n",
      "\n",
      " [[1 0 0 1 0 0]\n",
      "  [0 1 0 0 1 0]]\n",
      "\n",
      " [[1 0 0 1 0 0]\n",
      "  [0 1 0 0 1 0]]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "display_tensor(result_cah, \"input tensor\")\n",
    "result_cao = compute_attention_output_closure(2,3)(result_cah)\n",
    "display_tensor(result_cao, \"output tensor\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "input tensor shape: (6, 2, 3)\n",
    "\n",
    "[[[1 0 0]\n",
    "  [0 1 0]]\n",
    "\n",
    " [[1 0 0]\n",
    "  [0 1 0]]\n",
    "\n",
    " [[1 0 0]\n",
    "  [0 1 0]]\n",
    "\n",
    " [[1 0 0]\n",
    "  [0 1 0]]\n",
    "\n",
    " [[1 0 0]\n",
    "  [0 1 0]]\n",
    "\n",
    " [[1 0 0]\n",
    "  [0 1 0]]]\n",
    "\n",
    "output tensor shape: (3, 2, 6)\n",
    "\n",
    "[[[1 0 0 1 0 0]\n",
    "  [0 1 0 0 1 0]]\n",
    "\n",
    " [[1 0 0 1 0 0]\n",
    "  [0 1 0 0 1 0]]\n",
    "\n",
    " [[1 0 0 1 0 0]\n",
    "  [0 1 0 0 1 0]]]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Causal Attention Function\n",
    "\n",
    "Now it is time for you to put everything together within the `CausalAttention` or Masked multi-head attention function:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src = \"masked-attention.png\"> \n",
    "\n",
    "**Instructions:** Implement the causal attention.\n",
    "Your model returns the causal attention through a $tl.Serial$ with the following:\n",
    "\n",
    "- <span style='color:blue'> [tl.Branch](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.combinators.Branch) </span>: consisting of 3 [tl.Dense(d_feature), ComputeAttentionHeads] to account for the queries, keys, and values.\n",
    "- <span style='color:blue'> [tl.Fn](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.base.Fn)</span>: Takes in dot_product_self_attention function and uses it to compute the dot product using $Q$, $K$, $V$.\n",
    "- <span style='color:blue'> [tl.Fn](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.base.Fn)</span>: Takes in compute_attention_output_closure to allow for parallel computing.\n",
    "- <span style='color:blue'> [tl.Dense](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dense)</span>: Final Dense layer, with dimension `d_feature`.\n",
    "\n",
    "Remember that in order for trax to properly handle the functions you just defined, they need to be added as layers using the [`tl.Fn()`](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.base.Fn) function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "B9Adn6DtRWRG"
   },
   "outputs": [],
   "source": [
    "# UNQ_C5\n",
    "# GRADED FUNCTION: CausalAttention\n",
    "def CausalAttention(d_feature, \n",
    "                    n_heads, \n",
    "                    compute_attention_heads_closure=compute_attention_heads_closure,\n",
    "                    dot_product_self_attention=dot_product_self_attention,\n",
    "                    compute_attention_output_closure=compute_attention_output_closure,\n",
    "                    mode='train'):\n",
    "    \"\"\"Transformer-style multi-headed causal attention.\n",
    "\n",
    "    Args:\n",
    "        d_feature (int):  dimensionality of feature embedding.\n",
    "        n_heads (int): number of attention heads.\n",
    "        compute_attention_heads_closure (function): Closure around compute_attention heads.\n",
    "        dot_product_self_attention (function): dot_product_self_attention function. \n",
    "        compute_attention_output_closure (function): Closure around compute_attention_output. \n",
    "        mode (str): 'train' or 'eval'.\n",
    "\n",
    "    Returns:\n",
    "        trax.layers.combinators.Serial: Multi-headed self-attention model.\n",
    "    \"\"\"\n",
    "    \n",
    "    assert d_feature % n_heads == 0\n",
    "    d_head = d_feature // n_heads\n",
    "\n",
    "    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###\n",
    "    \n",
    "    # HINT: The second argument to tl.Fn() is an uncalled function (without the parentheses)\n",
    "    # Since you are dealing with closures you might need to call the outer \n",
    "    # function with the correct parameters to get the actual uncalled function.\n",
    "    ComputeAttentionHeads = tl.Fn('AttnHeads', compute_attention_heads_closure(n_heads, d_head), n_out=1)\n",
    "        \n",
    "\n",
    "    return tl.Serial(\n",
    "        tl.Branch( # creates three towers for one input, takes activations and creates queries keys and values\n",
    "            [tl.Dense(d_feature), ComputeAttentionHeads], # queries\n",
    "            [tl.Dense(d_feature), ComputeAttentionHeads], # keys\n",
    "            [tl.Dense(d_feature), ComputeAttentionHeads], # values\n",
    "        ),\n",
    "        \n",
    "        tl.Fn('DotProductAttn', dot_product_self_attention, n_out=1), # takes QKV\n",
    "        # HINT: The second argument to tl.Fn() is an uncalled function\n",
    "        # Since you are dealing with closures you might need to call the outer \n",
    "        # function with the correct parameters to get the actual uncalled function.\n",
    "        tl.Fn('AttnOutput', compute_attention_output_closure(n_heads, d_head), n_out=1), # to allow for parallel\n",
    "        tl.Dense(d_feature)  # Final dense layer\n",
    "    )\n",
    "\n",
    "    ### END CODE HERE ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serial[\n",
      "  Branch_out3[\n",
      "    [Dense_512, AttnHeads]\n",
      "    [Dense_512, AttnHeads]\n",
      "    [Dense_512, AttnHeads]\n",
      "  ]\n",
      "  DotProductAttn_in3\n",
      "  AttnOutput\n",
      "  Dense_512\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "# Take a look at the causal attention model\n",
    "print(CausalAttention(d_feature=512, n_heads=8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "Serial[\n",
    "  Branch_out3[\n",
    "    [Dense_512, AttnHeads]\n",
    "    [Dense_512, AttnHeads]\n",
    "    [Dense_512, AttnHeads]\n",
    "  ]\n",
    "  DotProductAttn_in3\n",
    "  AttnOutput\n",
    "  Dense_512\n",
    "]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "W6zwtPjqRWRJ"
   },
   "source": [
    "<a name='2.3'></a>\n",
    "\n",
    "## 2.3 Transformer decoder block\n",
    "\n",
    "Now that you have implemented the causal part of the transformer, you will implement the transformer decoder block. Concretely you will be implementing this image now.\n",
    "\n",
    "<img src = \"transformer_decoder_1.png\" style = \"height:300px\"> \n",
    "\n",
    "To implement this function, you will have to call the `CausalAttention` or Masked multi-head attention function you implemented above. You will have to add a feedforward which consists of: \n",
    "\n",
    "- <span style='color:blue'> [tl.LayerNorm](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.normalization.LayerNorm) </span>: used to layer normalize\n",
    "- <span style='color:blue'> [tl.Dense](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dense) </span>: the dense layer\n",
    "- <span style='color:blue'> [ff_activation](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.activation_fns.Relu) </span>: feed forward activation (we use ReLu) here.\n",
    "- <span style='color:blue'> [tl.Dropout](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dropout) </span>: dropout layer\n",
    "- <span style='color:blue'> [tl.Dense](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dense) </span>: dense layer\n",
    "- <span style='color:blue'> [tl.Dropout](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dropout) </span>: dropout layer\n",
    "\n",
    "Finally once you implement the feedforward, you can go ahead and implement the entire block using: \n",
    "\n",
    "- <span style='color:blue'> [tl.Residual](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.combinators.Residual) </span>: takes in the tl.LayerNorm(), causal attention block, tl.dropout. \n",
    "\n",
    "- <span style='color:blue'> [tl.Residual](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.combinators.Residual) </span>: takes in the feedforward block you will implement. \n",
    "\n",
    "<a name='ex03'></a>\n",
    "### Exercise 03\n",
    "**Instructions:** Implement the transformer decoder block. Good luck!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gKOxnRbp1K5U"
   },
   "outputs": [],
   "source": [
    "# UNQ_C6\n",
    "# GRADED FUNCTION: DecoderBlock\n",
    "def DecoderBlock(d_model, d_ff, n_heads,\n",
    "                 dropout, mode, ff_activation):\n",
    "    \"\"\"Returns a list of layers that implements a Transformer decoder block.\n",
    "\n",
    "    The input is an activation tensor.\n",
    "\n",
    "    Args:\n",
    "        d_model (int):  depth of embedding.\n",
    "        d_ff (int): depth of feed-forward layer.\n",
    "        n_heads (int): number of attention heads.\n",
    "        dropout (float): dropout rate (how much to drop out).\n",
    "        mode (str): 'train' or 'eval'.\n",
    "        ff_activation (function): the non-linearity in feed-forward layer.\n",
    "\n",
    "    Returns:\n",
    "        list: list of trax.layers.combinators.Serial that maps an activation tensor to an activation tensor.\n",
    "    \"\"\"\n",
    "    \n",
    "    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###\n",
    "    \n",
    "    # Create masked multi-head attention block using CausalAttention function\n",
    "    causal_attention = CausalAttention( \n",
    "                        d_model,\n",
    "                        n_heads=n_heads,\n",
    "                        mode=mode\n",
    "                        )\n",
    "\n",
    "    # Create feed-forward block (list) with two dense layers with dropout and input normalized\n",
    "    feed_forward = [ \n",
    "        # Normalize layer inputs\n",
    "        tl.LayerNorm(),\n",
    "        # Add first feed forward (dense) layer (don't forget to set the correct value for n_units)\n",
    "        tl.Dense(d_ff),\n",
    "        # Add activation function passed in as a parameter (you need to call it!)\n",
    "        ff_activation(), # Generally ReLU\n",
    "        # Add dropout with rate and mode specified (i.e., don't use dropout during evaluation)\n",
    "        tl.Dropout(rate=dropout, mode=mode),\n",
    "        # Add second feed forward layer (don't forget to set the correct value for n_units)\n",
    "        tl.Dense(d_model),\n",
    "        # Add dropout with rate and mode specified (i.e., don't use dropout during evaluation)\n",
    "        tl.Dropout(rate=dropout,mode=mode)\n",
    "    ]\n",
    "\n",
    "    # Add list of two Residual blocks: the attention with normalization and dropout and feed-forward blocks\n",
    "    return [\n",
    "      tl.Residual(\n",
    "          # Normalize layer input\n",
    "          tl.LayerNorm(),\n",
    "          # Add causal attention block previously defined (without parentheses)\n",
    "          causal_attention,\n",
    "          # Add dropout with rate and mode specified\n",
    "          tl.Dropout(rate=dropout, mode=mode)\n",
    "        ),\n",
    "      tl.Residual(\n",
    "          # Add feed forward block (without parentheses)\n",
    "          feed_forward\n",
    "        ),\n",
    "      ]\n",
    "    ### END CODE HERE ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Serial[\n",
      "  Branch_out2[\n",
      "    None\n",
      "    Serial[\n",
      "      LayerNorm\n",
      "      Serial[\n",
      "        Branch_out3[\n",
      "          [Dense_512, AttnHeads]\n",
      "          [Dense_512, AttnHeads]\n",
      "          [Dense_512, AttnHeads]\n",
      "        ]\n",
      "        DotProductAttn_in3\n",
      "        AttnOutput\n",
      "        Dense_512\n",
      "      ]\n",
      "      Dropout\n",
      "    ]\n",
      "  ]\n",
      "  Add_in2\n",
      "], Serial[\n",
      "  Branch_out2[\n",
      "    None\n",
      "    Serial[\n",
      "      LayerNorm\n",
      "      Dense_2048\n",
      "      Relu\n",
      "      Dropout\n",
      "      Dense_512\n",
      "      Dropout\n",
      "    ]\n",
      "  ]\n",
      "  Add_in2\n",
      "]]\n"
     ]
    }
   ],
   "source": [
    "# Take a look at the decoder block\n",
    "print(DecoderBlock(d_model=512, d_ff=2048, n_heads=8, dropout=0.1, mode='train', ff_activation=tl.Relu))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "[Serial[\n",
    "  Branch_out2[\n",
    "    None\n",
    "    Serial[\n",
    "      LayerNorm\n",
    "      Serial[\n",
    "        Branch_out3[\n",
    "          [Dense_512, AttnHeads]\n",
    "          [Dense_512, AttnHeads]\n",
    "          [Dense_512, AttnHeads]\n",
    "        ]\n",
    "        DotProductAttn_in3\n",
    "        AttnOutput\n",
    "        Dense_512\n",
    "      ]\n",
    "      Dropout\n",
    "    ]\n",
    "  ]\n",
    "  Add_in2\n",
    "], Serial[\n",
    "  Branch_out2[\n",
    "    None\n",
    "    Serial[\n",
    "      LayerNorm\n",
    "      Dense_2048\n",
    "      Relu\n",
    "      Dropout\n",
    "      Dense_512\n",
    "      Dropout\n",
    "    ]\n",
    "  ]\n",
    "  Add_in2\n",
    "]]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SoFv-nfLRWRN",
    "lines_to_next_cell": 0
   },
   "source": [
    "<a name='2.4'></a>\n",
    "## 2.4 Transformer Language Model\n",
    "\n",
    "You will now bring it all together. In this part you will use all the subcomponents you previously built to make the final model. Concretely, here is the image you will be implementing. \n",
    "<img src = \"transformer_decoder.png\" style = \"height:400px\">\n",
    "\n",
    "    \n",
    "<a name='ex04'></a>\n",
    "### Exercise 04\n",
    "**Instructions:** Previously you coded the decoder block. Now you will code the transformer language model. Here is what you will need. \n",
    "\n",
    "- <span style=\"color:blue\"> positional_enconder </span>- a list containing the following layers:\n",
    "    - <span style=\"color:blue\"> [tl.Embedding](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Embedding)\n",
    "    - <span style=\"color:blue\"> [tl.Dropout](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dropout)\n",
    "    - <span style=\"color:blue\"> [tl.PositionalEncoding](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.attention.PositionalEncoding)\n",
    "\n",
    "- A list of `n_layers` <span style=\"color:blue\"> decoder blocks</span>.\n",
    "- <span style=\"color:blue\"> [tl.Serial](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.combinators.Serial): </span> takes in the following layers or lists of layers:\n",
    "    - <span style=\"color:blue\"> [tl.ShiftRight](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.attention.ShiftRight): </span>: shift the tensor to the right by padding on axis 1.\n",
    "    - <span style=\"color:blue\"> positional_encoder </span>: encodes the text positions.\n",
    "    - <span style=\"color:blue\"> decoder_blocks </span>: the ones you created.\n",
    "    - <span style=\"color:blue\"> [tl.LayerNorm](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.normalization.LayerNorm) </span>: a layer norm.\n",
    "    - <span style=\"color:blue\"> [tl.Dense](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dense) </span>: takes in the vocab_size.\n",
    "    - <span style=\"color:blue\"> [tl.LogSoftmax](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.LogSoftmax) </span>: to predict.\n",
    "    \n",
    "Go go go!! You can do it :)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0yi4LJO1RWRS"
   },
   "outputs": [],
   "source": [
    "# UNQ_C7\n",
    "# GRADED FUNCTION: TransformerLM\n",
    "def TransformerLM(vocab_size=33300,\n",
    "                  d_model=512,\n",
    "                  d_ff=2048,\n",
    "                  n_layers=6,\n",
    "                  n_heads=8,\n",
    "                  dropout=0.1,\n",
    "                  max_len=4096,\n",
    "                  mode='train',\n",
    "                  ff_activation=tl.Relu):\n",
    "    \"\"\"Returns a Transformer language model.\n",
    "\n",
    "    The input to the model is a tensor of tokens. (This model uses only the\n",
    "    decoder part of the overall Transformer.)\n",
    "\n",
    "    Args:\n",
    "        vocab_size (int): vocab size.\n",
    "        d_model (int):  depth of embedding.\n",
    "        d_ff (int): depth of feed-forward layer.\n",
    "        n_layers (int): number of decoder layers.\n",
    "        n_heads (int): number of attention heads.\n",
    "        dropout (float): dropout rate (how much to drop out).\n",
    "        max_len (int): maximum symbol length for positional encoding.\n",
    "        mode (str): 'train', 'eval' or 'predict', predict mode is for fast inference.\n",
    "        ff_activation (function): the non-linearity in feed-forward layer.\n",
    "\n",
    "    Returns:\n",
    "        trax.layers.combinators.Serial: A Transformer language model as a layer that maps from a tensor of tokens\n",
    "        to activations over a vocab set.\n",
    "    \"\"\"\n",
    "    \n",
    "    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###\n",
    "    \n",
    "    # Embedding inputs and positional encoder\n",
    "    positional_encoder = [ \n",
    "        # Add embedding layer of dimension (vocab_size, d_model)\n",
    "        tl.Embedding(vocab_size, d_model),\n",
    "        # Use dropout with rate and mode specified\n",
    "        tl.Dropout(rate=dropout, mode=mode),\n",
    "        # Add positional encoding layer with maximum input length and mode specified\n",
    "        tl.PositionalEncoding(max_len=max_len, mode=mode)]\n",
    "\n",
    "    # Create stack (list) of decoder blocks with n_layers with necessary parameters\n",
    "    decoder_blocks = [ \n",
    "        DecoderBlock(d_model, d_ff, n_heads,\n",
    "                    dropout, mode, ff_activation) for _ in range(n_layers)]\n",
    "\n",
    "    # Create the complete model as written in the figure\n",
    "    return tl.Serial(\n",
    "        # Use teacher forcing (feed output of previous step to current step)\n",
    "        tl.ShiftRight(mode=mode), # Specify the mode!\n",
    "        # Add positional encoder\n",
    "        positional_encoder,\n",
    "        # Add decoder blocks\n",
    "        decoder_blocks,\n",
    "        # Normalize layer\n",
    "        tl.LayerNorm(),\n",
    "\n",
    "        # Add dense layer of vocab_size (since need to select a word to translate to)\n",
    "        # (a.k.a., logits layer. Note: activation already set by ff_activation)\n",
    "        tl.Dense(vocab_size),\n",
    "        # Get probabilities with Logsoftmax\n",
    "        tl.LogSoftmax(),\n",
    "    )\n",
    "\n",
    "    ### END CODE HERE ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serial[\n",
      "  ShiftRight(1)\n",
      "  Embedding_33300_512\n",
      "  Dropout\n",
      "  PositionalEncoding\n",
      "  Serial[\n",
      "    Branch_out2[\n",
      "      None\n",
      "      Serial[\n",
      "        LayerNorm\n",
      "        Serial[\n",
      "          Branch_out3[\n",
      "            [Dense_512, AttnHeads]\n",
      "            [Dense_512, AttnHeads]\n",
      "            [Dense_512, AttnHeads]\n",
      "          ]\n",
      "          DotProductAttn_in3\n",
      "          AttnOutput\n",
      "          Dense_512\n",
      "        ]\n",
      "        Dropout\n",
      "      ]\n",
      "    ]\n",
      "    Add_in2\n",
      "  ]\n",
      "  Serial[\n",
      "    Branch_out2[\n",
      "      None\n",
      "      Serial[\n",
      "        LayerNorm\n",
      "        Dense_2048\n",
      "        Relu\n",
      "        Dropout\n",
      "        Dense_512\n",
      "        Dropout\n",
      "      ]\n",
      "    ]\n",
      "    Add_in2\n",
      "  ]\n",
      "  LayerNorm\n",
      "  Dense_33300\n",
      "  LogSoftmax\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "# Take a look at the Transformer\n",
    "print(TransformerLM(n_layers=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "Serial[\n",
    "  ShiftRight(1)\n",
    "  Embedding_33300_512\n",
    "  Dropout\n",
    "  PositionalEncoding\n",
    "  Serial[\n",
    "    Branch_out2[\n",
    "      None\n",
    "      Serial[\n",
    "        LayerNorm\n",
    "        Serial[\n",
    "          Branch_out3[\n",
    "            [Dense_512, AttnHeads]\n",
    "            [Dense_512, AttnHeads]\n",
    "            [Dense_512, AttnHeads]\n",
    "          ]\n",
    "          DotProductAttn_in3\n",
    "          AttnOutput\n",
    "          Dense_512\n",
    "        ]\n",
    "        Dropout\n",
    "      ]\n",
    "    ]\n",
    "    Add_in2\n",
    "  ]\n",
    "  Serial[\n",
    "    Branch_out2[\n",
    "      None\n",
    "      Serial[\n",
    "        LayerNorm\n",
    "        Dense_2048\n",
    "        Relu\n",
    "        Dropout\n",
    "        Dense_512\n",
    "        Dropout\n",
    "      ]\n",
    "    ]\n",
    "    Add_in2\n",
    "  ]\n",
    "  LayerNorm\n",
    "  Dense_33300\n",
    "  LogSoftmax\n",
    "]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dRRKnoAdvmJ7"
   },
   "source": [
    "<a name='3'></a>\n",
    "# Part 3: Training\n",
    "\n",
    "Now you are going to train your model. As usual, you have to define the cost function, the optimizer, and decide whether you will be training it on a `gpu` or `cpu`. In this case, you will train your model on a cpu for a few steps and we will load in a pre-trained model that you can use to predict with your own words."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "l1lkVebQRWRV"
   },
   "source": [
    "<a name='3.1'></a>\n",
    "### 3.1 Training the model\n",
    "\n",
    "You will now write a function that takes in your model and trains it. To train your model you have to decide how many times you want to iterate over the entire data set. Each iteration is defined as an `epoch`. For each epoch, you have to go over all the data, using your training iterator.\n",
    "\n",
    "<a name='ex05'></a>\n",
    "### Exercise 05\n",
    "**Instructions:** Implement the `train_model` program below to train the neural network above. Here is a list of things you should do:\n",
    "\n",
    "- Create the train task by calling [`trax.supervised.training.TrainTask`](https://trax-ml.readthedocs.io/en/latest/trax.supervised.html#trax.supervised.training.TrainTask) and pass in the following: \n",
    "    - <span style='color:blue'> labeled_data </span> = train_gen\n",
    "    - <span style='color:blue'> loss_fn </span> = [tl.CrossEntropyLoss()](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.metrics.CrossEntropyLoss)\n",
    "    - <span style='color:blue'> optimizer </span> = [trax.optimizers.Adam(0.01)](https://trax-ml.readthedocs.io/en/latest/trax.optimizers.html#trax.optimizers.adam.Adam)\n",
    "    - <span style='color:blue'> lr_schedule </span> = [lr_schedule](https://trax-ml.readthedocs.io/en/latest/trax.supervised.html#trax.supervised.lr_schedules.warmup_and_rsqrt_decay)\n",
    "\n",
    "\n",
    "- Create the eval task by calling [`trax.supervised.training.EvalTask`](https://trax-ml.readthedocs.io/en/latest/trax.supervised.html#trax.supervised.training.EvalTask) and pass in the following: \n",
    "    - <span style='color:blue'> labeled_data </span> = eval_gen\n",
    "    - <span style='color:blue'> metrics </span> = tl.CrossEntropyLoss() and [tl.Accuracy()](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.metrics.Accuracy)\n",
    "    \n",
    "    \n",
    "- Create the training loop by calling [`trax.supervised.Training.Loop`](https://trax-ml.readthedocs.io/en/latest/trax.supervised.html#trax.supervised.training.Loop) and pass in the following: \n",
    "    - <span style='color:blue'> TransformerLM </span> \n",
    "    - <span style='color:blue'> train_task </span> \n",
    "    - <span style='color:blue'> eval_task </span> = [eval_task]\n",
    "    - <span style='color:blue'> output_dir</span> = output_dir\n",
    "    \n",
    "You will be using a cross entropy loss, with Adam optimizer. Please read the [Trax](https://trax-ml.readthedocs.io/en/latest/index.html) documentation to get a full understanding. \n",
    "\n",
    "The training loop that this function returns can be runned using the `run()` method by passing in the desired number of steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gM2gpu4xvjtX"
   },
   "outputs": [],
   "source": [
    "from trax.supervised import training\n",
    "\n",
    "# UNQ_C8\n",
    "# GRADED FUNCTION: train_model\n",
    "def training_loop(TransformerLM, train_gen, eval_gen, output_dir = \"~/model\"):\n",
    "    '''\n",
    "    Input:\n",
    "        TransformerLM (trax.layers.combinators.Serial): The model you are building.\n",
    "        train_gen (generator): Training stream of data.\n",
    "        eval_gen (generator): Evaluation stream of data.\n",
    "        output_dir (str): folder to save your file.\n",
    "        \n",
    "    Returns:\n",
    "        trax.supervised.training.Loop: Training loop.\n",
    "    '''\n",
    "    output_dir = os.path.expanduser(output_dir)  # trainer is an object\n",
    "    lr_schedule = trax.lr.warmup_and_rsqrt_decay(n_warmup_steps=1000, max_value=0.01)\n",
    "\n",
    "    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###\n",
    "    train_task = training.TrainTask( \n",
    "      labeled_data=train_gen, # The training generator\n",
    "      loss_layer=tl.CrossEntropyLoss(), # Loss function \n",
    "      optimizer=trax.optimizers.Adam(0.01), # Optimizer (Don't forget to set LR to 0.01)\n",
    "      lr_schedule=lr_schedule,\n",
    "      n_steps_per_checkpoint=10\n",
    "    )\n",
    "\n",
    "    eval_task = training.EvalTask( \n",
    "      labeled_data=eval_gen, # The evaluation generator\n",
    "      metrics=[tl.CrossEntropyLoss(), tl.Accuracy()] # CrossEntropyLoss and Accuracy\n",
    "    )\n",
    "\n",
    "    ### END CODE HERE ###\n",
    "\n",
    "    loop = training.Loop(TransformerLM(d_model=4,\n",
    "                                       d_ff=16,\n",
    "                                       n_layers=1,\n",
    "                                       n_heads=2,\n",
    "                                       mode='train'),\n",
    "                         train_task,\n",
    "                         eval_tasks=[eval_task],\n",
    "                         output_dir=output_dir)\n",
    "    \n",
    "    return loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the model will be trained for only 10 steps. \n",
    "\n",
    "Even with this constraint the model with the original default arguments took a very long time to finish. Because of this some parameters are changed when defining the model that is fed into the training loop in the function above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 187
    },
    "colab_type": "code",
    "id": "BFRBTwSqRWRZ",
    "outputId": "aff859e5-8f4a-4d3b-f1d3-98e137581a77"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Step      1: Ran 1 train steps in 8.88 secs\n",
      "Step      1: train CrossEntropyLoss |  10.41675186\n",
      "Step      1: eval  CrossEntropyLoss |  10.41320705\n",
      "Step      1: eval          Accuracy |  0.00000000\n",
      "\n",
      "Step     10: Ran 9 train steps in 64.41 secs\n",
      "Step     10: train CrossEntropyLoss |  10.41439247\n",
      "Step     10: eval  CrossEntropyLoss |  10.41351700\n",
      "Step     10: eval          Accuracy |  0.00000000\n"
     ]
    }
   ],
   "source": [
    "# Should take around 1.5 minutes\n",
    "!rm -f ~/model/model.pkl.gz\n",
    "loop = training_loop(TransformerLM, train_batch_stream, eval_batch_stream)\n",
    "loop.run(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XKrEBjmskeWa"
   },
   "source": [
    " <a name='4'></a>\n",
    " # Part 4:  Evaluation  \n",
    "\n",
    "<a name='4.1'></a>\n",
    "### 4.1 Loading in a trained model\n",
    "\n",
    "In this part you will evaluate by loading in an almost exact version of the model you coded, but we trained it for you to save you time. Please run the cell below to load in the model.\n",
    "\n",
    "As you may have already noticed the model that you trained and the pretrained model share the same overall architecture but they have different values for some of the parameters:\n",
    "\n",
    "    \n",
    "   `Original (pretrained) model: `                                 \n",
    "                                       \n",
    "    TransformerLM(vocab_size=33300, d_model=512, d_ff=2048, n_layers=6, n_heads=8, \n",
    "                   dropout=0.1, max_len=4096, ff_activation=tl.Relu)\n",
    "                   \n",
    "   `Your model:`\n",
    "   \n",
    "    TransformerLM(d_model=4, d_ff=16, n_layers=1, n_heads=2)\n",
    "   \n",
    "   **Only the parameters shown for your model were changed. The others stayed the same.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 85
    },
    "colab_type": "code",
    "id": "zWoSzR5tkoAx",
    "outputId": "2b9f1cca-4778-4509-bd9e-bd1738625a4e"
   },
   "outputs": [],
   "source": [
    "# Get the model architecture\n",
    "model = TransformerLM(mode='eval')\n",
    "\n",
    "# Load the pre-trained weights\n",
    "model.init_from_file('model.pkl.gz', weights_only=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ilM9C8P3RWRf"
   },
   "source": [
    "<a name='5'></a>\n",
    "# Part 5: Testing with your own input\n",
    "\n",
    "You will now test your input. You are going to implement greedy decoding. This consists of two functions. The first one allows you to identify the next symbol. It gets the argmax of the output of your model and then returns that index. \n",
    "\n",
    "<a name='ex06'></a>\n",
    "### Exercise 06\n",
    "**Instructions:** Implement the next symbol function that takes in the cur_output_tokens and the trained model to return the index of the next word. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rD_bXRCpRWRg"
   },
   "outputs": [],
   "source": [
    "# UNQ_C9\n",
    "def next_symbol(cur_output_tokens, model):\n",
    "    \"\"\"Returns the next symbol for a given sentence.\n",
    "\n",
    "    Args:\n",
    "        cur_output_tokens (list): tokenized sentence with EOS and PAD tokens at the end.\n",
    "        model (trax.layers.combinators.Serial): The transformer model.\n",
    "\n",
    "    Returns:\n",
    "        int: tokenized symbol.\n",
    "    \"\"\"\n",
    "    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###\n",
    "    \n",
    "    # current output tokens length\n",
    "    token_length = len(cur_output_tokens)\n",
    "    # calculate the minimum power of 2 big enough to store token_length\n",
    "    # HINT: use np.ceil() and np.log2()\n",
    "    # add 1 to token_length so np.log2() doesn't receive 0 when token_length is 0\n",
    "    padded_length = 2**int(np.ceil(np.log2(token_length  + 1)))\n",
    "\n",
    "    # Fill cur_output_tokens with 0's until it reaches padded_length\n",
    "    padded = cur_output_tokens  + [0] * (padded_length  - token_length)\n",
    "    padded_with_batch = np.array(padded)[None, :] # Don't replace this 'None'! This is a way of setting the batch dim\n",
    "\n",
    "    # model expects a tuple containing two padded tensors (with batch)\n",
    "    output, _ = model((padded_with_batch, padded_with_batch)) \n",
    "    # HINT: output has shape (1, padded_length, vocab_size)\n",
    "    # To get log_probs you need to index output with 0 in the first dim\n",
    "    # token_length in the second dim and all of the entries for the last dim.\n",
    "    log_probs = output[0, token_length, :]\n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "    \n",
    "    return int(np.argmax(log_probs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The'"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test it out!\n",
    "sentence_test_nxt_symbl = \"I want to fly in the sky.\"\n",
    "detokenize([next_symbol(tokenize(sentence_test_nxt_symbl)+[0], model)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "'The'\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2AwrQFglRWRj"
   },
   "source": [
    "<a name='5.1'></a>\n",
    "### 5.1 Greedy decoding\n",
    "\n",
    "Now you will implement the greedy_decode algorithm that will call the `next_symbol` function. It takes in the input_sentence, the trained model and returns the decoded sentence. \n",
    "\n",
    "<a name='ex07'></a>\n",
    "### Exercise 07\n",
    "\n",
    "**Instructions**: Implement the greedy_decode algorithm. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6HwIdimiN0k2"
   },
   "outputs": [],
   "source": [
    "# UNQ_C10\n",
    "# Decoding functions.\n",
    "def greedy_decode(input_sentence, model):\n",
    "    \"\"\"Greedy decode function.\n",
    "\n",
    "    Args:\n",
    "        input_sentence (string): a sentence or article.\n",
    "        model (trax.layers.combinators.Serial): Transformer model.\n",
    "\n",
    "    Returns:\n",
    "        string: summary of the input.\n",
    "    \"\"\"\n",
    "    \n",
    "    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###\n",
    "    # Use tokenize()\n",
    "    cur_output_tokens = tokenize(input_sentence) + [0]\n",
    "    generated_output = [] \n",
    "    cur_output = 0 \n",
    "    EOS = 1 \n",
    "    \n",
    "    while cur_output != EOS:\n",
    "        # Get next symbol\n",
    "        cur_output = next_symbol(cur_output_tokens, model)\n",
    "        # Append next symbol to original sentence\n",
    "        cur_output_tokens.append(cur_output)\n",
    "        # Append next symbol to generated sentence\n",
    "        generated_output.append(cur_output)\n",
    "        print(detokenize(generated_output))\n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "    \n",
    "    return detokenize(generated_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 306
    },
    "colab_type": "code",
    "id": "9kHuIDGW1sOr",
    "outputId": "2525ca2c-4625-47c0-8456-f75598581993"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "It was a sunny day when I went to the market to buy some flowers. But\n",
      "I only found roses, not tulips. \n",
      "\n",
      ":\n",
      ": I\n",
      ": I just\n",
      ": I just found\n",
      ": I just found ros\n",
      ": I just found roses\n",
      ": I just found roses,\n",
      ": I just found roses, not\n",
      ": I just found roses, not tu\n",
      ": I just found roses, not tulips\n",
      ": I just found roses, not tulips\n",
      ": I just found roses, not tulips.\n",
      ": I just found roses, not tulips.<EOS>\n",
      ": I just found roses, not tulips.<EOS>\n"
     ]
    }
   ],
   "source": [
    "# Test it out on a sentence!\n",
    "test_sentence = \"It was a sunny day when I went to the market to buy some flowers. But I only found roses, not tulips.\"\n",
    "print(wrapper.fill(test_sentence), '\\n')\n",
    "print(greedy_decode(test_sentence, model))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "CA-279WI2D3G"
   },
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    ":\n",
    ": I\n",
    ": I just\n",
    ": I just found\n",
    ": I just found ros\n",
    ": I just found roses\n",
    ": I just found roses,\n",
    ": I just found roses, not\n",
    ": I just found roses, not tu\n",
    ": I just found roses, not tulips\n",
    ": I just found roses, not tulips\n",
    ": I just found roses, not tulips.\n",
    ": I just found roses, not tulips.<EOS>\n",
    ": I just found roses, not tulips.<EOS>\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "DYgX-mzjyUia",
    "outputId": "b901e164-48b3-4124-d21a-fe7443d15b79"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "It’s the posing craze sweeping the U.S. after being brought to fame by\n",
      "skier Lindsey Vonn, soccer star Omar Cummings, baseball player Albert\n",
      "Pujols - and even Republican politician Rick Perry. But now four\n",
      "students at Riverhead High School on Long Island, New York, have been\n",
      "suspended for dropping to a knee and taking up a prayer pose to mimic\n",
      "Denver Broncos quarterback Tim Tebow. Jordan Fulcoly, Wayne Drexel,\n",
      "Tyler Carroll and Connor Carroll were all suspended for one day\n",
      "because the ‘Tebowing’ craze was blocking the hallway and presenting a\n",
      "safety hazard to students. Scroll down for video. Banned: Jordan\n",
      "Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll (all pictured\n",
      "left) were all suspended for one day by Riverhead High School on Long\n",
      "Island, New York, for their tribute to Broncos quarterback Tim Tebow.\n",
      "Issue: Four of the pupils were suspended for one day because they\n",
      "allegedly did not heed to warnings that the 'Tebowing' craze at the\n",
      "school was blocking the hallway and presenting a safety hazard to\n",
      "students. \n",
      "\n",
      "Jordan\n",
      "Jordan Ful\n",
      "Jordan Fulcol\n",
      "Jordan Fulcoly\n",
      "Jordan Fulcoly,\n",
      "Jordan Fulcoly, Wayne\n",
      "Jordan Fulcoly, Wayne Dre\n",
      "Jordan Fulcoly, Wayne Drexe\n",
      "Jordan Fulcoly, Wayne Drexel\n",
      "Jordan Fulcoly, Wayne Drexel,\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day.\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not hee\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warn\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the '\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Te\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebow\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "cra\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocki\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the hall\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the hallway\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the hallway and\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the hallway and presenting\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the hallway and presenting a\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the hallway and presenting a safety\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the hallway and presenting a safety hazard\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the hallway and presenting a safety hazard to\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the hallway and presenting a safety hazard to\n",
      "students\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the hallway and presenting a safety hazard to\n",
      "students.\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the hallway and presenting a safety hazard to\n",
      "students.<EOS>\n",
      "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
      "suspended for one day. Four students were suspended for one day\n",
      "because they allegedly did not heed to warnings that the 'Tebowing'\n",
      "craze was blocking the hallway and presenting a safety hazard to\n",
      "students.<EOS>\n"
     ]
    }
   ],
   "source": [
    "# Test it out with a whole article!\n",
    "article = \"It’s the posing craze sweeping the U.S. after being brought to fame by skier Lindsey Vonn, soccer star Omar Cummings, baseball player Albert Pujols - and even Republican politician Rick Perry. But now four students at Riverhead High School on Long Island, New York, have been suspended for dropping to a knee and taking up a prayer pose to mimic Denver Broncos quarterback Tim Tebow. Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were all suspended for one day because the ‘Tebowing’ craze was blocking the hallway and presenting a safety hazard to students. Scroll down for video. Banned: Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll (all pictured left) were all suspended for one day by Riverhead High School on Long Island, New York, for their tribute to Broncos quarterback Tim Tebow. Issue: Four of the pupils were suspended for one day because they allegedly did not heed to warnings that the 'Tebowing' craze at the school was blocking the hallway and presenting a safety hazard to students.\"\n",
    "print(wrapper.fill(article), '\\n')\n",
    "print(greedy_decode(article, model))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "Jordan\n",
    "Jordan Ful\n",
    "Jordan Fulcol\n",
    "Jordan Fulcoly\n",
    "Jordan Fulcoly,\n",
    "Jordan Fulcoly, Wayne\n",
    "Jordan Fulcoly, Wayne Dre\n",
    "Jordan Fulcoly, Wayne Drexe\n",
    "Jordan Fulcoly, Wayne Drexel\n",
    "Jordan Fulcoly, Wayne Drexel,\n",
    ".\n",
    ".\n",
    ".\n",
    "\n",
    "Final summary:\n",
    "\n",
    "Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were\n",
    "suspended for one day. Four students were suspended for one day\n",
    "because they allegedly did not heed to warnings that the 'Tebowing'\n",
    "craze was blocking the hallway and presenting a safety hazard to\n",
    "students.<EOS>\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Congratulations on finishing this week's assignment!** You did a lot of work and now you should have a better understanding of the encoder part of Transformers and how Transformers can be used for text summarization.\n",
    "\n",
    "**Keep it up!**"
   ]
  }
 ],
 "metadata": {
  "coursera": {
   "schema_names": [
    "NLPC4-2"
   ]
  },
  "jupytext": {
   "encoding": "# -*- coding: utf-8 -*-",
   "formats": "ipynb,py:percent"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}