NLP_C4W1_Assignment_backup.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9cb49525",
   "metadata": {},
   "source": [
    "# Assignment 1: Neural Machine Translation\n",
    "\n",
    "Welcome to the first assignment of Course 4. Here, you will build an English-to-Portuguese neural machine translation (NMT) model using Long Short-Term Memory (LSTM) networks with attention. Machine translation is an important task in natural language processing and could be useful not only for translating one language to another but also for word sense disambiguation (e.g. determining whether the word \"bank\" refers to the financial bank, or the land alongside a river). Implementing this using just a Recurrent Neural Network (RNN) with LSTMs can work for short to medium length sentences but can result in vanishing gradients for very long sequences. To help with this, you will be adding an attention mechanism to allow the decoder to access all relevant parts of the input sentence regardless of its length. By completing this assignment, you will:\n",
    "\n",
    "- Implement an encoder-decoder system with attention\n",
    "- Build the NMT model from scratch using Tensorflow\n",
    "- Generate translations using greedy and Minimum Bayes Risk (MBR) decoding\n",
    "\n",
    "## Table of Contents\n",
    "- [1 - Data Preparation](#1)\n",
    "- [2 - NMT model with attention](#2)\n",
    "    - [Exercise 1 - Encoder](#ex1)\n",
    "    - [Exercise 2 - CrossAttention](#ex2)\n",
    "    - [Exercise 3 - Decoder](#ex3)    \n",
    "    - [Exercise 4 - Translator](#ex4)\n",
    "- [3 - Training](#3)\n",
    "- [4 - Using the model for inference ](#4)\n",
    "    - [Exercise 5 - translate](#ex5)\n",
    "- [5 - Minimum Bayes-Risk Decoding](#5)\n",
    "    - [Exercise 6 - rouge1_similarity](#ex6)\n",
    "    - [Exercise 7 - average_overlap](#ex7)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9ef370d",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # Setting this env variable prevents TF warnings from showing up\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from collections import Counter\n",
    "from utils import (sentences, train_data, val_data, english_vectorizer, portuguese_vectorizer, \n",
    "                   masked_loss, masked_acc, tokens_to_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8adb8fd6",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import w1_unittest"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e76be1dc",
   "metadata": {},
   "source": [
    "<a name=\"1\"></a>\n",
    "## 1. Data Preparation\n",
    "\n",
    "The text pre-processing bits have already been taken care of (if you are interested in this be sure to check the `utils.py` file). The steps performed can be summarized as:\n",
    "\n",
    "- Reading the raw data from the text files\n",
    "- Cleaning the data (using lowercase, adding space around punctuation, trimming whitespaces, etc)\n",
    "- Splitting it into training and validation sets\n",
    "- Adding the start-of-sentence and end-of-sentence tokens to every sentence\n",
    "- Tokenizing the sentences\n",
    "- Creating a Tensorflow dataset out of the tokenized sentences\n",
    "\n",
    "Take a moment to inspect the raw sentences:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "226033a1",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "portuguese_sentences, english_sentences = sentences\n",
    "\n",
    "print(f\"English (to translate) sentence:\\n\\n{english_sentences[-5]}\\n\")\n",
    "print(f\"Portuguese (translation) sentence:\\n\\n{portuguese_sentences[-5]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ba90eb9",
   "metadata": {},
   "source": [
    "You don't have much use for the raw sentences so delete them to save memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9f081b0",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "del portuguese_sentences\n",
    "del english_sentences\n",
    "del sentences"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2ff83d2",
   "metadata": {},
   "source": [
    "Notice that you imported an `english_vectorizer` and a `portuguese_vectorizer` from `utils.py`. These were created using [tf.keras.layers.TextVectorization](https://www.tensorflow.org/api_docs/python/tf/keras/layers/TextVectorization) and they provide interesting features such as ways to visualize the vocabulary and convert text into tokenized ids and vice versa. In fact, you can inspect the first ten words of the vocabularies for both languages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c1cfc17",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "print(f\"First 10 words of the english vocabulary:\\n\\n{english_vectorizer.get_vocabulary()[:10]}\\n\")\n",
    "print(f\"First 10 words of the portuguese vocabulary:\\n\\n{portuguese_vectorizer.get_vocabulary()[:10]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3152b075",
   "metadata": {},
   "source": [
    "Notice that the first 4 words are reserved for special words. In order, these are:\n",
    "\n",
    "- the empty string\n",
    "- a special token to represent an unknown word\n",
    "- a special token to represent the start of a sentence\n",
    "- a special token to represent the end of a sentence\n",
    "\n",
    "You can see how many words are in a vocabulary by using the `vocabulary_size` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5facaa0c",
   "metadata": {
    "deletable": false,
    "editable": false,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# Size of the vocabulary\n",
    "vocab_size_por = portuguese_vectorizer.vocabulary_size()\n",
    "vocab_size_eng = english_vectorizer.vocabulary_size()\n",
    "\n",
    "print(f\"Portuguese vocabulary is made up of {vocab_size_por} words\")\n",
    "print(f\"English vocabulary is made up of {vocab_size_eng} words\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53e4b615",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "You can define [tf.keras.layers.StringLookup](https://www.tensorflow.org/api_docs/python/tf/keras/layers/StringLookup) objects that will help you map from words to ids and vice versa. Do this for the portuguese vocabulary since this will be useful later on when you decode the predictions from your model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "218f7a36",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# This helps you convert from words to ids\n",
    "word_to_id = tf.keras.layers.StringLookup(\n",
    "    vocabulary=portuguese_vectorizer.get_vocabulary(), \n",
    "    mask_token=\"\", \n",
    "    oov_token=\"[UNK]\"\n",
    ")\n",
    "\n",
    "# This helps you convert from ids to words\n",
    "id_to_word = tf.keras.layers.StringLookup(\n",
    "    vocabulary=portuguese_vectorizer.get_vocabulary(),\n",
    "    mask_token=\"\",\n",
    "    oov_token=\"[UNK]\",\n",
    "    invert=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4af8b623",
   "metadata": {},
   "source": [
    "Try it out for the special tokens and a random word:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20076b9a",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "unk_id = word_to_id(\"[UNK]\")\n",
    "sos_id = word_to_id(\"[SOS]\")\n",
    "eos_id = word_to_id(\"[EOS]\")\n",
    "baunilha_id = word_to_id(\"baunilha\")\n",
    "\n",
    "print(f\"The id for the [UNK] token is {unk_id}\")\n",
    "print(f\"The id for the [SOS] token is {sos_id}\")\n",
    "print(f\"The id for the [EOS] token is {eos_id}\")\n",
    "print(f\"The id for baunilha (vanilla) is {baunilha_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f1d744c",
   "metadata": {},
   "source": [
    "Finally take a look at how the data that is going to be fed to the neural network looks like. Both `train_data` and `val_data` are of type `tf.data.Dataset` and are already arranged in batches of 64 examples. To get the first batch out of a tf dataset you can use the `take` method. To get the first example out of the batch you can slice the tensor and use the `numpy` method for nicer printing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "739777eb",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "for (to_translate, sr_translation), translation in train_data.take(1):\n",
    "    print(f\"Tokenized english sentence:\\n{to_translate[0, :].numpy()}\\n\\n\")\n",
    "    print(f\"Tokenized portuguese sentence (shifted to the right):\\n{sr_translation[0, :].numpy()}\\n\\n\")\n",
    "    print(f\"Tokenized portuguese sentence:\\n{translation[0, :].numpy()}\\n\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdd9ee3c",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "There are a couple of important details to notice.\n",
    "\n",
    "- Padding has already been applied to the tensors and the value used for this is 0\n",
    "- Each example consists of 3 different tensors:\n",
    "    - The sentence to translate\n",
    "    - The shifted-to-the-right translation\n",
    "    - The translation\n",
    "    \n",
    "The first two can be considered as the features, while the third one as the target. By doing this your model can perform Teacher Forcing as you saw in the lectures.\n",
    "\n",
    "Now it is time to begin coding!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd41cb52",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "<a name=\"2\"></a>\n",
    "## 2. NMT model with attention\n",
    "\n",
    "The model you will build uses an encoder-decoder architecture. This Recurrent Neural Network (RNN) takes in a tokenized version of a sentence in its encoder, then passes it on to the decoder for translation. As mentioned in the lectures, just using a a regular sequence-to-sequence model with LSTMs will work effectively for short to medium sentences but will start to degrade for longer ones. You can picture it like the figure below where all of the context of the input sentence is compressed into one vector that is passed into the decoder block. You can see how this will be an issue for very long sentences (e.g. 100 tokens or more) because the context of the first parts of the input will have very little effect on the final vector passed to the decoder.\n",
    "\n",
    "<img src='images/plain_rnn.png'>\n",
    "\n",
    "Adding an attention layer to this model avoids this problem by giving the decoder access to all parts of the input sentence. To illustrate, let's just use a 4-word input sentence as shown below. Remember that a hidden state is produced at each timestep of the encoder (represented by the orange rectangles). These are all passed to the attention layer and each are given a score given the current activation (i.e. hidden state) of the decoder. For instance, let's consider the figure below where the first prediction \"como\" is already made. To produce the next prediction, the attention layer will first receive all the encoder hidden states (i.e. orange rectangles) as well as the decoder hidden state when producing the word \"como\" (i.e. first green rectangle). Given this information, it will score each of the encoder hidden states to know which one the decoder should focus on to produce the next word. As a result of training, the model might have learned that it should align to the second encoder hidden state and subsequently assigns a high probability to the word \"você\". If we are using greedy decoding, we will output the said word as the next symbol, then restart the process to produce the next word until we reach an end-of-sentence prediction.\n",
    "\n",
    "<img src='images/attention_overview.png'>\n",
    "\n",
    "\n",
    "There are different ways to implement attention and the one we'll use for this assignment is the Scaled Dot Product Attention which has the form:\n",
    "\n",
    "$$Attention(Q, K, V) = softmax(\\frac{QK^T}{\\sqrt{d_k}})V$$\n",
    "\n",
    "You will dive deeper into this equation in the next week but for now, you can think of it as computing scores using queries (Q) and keys (K), followed by a multiplication of values (V) to get a context vector at a particular timestep of the decoder. This context vector is fed to the decoder RNN to get a set of probabilities for the next predicted word. The division by square root of the keys dimensionality ($\\sqrt{d_k}$) is for improving model performance and you'll also learn more about it next week. For our machine translation application, the encoder activations (i.e. encoder hidden states) will be the keys and values, while the decoder activations (i.e. decoder hidden states) will be the queries.\n",
    "\n",
    "You will see in the upcoming sections that this complex architecture and mechanism can be implemented with just a few lines of code. \n",
    "\n",
    "First you will define two important global variables:\n",
    "\n",
    "- The size of the vocabulary\n",
    "- The number of units in the LSTM layers (the same number will be used for all LSTM layers)\n",
    "\n",
    "In this assignment, the vocabulary sizes for English and Portuguese are the same. Therefore, we use a single constant VOCAB_SIZE throughout the notebook. While in other settings, vocabulary sizes could differ, that is not the case in our assignment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e484abf",
   "metadata": {
    "deletable": false,
    "editable": false,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "VOCAB_SIZE = 12000\n",
    "UNITS = 256"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc251965",
   "metadata": {},
   "source": [
    "<a name=\"ex1\"></a>\n",
    "## Exercise 1 - Encoder\n",
    "\n",
    "Your first exercise is to code the encoder part of the neural network. For this, complete the `Encoder` class below. Notice that in the constructor (the `__init__` method) you need to define all of the sublayers of the encoder and then use these sublayers during the forward pass (the `call` method).\n",
    "\n",
    "The encoder consists of the following layers:\n",
    "\n",
    "- [Embedding](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Embedding). For this layer you need to define the appropriate `input_dim` and `output_dim` and let it know that you are using '0' as padding, which can be done by using the appropriate value for the `mask_zero` parameter.\n",
    "    \n",
    "+ [Bidirectional](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Bidirectional) [LSTM](https://www.tensorflow.org/api_docs/python/tf/keras/layers/LSTM). In TF you can implement bidirectional behaviour for RNN-like layers. This part is already taken care of but you will need to specify the appropriate type of layer as well as its parameters. In particular you need to set the appropriate number of units and make sure that the LSTM returns the full sequence and not only the last output, which can be done by using the appropriate value for the `return_sequences` parameter.\n",
    "\n",
    "\n",
    "You need to define the forward pass using the syntax of TF's [functional API](https://www.tensorflow.org/guide/keras/functional_api). What this means is that you chain function calls together to define your network like this:\n",
    "\n",
    "```python\n",
    "encoder_input = keras.Input(shape=(28, 28, 1), name=\"original_img\")\n",
    "x = layers.Conv2D(16, 3, activation=\"relu\")(encoder_input)\n",
    "x = layers.MaxPooling2D(3)(x)\n",
    "x = layers.Conv2D(16, 3, activation=\"relu\")(x)\n",
    "encoder_output = layers.GlobalMaxPooling2D()(x)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1db0a1d",
   "metadata": {
    "deletable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# GRADED CLASS: Encoder\n",
    "class Encoder(tf.keras.layers.Layer):\n",
    "    def __init__(self, vocab_size, units):\n",
    "        \"\"\"Initializes an instance of this class\n",
    "\n",
    "        Args:\n",
    "            vocab_size (int): Size of the vocabulary\n",
    "            units (int): Number of units in the LSTM layer\n",
    "        \"\"\"\n",
    "        super(Encoder, self).__init__()\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "\n",
    "        self.embedding = tf.keras.layers.Embedding(  \n",
    "            input_dim=None,\n",
    "            output_dim=None,\n",
    "            mask_zero=None\n",
    "        )  \n",
    "\n",
    "        self.rnn = tf.keras.layers.Bidirectional(  \n",
    "            merge_mode=\"sum\",  \n",
    "            layer=tf.keras.layers.None(\n",
    "                units=None,\n",
    "                return_sequences=None\n",
    "            ),  \n",
    "        )  \n",
    "\n",
    "        ### END CODE HERE ###\n",
    "\n",
    "    def call(self, context):\n",
    "        \"\"\"Forward pass of this layer\n",
    "\n",
    "        Args:\n",
    "            context (tf.Tensor): The sentence to translate\n",
    "\n",
    "        Returns:\n",
    "            tf.Tensor: Encoded sentence to translate\n",
    "        \"\"\"\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "\n",
    "        # Pass the context through the embedding layer\n",
    "        x = None\n",
    "\n",
    "        # Pass the output of the embedding through the RNN\n",
    "        x = None\n",
    "\n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65034ffd",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# Do a quick check of your implementation\n",
    "\n",
    "# Create an instance of your class\n",
    "encoder = Encoder(VOCAB_SIZE, UNITS)\n",
    "\n",
    "# Pass a batch of sentences to translate from english to portuguese\n",
    "encoder_output = encoder(to_translate)\n",
    "\n",
    "print(f'Tensor of sentences in english has shape: {to_translate.shape}\\n')\n",
    "print(f'Encoder output has shape: {encoder_output.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a909aea1",
   "metadata": {},
   "source": [
    "##### __Expected Output__\n",
    "\n",
    "```\n",
    "Tensor of sentences in english has shape: (64, 14)\n",
    "\n",
    "Encoder output has shape: (64, 14, 256)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3031bb14",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Test your code!\n",
    "\n",
    "w1_unittest.test_encoder(Encoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1afe83f4",
   "metadata": {},
   "source": [
    "<a name=\"ex2\"></a>\n",
    "## Exercise 2 - CrossAttention\n",
    "\n",
    "Your next exercise is to code the layer that will perform cross attention between the original sentences and the translations. For this, complete the `CrossAttention` class below. Notice that in the constructor (the `__init__` method) you need to define all of the sublayers and then use these sublayers during the forward pass (the `call` method). For this particular case some of these bits are already taken care of.\n",
    "\n",
    "The cross attention consists of the following layers:\n",
    "\n",
    "- [MultiHeadAttention](https://www.tensorflow.org/api_docs/python/tf/keras/layers/MultiHeadAttention). For this layer you need to define the appropriate `key_dim`, which is the size of the key and query tensors. You will also need to set the number of heads to 1 since you aren't implementing multi head attention but attention between two tensors. The reason why this layer is preferred over [Attention](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Attention) is that it allows simpler code during the forward pass.\n",
    "    \n",
    "A couple of things to notice:\n",
    "- You need a way to pass both the output of the attention alongside the shifted-to-the-right translation (since this cross attention happens in the decoder side). For this you will use an [Add](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Add) layer so that the original dimension is preserved, which would not happen if you use something like a [Concatenate](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Concatenate) layer.\n",
    "\n",
    "+ Layer normalization is also performed for better stability of the network by using a [LayerNormalization](https://www.tensorflow.org/api_docs/python/tf/keras/layers/LayerNormalization) layer.\n",
    "\n",
    "- You don't need to worry about these last steps as these are already solved.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74e71f3d",
   "metadata": {
    "deletable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# GRADED CLASS: CrossAttention\n",
    "class CrossAttention(tf.keras.layers.Layer):\n",
    "    def __init__(self, units):\n",
    "        \"\"\"Initializes an instance of this class\n",
    "\n",
    "        Args:\n",
    "            units (int): Number of units in the LSTM layer\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "\n",
    "        self.mha = ( \n",
    "            tf.keras.layers.None(\n",
    "                key_dim=None,\n",
    "                num_heads=None\n",
    "            ) \n",
    "        )  \n",
    "\n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        self.layernorm = tf.keras.layers.LayerNormalization()\n",
    "        self.add = tf.keras.layers.Add()\n",
    "\n",
    "    def call(self, context, target):\n",
    "        \"\"\"Forward pass of this layer\n",
    "\n",
    "        Args:\n",
    "            context (tf.Tensor): Encoded sentence to translate\n",
    "            target (tf.Tensor): The embedded shifted-to-the-right translation\n",
    "\n",
    "        Returns:\n",
    "            tf.Tensor: Cross attention between context and target\n",
    "        \"\"\"\n",
    "        ### START CODE HERE ###\n",
    "\n",
    "        # Call the MH attention by passing in the query and value\n",
    "        # For this case the query should be the translation and the value the encoded sentence to translate\n",
    "        # Hint: Check the call arguments of MultiHeadAttention in the docs\n",
    "        attn_output = None(\n",
    "            query=None,\n",
    "            value=None\n",
    "        )  \n",
    "\n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        x = self.add([target, attn_output])\n",
    "\n",
    "        x = self.layernorm(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c62796f",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# Do a quick check of your implementation\n",
    "\n",
    "# Create an instance of your class\n",
    "attention_layer = CrossAttention(UNITS)\n",
    "\n",
    "# The attention layer expects the embedded sr-translation and the context\n",
    "# The context (encoder_output) is already embedded so you need to do this for sr_translation:\n",
    "sr_translation_embed = tf.keras.layers.Embedding(VOCAB_SIZE, output_dim=UNITS, mask_zero=True)(sr_translation)\n",
    "\n",
    "# Compute the cross attention\n",
    "attention_result = attention_layer(encoder_output, sr_translation_embed)\n",
    "\n",
    "print(f'Tensor of contexts has shape: {encoder_output.shape}')\n",
    "print(f'Tensor of translations has shape: {sr_translation_embed.shape}')\n",
    "print(f'Tensor of attention scores has shape: {attention_result.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41d4f99a",
   "metadata": {},
   "source": [
    "##### __Expected Output__\n",
    "\n",
    "```\n",
    "Tensor of contexts has shape: (64, 14, 256)\n",
    "Tensor of translations has shape: (64, 15, 256)\n",
    "Tensor of attention scores has shape: (64, 15, 256)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f658975",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Test your code!\n",
    "\n",
    "w1_unittest.test_cross_attention(CrossAttention)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa296ee2",
   "metadata": {},
   "source": [
    "<a name=\"ex3\"></a>\n",
    "## Exercise 3 - Decoder\n",
    "\n",
    "\n",
    "Now you will implement the decoder part of the neural network by completing the `Decoder` class below. Notice that in the constructor (the `__init__` method) you need to define all of the sublayers of the decoder and then use these sublayers during the forward pass (the `call` method).\n",
    "\n",
    "The decoder consists of the following layers:\n",
    "\n",
    "- [Embedding](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Embedding). For this layer you need to define the appropriate `input_dim` and `output_dim` and let it know that you are using '0' as padding, which can be done by using the appropriate value for the `mask_zero` parameter.\n",
    "  \n",
    "  \n",
    "+ Pre-attention [LSTM](https://www.tensorflow.org/api_docs/python/tf/keras/layers/LSTM). Unlike in the encoder in which you used a Bidirectional LSTM, here you will use a vanilla LSTM. Don't forget to set the appropriate number of units and make sure that the LSTM returns the full sequence and not only the last output, which can be done by using the appropriate value for the `return_sequences` parameter. It is very important that this layer returns the state since this will be needed for inference so make sure to set the `return_state` parameter accordingly. Notice that LSTM layers return state as a tuple of two tensors called `memory_state` and `carry_state`, **however these names have been changed to better reflect what you have seen in the lectures to `hidden_state` and `cell_state` respectively**.\n",
    "\n",
    "- The attention layer that performs cross attention between the sentence to translate and the right-shifted translation. Here you need to use the `CrossAttention` layer you defined in the previous exercise.\n",
    "\n",
    "+ Post-attention [LSTM](https://www.tensorflow.org/api_docs/python/tf/keras/layers/LSTM). Another LSTM layer. For this one you don't need it to return the state.\n",
    "\n",
    "- Finally a [Dense](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Dense) layer. This one should have the same number of units as the size of the vocabulary since you expect it to compute the logits for every possible word in the vocabulary. Make sure to use a `logsoftmax` activation function for this one, which you can get as [tf.nn.log_softmax](https://www.tensorflow.org/api_docs/python/tf/nn/log_softmax).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9639bdb",
   "metadata": {
    "deletable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# GRADED CLASS: Decoder\n",
    "class Decoder(tf.keras.layers.Layer):\n",
    "    def __init__(self, vocab_size, units):\n",
    "        \"\"\"Initializes an instance of this class\n",
    "\n",
    "        Args:\n",
    "            vocab_size (int): Size of the vocabulary\n",
    "            units (int): Number of units in the LSTM layer\n",
    "        \"\"\"\n",
    "        super(Decoder, self).__init__()\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "\n",
    "        # The embedding layer\n",
    "        self.embedding = tf.keras.layers.None(\n",
    "            input_dim=None,\n",
    "            output_dim=None,\n",
    "            mask_zero=None\n",
    "        )  \n",
    "\n",
    "        # The RNN before attention\n",
    "        self.pre_attention_rnn = tf.keras.layers.None(\n",
    "            units=None,\n",
    "            return_sequences=None,\n",
    "            return_state=None\n",
    "        )  \n",
    "\n",
    "        # The attention layer\n",
    "        self.attention = None(None)\n",
    "\n",
    "        # The RNN after attention\n",
    "        self.post_attention_rnn = tf.keras.layers.None(\n",
    "            units=None,\n",
    "            return_sequences=None\n",
    "        )  \n",
    "\n",
    "        # The dense layer with logsoftmax activation\n",
    "        self.output_layer = tf.keras.layers.None(\n",
    "            units=None,\n",
    "            activation=None\n",
    "        )  \n",
    "\n",
    "        ### END CODE HERE ###\n",
    "\n",
    "    def call(self, context, target, state=None, return_state=False):\n",
    "        \"\"\"Forward pass of this layer\n",
    "\n",
    "        Args:\n",
    "            context (tf.Tensor): Encoded sentence to translate\n",
    "            target (tf.Tensor): The shifted-to-the-right translation\n",
    "            state (list[tf.Tensor, tf.Tensor], optional): Hidden state of the pre-attention LSTM. Defaults to None.\n",
    "            return_state (bool, optional): If set to true return the hidden states of the LSTM. Defaults to False.\n",
    "\n",
    "        Returns:\n",
    "            tf.Tensor: The log_softmax probabilities of predicting a particular token\n",
    "        \"\"\"\n",
    "        ### START CODE HERE ###\n",
    "\n",
    "        # Get the embedding of the input\n",
    "        x = self.None(None)\n",
    "\n",
    "        # Pass the embedded input into the pre attention LSTM\n",
    "        # Hints:\n",
    "        # - The LSTM you defined earlier should return the output alongside the state (made up of two tensors)\n",
    "        # - Pass in the state to the LSTM (needed for inference)\n",
    "        x, hidden_state, cell_state = self.None(None, initial_state=None)\n",
    "\n",
    "        # Perform cross attention between the context and the output of the LSTM (in that order)\n",
    "        x = self.None(None, None)\n",
    "\n",
    "        # Do a pass through the post attention LSTM\n",
    "        x = self.None(None)\n",
    "\n",
    "        # Compute the logits\n",
    "        logits = self.None(None)\n",
    "\n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        if return_state:\n",
    "            return logits, [hidden_state, cell_state]\n",
    "\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6165cf2",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# Do a quick check of your implementation\n",
    "\n",
    "# Create an instance of your class\n",
    "decoder = Decoder(VOCAB_SIZE, UNITS)\n",
    "\n",
    "# Notice that you don't need the embedded version of sr_translation since this is done inside the class\n",
    "logits = decoder(encoder_output, sr_translation)\n",
    "\n",
    "print(f'Tensor of contexts has shape: {encoder_output.shape}')\n",
    "print(f'Tensor of right-shifted translations has shape: {sr_translation.shape}')\n",
    "print(f'Tensor of logits has shape: {logits.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f2b5d7d",
   "metadata": {},
   "source": [
    "##### __Expected Output__\n",
    "\n",
    "```\n",
    "Tensor of contexts has shape: (64, 14, 256)\n",
    "Tensor of right-shifted translations has shape: (64, 15)\n",
    "Tensor of logits has shape: (64, 15, 12000)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b61093a",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Test your code!\n",
    "\n",
    "w1_unittest.test_decoder(Decoder, CrossAttention)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dcce3a7",
   "metadata": {},
   "source": [
    "<a name=\"ex4\"></a>\n",
    "## Exercise 4 - Translator\n",
    "\n",
    "Now you have to put together all of the layers you previously coded into an actual model. For this, complete the `Translator` class below. Notice how unlike the Encoder and Decoder classes inherited from `tf.keras.layers.Layer`, the Translator class inherits from `tf.keras.Model`.\n",
    "\n",
    "Remember that `train_data` will yield a tuple with the sentence to translate and the shifted-to-the-right translation, which are the \"features\" of the model. This means that the inputs of your network will be tuples containing context and targets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "205fcf31",
   "metadata": {
    "deletable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# GRADED CLASS: Translator\n",
    "class Translator(tf.keras.Model):\n",
    "    def __init__(self, vocab_size, units):\n",
    "        \"\"\"Initializes an instance of this class\n",
    "\n",
    "        Args:\n",
    "            vocab_size (int): Size of the vocabulary\n",
    "            units (int): Number of units in the LSTM layer\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "\n",
    "        # Define the encoder with the appropriate vocab_size and number of units\n",
    "        self.encoder = None\n",
    "\n",
    "        # Define the decoder with the appropriate vocab_size and number of units\n",
    "        self.decoder = None\n",
    "\n",
    "        ### END CODE HERE ###\n",
    "\n",
    "    def call(self, inputs):\n",
    "        \"\"\"Forward pass of this layer\n",
    "\n",
    "        Args:\n",
    "            inputs (tuple(tf.Tensor, tf.Tensor)): Tuple containing the context (sentence to translate) and the target (shifted-to-the-right translation)\n",
    "\n",
    "        Returns:\n",
    "            tf.Tensor: The log_softmax probabilities of predicting a particular token\n",
    "        \"\"\"\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "\n",
    "        # In this case inputs is a tuple consisting of the context and the target, unpack it into single variables\n",
    "        context, target = None\n",
    "\n",
    "        # Pass the context through the encoder\n",
    "        encoded_context = None\n",
    "\n",
    "        # Compute the logits by passing the encoded context and the target to the decoder\n",
    "        logits = None\n",
    "\n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d4a231c",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# Do a quick check of your implementation\n",
    "\n",
    "# Create an instance of your class\n",
    "translator = Translator(VOCAB_SIZE, UNITS)\n",
    "\n",
    "# Compute the logits for every word in the vocabulary\n",
    "logits = translator((to_translate, sr_translation))\n",
    "\n",
    "print(f'Tensor of sentences to translate has shape: {to_translate.shape}')\n",
    "print(f'Tensor of right-shifted translations has shape: {sr_translation.shape}')\n",
    "print(f'Tensor of logits has shape: {logits.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3a162dd",
   "metadata": {},
   "source": [
    "##### __Expected Output__\n",
    "\n",
    "```\n",
    "Tensor of sentences to translate has shape: (64, 14)\n",
    "Tensor of right-shifted translations has shape: (64, 15)\n",
    "Tensor of logits has shape: (64, 15, 12000)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37009022",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "w1_unittest.test_translator(Translator, Encoder, Decoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f81bc228",
   "metadata": {},
   "source": [
    "<a name=\"3\"></a>\n",
    "## 3. Training\n",
    "\n",
    "Now that you have an untrained instance of the NMT model, it is time to train it. You can use the `compile_and_train` function below to achieve this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a61ef65",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "def compile_and_train(model, epochs=20, steps_per_epoch=500):\n",
    "    model.compile(optimizer=\"adam\", loss=masked_loss, metrics=[masked_acc, masked_loss])\n",
    "\n",
    "    history = model.fit(\n",
    "        train_data.repeat(),\n",
    "        epochs=epochs,\n",
    "        steps_per_epoch=steps_per_epoch,\n",
    "        validation_data=val_data,\n",
    "        validation_steps=50,\n",
    "        callbacks=[tf.keras.callbacks.EarlyStopping(patience=3)],\n",
    "    )\n",
    "\n",
    "    return model, history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87d9bf9f",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Train the translator (this takes some minutes so feel free to take a break)\n",
    "\n",
    "trained_translator, history = compile_and_train(translator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d23b9301",
   "metadata": {},
   "source": [
    "<a name=\"4\"></a>\n",
    "## 4. Using the model for inference \n",
    "\n",
    "\n",
    "Now that your model is trained you can use it for inference. To help you with this the `generate_next_token` function is provided. Notice that this function is meant to be used inside a for-loop, so you feed to it the information of the previous step to generate the information of the next step. In particular you need to keep track of the state of the pre-attention LSTM in the decoder and if you are done with the translation. Also notice that a `temperature` variable is introduced which determines how to select the next token given the predicted logits:  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "522f6b6f",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "def generate_next_token(decoder, context, next_token, done, state, temperature=0.0):\n",
    "    \"\"\"Generates the next token in the sequence\n",
    "\n",
    "    Args:\n",
    "        decoder (Decoder): The decoder\n",
    "        context (tf.Tensor): Encoded sentence to translate\n",
    "        next_token (tf.Tensor): The predicted next token\n",
    "        done (bool): True if the translation is complete\n",
    "        state (list[tf.Tensor, tf.Tensor]): Hidden states of the pre-attention LSTM layer\n",
    "        temperature (float, optional): The temperature that controls the randomness of the predicted tokens. Defaults to 0.0.\n",
    "\n",
    "    Returns:\n",
    "        tuple(tf.Tensor, np.float, list[tf.Tensor, tf.Tensor], bool): The next token, log prob of said token, hidden state of LSTM and if translation is done\n",
    "    \"\"\"\n",
    "    # Get the logits and state from the decoder\n",
    "    logits, state = decoder(context, next_token, state=state, return_state=True)\n",
    "    \n",
    "    # Trim the intermediate dimension \n",
    "    logits = logits[:, -1, :]\n",
    "        \n",
    "    # If temp is 0 then next_token is the argmax of logits\n",
    "    if temperature == 0.0:\n",
    "        next_token = tf.argmax(logits, axis=-1)\n",
    "        \n",
    "    # If temp is not 0 then next_token is sampled out of logits\n",
    "    else:\n",
    "        logits = logits / temperature\n",
    "        next_token = tf.random.categorical(logits, num_samples=1)\n",
    "    \n",
    "    # Trim dimensions of size 1\n",
    "    logits = tf.squeeze(logits)\n",
    "    next_token = tf.squeeze(next_token)\n",
    "    \n",
    "    # Get the logit of the selected next_token\n",
    "    logit = logits[next_token].numpy()\n",
    "    \n",
    "    # Reshape to (1,1) since this is the expected shape for text encoded as TF tensors\n",
    "    next_token = tf.reshape(next_token, shape=(1,1))\n",
    "    \n",
    "    # If next_token is End-of-Sentence token you are done\n",
    "    if next_token == eos_id:\n",
    "        done = True\n",
    "    \n",
    "    return next_token, logit, state, done"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "190d2d76",
   "metadata": {},
   "source": [
    "See how it works by running the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9937547a",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# PROCESS SENTENCE TO TRANSLATE AND ENCODE\n",
    "\n",
    "# A sentence you wish to translate\n",
    "eng_sentence = \"I love languages\"\n",
    "\n",
    "# Convert it to a tensor\n",
    "texts = tf.convert_to_tensor(eng_sentence)[tf.newaxis]\n",
    "\n",
    "# Vectorize it and pass it through the encoder\n",
    "context = english_vectorizer(texts).to_tensor()\n",
    "context = encoder(context)\n",
    "\n",
    "# SET STATE OF THE DECODER\n",
    "\n",
    "# Next token is Start-of-Sentence since you are starting fresh\n",
    "next_token = tf.fill((1,1), sos_id)\n",
    "\n",
    "# Hidden and Cell states of the LSTM can be mocked using uniform samples\n",
    "state = [tf.random.uniform((1, UNITS)), tf.random.uniform((1, UNITS))]\n",
    "\n",
    "# You are not done until next token is EOS token\n",
    "done = False\n",
    "\n",
    "# Generate next token\n",
    "next_token, logit, state, done = generate_next_token(decoder, context, next_token, done, state, temperature=0.5)\n",
    "print(f\"Next token: {next_token}\\nLogit: {logit:.4f}\\nDone? {done}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "170323dd",
   "metadata": {},
   "source": [
    "<a name=\"ex5\"></a>\n",
    "## Exercise 5 - translate\n",
    "\n",
    "Now you can put everything together to translate a given sentence. For this, complete the `translate` function below. This function will take care of the following steps: \n",
    "- Process the sentence to translate and encode it\n",
    "\n",
    "+ Set the initial state of the decoder\n",
    "\n",
    "- Get predictions of the next token (starting with the \\<SOS> token) for a maximum of iterations (in case the \\<EOS> token is never returned)\n",
    "    \n",
    "+ Return the translated text (as a string), the logit of the last iteration (this helps measure how certain was that the sequence was translated in its totality) and the translation in token format.\n",
    "\n",
    "\n",
    "Hints: \n",
    "\n",
    "- The previous cell provides a lot of insights on how this function should work, so if you get stuck refer to it.\n",
    "\n",
    "+ Some useful docs:\n",
    "    + [tf.newaxis](https://www.tensorflow.org/api_docs/python/tf#newaxis)\n",
    "\n",
    "    - [tf.fill](https://www.tensorflow.org/api_docs/python/tf/fill)\n",
    "\n",
    "    + [tf.zeros](https://www.tensorflow.org/api_docs/python/tf/zeros)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42c74f1f",
   "metadata": {
    "deletable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# GRADED FUNCTION: translate\n",
    "def translate(model, text, max_length=50, temperature=0.0):\n",
    "    \"\"\"Translate a given sentence from English to Portuguese\n",
    "\n",
    "    Args:\n",
    "        model (tf.keras.Model): The trained translator\n",
    "        text (string): The sentence to translate\n",
    "        max_length (int, optional): The maximum length of the translation. Defaults to 50.\n",
    "        temperature (float, optional): The temperature that controls the randomness of the predicted tokens. Defaults to 0.0.\n",
    "\n",
    "    Returns:\n",
    "        tuple(str, np.float, tf.Tensor): The translation, logit that predicted <EOS> token and the tokenized translation\n",
    "    \"\"\"\n",
    "    # Lists to save tokens and logits\n",
    "    tokens, logits = [], []\n",
    "\n",
    "    ### START CODE HERE ###\n",
    "    \n",
    "    # PROCESS THE SENTENCE TO TRANSLATE\n",
    "    \n",
    "    # Convert the original string into a tensor\n",
    "    text = tf.None(None)[tf.None]\n",
    "    \n",
    "    # Vectorize the text using the correct vectorizer\n",
    "    context = None(None).to_tensor()\n",
    "    \n",
    "    # Get the encoded context (pass the context through the encoder)\n",
    "    # Hint: Remember you can get the encoder by using model.encoder\n",
    "    context = None.None(None)\n",
    "    \n",
    "    # INITIAL STATE OF THE DECODER\n",
    "    \n",
    "    # First token should be SOS token with shape (1,1)\n",
    "    next_token = tf.None((None, None), None)\n",
    "    \n",
    "    # Initial hidden and cell states should be tensors of zeros with shape (1, UNITS)\n",
    "    state = [tf.None((None, None)), tf.None((None, None))]\n",
    "    \n",
    "    # You are done when you draw a EOS token as next token (initial state is False)\n",
    "    done = None\n",
    "\n",
    "    # Iterate for max_length iterations\n",
    "    for None in None(None):\n",
    "        # Generate the next token\n",
    "        try:\n",
    "            next_token, logit, state, done = None(\n",
    "                decoder=None,\n",
    "                context=None,\n",
    "                next_token=None,\n",
    "                done=None,\n",
    "                state=None,\n",
    "                temperature=None\n",
    "            )\n",
    "        except:\n",
    "             raise Exception(\"Problem generating the next token\")\n",
    "        \n",
    "        # If done then break out of the loop\n",
    "        if None:\n",
    "            None\n",
    "        \n",
    "        # Add next_token to the list of tokens\n",
    "        None.None(None)\n",
    "        \n",
    "        # Add logit to the list of logits\n",
    "        None.None(None)\n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "    \n",
    "    # Concatenate all tokens into a tensor\n",
    "    tokens = tf.concat(tokens, axis=-1)\n",
    "    \n",
    "    # Convert the translated tokens into text\n",
    "    translation = tf.squeeze(tokens_to_text(tokens, id_to_word))\n",
    "    translation = translation.numpy().decode()\n",
    "    \n",
    "    return translation, logits[-1], tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3525e8ba",
   "metadata": {},
   "source": [
    "Try your function with temperature of 0, which will yield a deterministic output and is equivalent to a greedy decoding:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daaea8c5",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Running this cell multiple times should return the same output since temp is 0\n",
    "\n",
    "temp = 0.0 \n",
    "original_sentence = \"I love languages\"\n",
    "\n",
    "translation, logit, tokens = translate(trained_translator, original_sentence, temperature=temp)\n",
    "\n",
    "print(f\"Temperature: {temp}\\n\\nOriginal sentence: {original_sentence}\\nTranslation: {translation}\\nTranslation tokens:{tokens}\\nLogit: {logit:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d05129b",
   "metadata": {},
   "source": [
    "Try your function with temperature of 0.7 (stochastic output):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e0697db",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Running this cell multiple times should return different outputs since temp is not 0\n",
    "# You can try different temperatures\n",
    "\n",
    "temp = 0.7\n",
    "original_sentence = \"I love languages\"\n",
    "\n",
    "translation, logit, tokens = translate(trained_translator, original_sentence, temperature=temp)\n",
    "\n",
    "print(f\"Temperature: {temp}\\n\\nOriginal sentence: {original_sentence}\\nTranslation: {translation}\\nTranslation tokens:{tokens}\\nLogit: {logit:.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3a9ea35",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "w1_unittest.test_translate(translate, trained_translator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba027524",
   "metadata": {},
   "source": [
    "<a name=\"5\"></a>\n",
    "## 5. Minimum Bayes-Risk Decoding\n",
    "\n",
    "As mentioned in the lectures, getting the most probable token at each step may not necessarily produce the best results. Another approach is to do Minimum Bayes Risk Decoding or MBR. The general steps to implement this are:\n",
    "\n",
    "- Take several random samples\n",
    "+ Score each sample against all other samples\n",
    "- Select the one with the highest score\n",
    "\n",
    "You will be building helper functions for these steps in the following sections.\n",
    "\n",
    "With the ability to generate different translations by setting different temperature values you can do what you saw in the lectures and generate a bunch of translations and then determine which one is the best candidate. You will now do this by using the provided `generate_samples` function. This function will return any desired number of candidate translations alongside the log-probability for each one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62301cd5",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "def generate_samples(model, text, n_samples=4, temperature=0.6):\n",
    "    \n",
    "    samples, log_probs = [], []\n",
    "\n",
    "    # Iterate for n_samples iterations\n",
    "    for _ in range(n_samples):\n",
    "        \n",
    "        # Save the logit and the translated tensor\n",
    "        _, logp, sample = translate(model, text, temperature=temperature)\n",
    "        \n",
    "        # Save the translated tensors\n",
    "        samples.append(np.squeeze(sample.numpy()).tolist())\n",
    "        \n",
    "        # Save the logits\n",
    "        log_probs.append(logp)\n",
    "                \n",
    "    return samples, log_probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06bd792c",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "samples, log_probs = generate_samples(trained_translator, 'I love languages')\n",
    "\n",
    "for s, l in zip(samples, log_probs):\n",
    "    print(f\"Translated tensor: {s} has logit: {l:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29b10677",
   "metadata": {},
   "source": [
    "## Comparing overlaps\n",
    "\n",
    "Now that you can generate multiple translations it is time to come up with a method to measure the goodness of each one. As you saw in the lectures, one way to achieve this is by comparing each sample against the others. \n",
    "\n",
    "There are several metrics you can use for this purpose, as shown in the lectures and you can try experimenting with any one of these. For this assignment, you will be calculating scores for **unigram overlaps**. \n",
    "\n",
    "One of these metrics is the widely used yet simple [Jaccard similarity](https://en.wikipedia.org/wiki/Jaccard_index) which gets the intersection over union of two sets. The `jaccard_similarity` function returns this metric for any pair of candidate and reference translations:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edb54a71",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "def jaccard_similarity(candidate, reference):\n",
    "        \n",
    "    # Convert the lists to sets to get the unique tokens\n",
    "    candidate_set = set(candidate)\n",
    "    reference_set = set(reference)\n",
    "    \n",
    "    # Get the set of tokens common to both candidate and reference\n",
    "    common_tokens = candidate_set.intersection(reference_set)\n",
    "    \n",
    "    # Get the set of all tokens found in either candidate or reference\n",
    "    all_tokens = candidate_set.union(reference_set)\n",
    "    \n",
    "    # Compute the percentage of overlap (divide the number of common tokens by the number of all tokens)\n",
    "    overlap = len(common_tokens) / len(all_tokens)\n",
    "        \n",
    "    return overlap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc3384bf",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "l1 = [1, 2, 3]\n",
    "l2 = [1, 2, 3, 4]\n",
    "\n",
    "js = jaccard_similarity(l1, l2)\n",
    "\n",
    "print(f\"jaccard similarity between lists: {l1} and {l2} is {js:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6997662",
   "metadata": {},
   "source": [
    "##### __Expected Output__\n",
    "\n",
    "```\n",
    "jaccard similarity between tensors: [1, 2, 3] and [1, 2, 3, 4] is 0.750\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2510e3d",
   "metadata": {},
   "source": [
    "<a name=\"ex6\"></a>\n",
    "## Exercise 6 - rouge1_similarity\n",
    "\n",
    "Jaccard similarity is good but a more commonly used metric in machine translation is the ROUGE score. For unigrams, this is called ROUGE-1 and as shown in the lectures, you can output the scores for both precision and recall when comparing two samples. To get the final score, you will want to compute the F1-score as given by:\n",
    "\n",
    "$$score = 2* \\frac{(precision * recall)}{(precision + recall)}$$\n",
    "\n",
    "For the implementation of the `rouge1_similarity` function you want to use the [Counter](https://docs.python.org/3/library/collections.html#collections.Counter) class from the Python standard library:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb2e0a00",
   "metadata": {
    "deletable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# GRADED FUNCTION: rouge1_similarity\n",
    "def rouge1_similarity(candidate, reference):\n",
    "    \"\"\"Computes the ROUGE 1 score between two token lists\n",
    "\n",
    "    Args:\n",
    "        candidate (list[int]): Tokenized candidate translation\n",
    "        reference (list[int]): Tokenized reference translation\n",
    "\n",
    "    Returns:\n",
    "        float: Overlap between the two token lists\n",
    "    \"\"\"\n",
    "    ### START CODE HERE ###\n",
    "    \n",
    "    # Make a frequency table of the candidate and reference tokens\n",
    "    # Hint: use the Counter class (already imported)\n",
    "    candidate_word_counts = None\n",
    "    reference_word_counts = None\n",
    "    \n",
    "    # Initialize overlap at 0\n",
    "    overlap = None\n",
    "    \n",
    "    # Iterate over the tokens in the candidate frequency table\n",
    "    # Hint: Counter is a subclass of dict and you can get the keys \n",
    "    #       out of a dict using the keys method like this: dict.keys()\n",
    "    for token in None:\n",
    "        \n",
    "        # Get the count of the current token in the candidate frequency table\n",
    "        # Hint: You can access the counts of a token as you would access values of a dictionary\n",
    "        token_count_candidate = None\n",
    "        \n",
    "        # Get the count of the current token in the reference frequency table\n",
    "        # Hint: You can access the counts of a token as you would access values of a dictionary\n",
    "        token_count_reference = None\n",
    "        \n",
    "        # Update the overlap by getting the minimum between the two token counts above\n",
    "        overlap += None\n",
    "    \n",
    "    # Compute the precision\n",
    "    # Hint: precision = overlap / (number of tokens in candidate list) \n",
    "    precision = None\n",
    "    \n",
    "    # Compute the recall\n",
    "    # Hint: recall = overlap / (number of tokens in reference list) \n",
    "    recall = None\n",
    "    \n",
    "    if precision + recall != 0:\n",
    "        # Compute the Rouge1 Score\n",
    "        # Hint: This is equivalent to the F1 score\n",
    "        f1_score = None\n",
    "        \n",
    "        return f1_score\n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "        \n",
    "    return 0 # If precision + recall = 0 then return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14bb5295",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "l1 = [1, 2, 3]\n",
    "l2 = [1, 2, 3, 4]\n",
    "\n",
    "r1s = rouge1_similarity(l1, l2)\n",
    "\n",
    "print(f\"rouge 1 similarity between lists: {l1} and {l2} is {r1s:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afb8c61a",
   "metadata": {},
   "source": [
    "##### __Expected Output__\n",
    "\n",
    "```\n",
    "rouge 1 similarity between lists: [1, 2, 3] and [1, 2, 3, 4] is 0.857\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a680132e",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "w1_unittest.test_rouge1_similarity(rouge1_similarity)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaf8a058",
   "metadata": {},
   "source": [
    "## Computing the Overall Score\n",
    "\n",
    "\n",
    "You will now build a function to generate the overall score for a particular sample. As mentioned in the lectures, you need to compare each sample with all other samples. For instance, if we generated 30 sentences, we will need to compare sentence 1 to sentences 2 through 30. Then, we compare sentence 2 to sentences 1 and 3 through 30, and so forth. At each step, we get the average score of all comparisons to get the overall score for a particular sample. To illustrate, these will be the steps to generate the scores of a 4-sample list.\n",
    "\n",
    "- Get similarity score between sample 1 and sample 2\n",
    "+ Get similarity score between sample 1 and sample 3\n",
    "- Get similarity score between sample 1 and sample 4\n",
    "+ Get average score of the first 3 steps. This will be the overall score of sample 1\n",
    "- Iterate and repeat until samples 1 to 4 have overall scores.\n",
    "\n",
    "\n",
    "The results will be stored in a dictionary for easy lookups.\n",
    "\n",
    "<a name=\"ex7\"></a>\n",
    "## Exercise 7 - average_overlap\n",
    "\n",
    "Complete the `average_overlap` function below which should implement the process described above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "142264ff",
   "metadata": {
    "deletable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# GRADED FUNCTION: average_overlap\n",
    "def average_overlap(samples, similarity_fn):\n",
    "    \"\"\"Computes the arithmetic mean of each candidate sentence in the samples\n",
    "\n",
    "    Args:\n",
    "        samples (list[list[int]]): Tokenized version of translated sentences\n",
    "        similarity_fn (Function): Similarity function used to compute the overlap\n",
    "\n",
    "    Returns:\n",
    "        dict[int, float]: A dictionary mapping the index of each translation to its score\n",
    "    \"\"\"\n",
    "    # Initialize dictionary\n",
    "    scores = {}\n",
    "    \n",
    "    # Iterate through all samples (enumerate helps keep track of indexes)\n",
    "    for index_candidate, candidate in enumerate(samples):    \n",
    "        \n",
    "        ### START CODE HERE ###\n",
    "                \n",
    "        # Initially overlap is zero\n",
    "        overlap = None\n",
    "        \n",
    "        # Iterate through all samples (enumerate helps keep track of indexes)\n",
    "        for index_sample, sample in enumerate(samples):\n",
    "\n",
    "            # Skip if the candidate index is the same as the sample index\n",
    "            if None == None:\n",
    "                None\n",
    "                \n",
    "            # Get the overlap between candidate and sample using the similarity function\n",
    "            sample_overlap = None(None, None)\n",
    "            \n",
    "            # Add the sample overlap to the total overlap\n",
    "            overlap += None\n",
    "\n",
    "        ### END CODE HERE ###\n",
    "        \n",
    "        # Get the score for the candidate by computing the average\n",
    "        score = overlap / (len(samples) - 1)\n",
    "\n",
    "        # Only use 3 decimal points\n",
    "        score = round(score, 3)\n",
    "        \n",
    "        # Save the score in the dictionary. use index as the key.\n",
    "        scores[index_candidate] = score\n",
    "        \n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f36cf403",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# Test with Jaccard similarity\n",
    "\n",
    "l1 = [1, 2, 3]\n",
    "l2 = [1, 2, 4]\n",
    "l3 = [1, 2, 4, 5]\n",
    "\n",
    "avg_ovlp = average_overlap([l1, l2, l3], jaccard_similarity)\n",
    "\n",
    "print(f\"average overlap between lists: {l1}, {l2} and {l3} using Jaccard similarity is:\\n\\n{avg_ovlp}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e277aed2-a5c9-4ed0-9ee2-614939f2df7b",
   "metadata": {},
   "source": [
    "##### __Expected Output__\n",
    "\n",
    "```\n",
    "average overlap between lists: [1, 2, 3], [1, 2, 4] and [1, 2, 4, 5] using Jaccard similarity is:\n",
    "\n",
    "{0: 0.45, 1: 0.625, 2: 0.575}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d961a304-7c03-4ecb-ba5f-c8747ed3ec39",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "# Test with Rouge1 similarity\n",
    "\n",
    "l1 = [1, 2, 3]\n",
    "l2 = [1, 4]\n",
    "l3 = [1, 2, 4, 5]\n",
    "l4 = [5,6]\n",
    "\n",
    "avg_ovlp = average_overlap([l1, l2, l3, l4], rouge1_similarity)\n",
    "\n",
    "print(f\"average overlap between lists: {l1}, {l2}, {l3} and {l4} using Rouge1 similarity is:\\n\\n{avg_ovlp}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30adc749-ffcb-4e82-a8f0-c04a7e39da0a",
   "metadata": {},
   "source": [
    "##### __Expected Output__\n",
    "\n",
    "```\n",
    "average overlap between lists: [1, 2, 3], [1, 4], [1, 2, 4, 5] and [5, 6] using Rouge1 similarity is:\n",
    "\n",
    "{0: 0.324, 1: 0.356, 2: 0.524, 3: 0.111}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c41b1fba-fd0f-41e6-9b07-746f64030fe3",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "w1_unittest.test_average_overlap(average_overlap)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4482249",
   "metadata": {},
   "source": [
    "In practice, it is also common to see the weighted mean being used to calculate the overall score instead of just the arithmetic mean. This is implemented in the `weighted_avg_overlap` function below and you can use it in your experiments to see which one will give better results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "398714be",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "def weighted_avg_overlap(samples, log_probs, similarity_fn):\n",
    "    \n",
    "    # Scores dictionary\n",
    "    scores = {}\n",
    "    \n",
    "    # Iterate over the samples\n",
    "    for index_candidate, candidate in enumerate(samples):    \n",
    "        \n",
    "        # Initialize overlap and weighted sum\n",
    "        overlap, weight_sum = 0.0, 0.0\n",
    "        \n",
    "        # Iterate over all samples and log probabilities\n",
    "        for index_sample, (sample, logp) in enumerate(zip(samples, log_probs)):\n",
    "\n",
    "            # Skip if the candidate index is the same as the sample index            \n",
    "            if index_candidate == index_sample:\n",
    "                continue\n",
    "                \n",
    "            # Convert log probability to linear scale\n",
    "            sample_p = float(np.exp(logp))\n",
    "\n",
    "            # Update the weighted sum\n",
    "            weight_sum += sample_p\n",
    "\n",
    "            # Get the unigram overlap between candidate and sample\n",
    "            sample_overlap = similarity_fn(candidate, sample)\n",
    "            \n",
    "            # Update the overlap\n",
    "            overlap += sample_p * sample_overlap\n",
    "            \n",
    "        # Compute the score for the candidate\n",
    "        score = overlap / weight_sum\n",
    "\n",
    "        # Only use 3 decimal points\n",
    "        score = round(score, 3)\n",
    "        \n",
    "        # Save the score in the dictionary. use index as the key.\n",
    "        scores[index_candidate] = score\n",
    "    \n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3dfd6d3",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "l1 = [1, 2, 3]\n",
    "l2 = [1, 2, 4]\n",
    "l3 = [1, 2, 4, 5]\n",
    "log_probs = [0.4, 0.2, 0.5]\n",
    "\n",
    "w_avg_ovlp = weighted_avg_overlap([l1, l2, l3], log_probs, jaccard_similarity)\n",
    "\n",
    "print(f\"weighted average overlap using Jaccard similarity is:\\n\\n{w_avg_ovlp}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdb0b4db",
   "metadata": {},
   "source": [
    "## mbr_decode\n",
    "\n",
    "You will now put everything together in the the `mbr_decode` function below. This final step is not graded as this function is just a wrapper around all the cool stuff you have coded so far! \n",
    "\n",
    "You can use it to play around, trying different numbers of samples, temperatures and similarity functions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fcfa640",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": [
     "graded"
    ]
   },
   "outputs": [],
   "source": [
    "def mbr_decode(model, text, n_samples=5, temperature=0.6, similarity_fn=jaccard_similarity):\n",
    "    \n",
    "    # Generate samples\n",
    "    samples, log_probs = generate_samples(model, text, n_samples=n_samples, temperature=temperature)\n",
    "    \n",
    "    # Compute the overlap scores\n",
    "    scores = weighted_avg_overlap(samples, log_probs, similarity_fn)\n",
    "\n",
    "    # Decode samples\n",
    "    decoded_translations = [tokens_to_text(s, id_to_word).numpy().decode('utf-8') for s in samples]\n",
    "    \n",
    "    # Find the key with the highest score\n",
    "    max_score_key = max(scores, key=lambda k: scores[k])\n",
    "    \n",
    "    # Get the translation \n",
    "    translation = decoded_translations[max_score_key]\n",
    "    \n",
    "    return translation, decoded_translations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99507fcc-7727-45e7-933b-d3a08034f731",
   "metadata": {
    "deletable": false,
    "editable": false,
    "tags": []
   },
   "outputs": [],
   "source": [
    "english_sentence = \"I love languages\"\n",
    "\n",
    "translation, candidates = mbr_decode(trained_translator, english_sentence, n_samples=10, temperature=0.6)\n",
    "\n",
    "print(\"Translation candidates:\")\n",
    "for c in candidates:\n",
    "    print(c)\n",
    "\n",
    "print(f\"\\nSelected translation: {translation}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "801b193f-4ea6-4ca1-ae29-a506cce656d9",
   "metadata": {},
   "source": [
    "**Congratulations!** Next week, you'll dive deeper into attention models and study the Transformer architecture. You will build another network but without the recurrent part. It will show that attention is all you need! It should be fun!\n",
    "\n",
    "**Keep up the good work!**"
   ]
  }
 ],
 "metadata": {
  "grader_version": "1",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}