NLP_C4W2_Masking.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c2b5b44f",
   "metadata": {},
   "source": [
    "# Masking\n",
    "\n",
    "In this lab, you will implement the masking, that is one of the essential building blocks of the transformer. You will see how to define the masks and test how they work. You will use the masks later in the programming assignment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff0def87",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c0b0358",
   "metadata": {},
   "source": [
    "## 1 - Masking\n",
    "\n",
    "There are two types of masks that are useful when building your Transformer network: the *padding mask* and the *look-ahead mask*. Both help the softmax computation give the appropriate weights to the words in your input sentence. \n",
    "\n",
    "### 1.1 - Padding Mask\n",
    "\n",
    "Oftentimes your input sequence will exceed the maximum length of a sequence your network can process. Let's say the maximum length of your model is five, it is fed the following sequences:\n",
    "\n",
    "    [[\"Do\", \"you\", \"know\", \"when\", \"Jane\", \"is\", \"going\", \"to\", \"visit\", \"Africa\"], \n",
    "     [\"Jane\", \"visits\", \"Africa\", \"in\", \"September\" ],\n",
    "     [\"Exciting\", \"!\"]\n",
    "    ]\n",
    "\n",
    "which might get vectorized as:\n",
    "\n",
    "    [[ 71, 121, 4, 56, 99, 2344, 345, 1284, 15],\n",
    "     [ 56, 1285, 15, 181, 545],\n",
    "     [ 87, 600]\n",
    "    ]\n",
    "    \n",
    "When passing sequences into a transformer model, it is important that they are of uniform length. You can achieve this by padding the sequence with zeros, and truncating sentences that exceed the maximum length of your model:\n",
    "\n",
    "    [[ 71, 121, 4, 56, 99],\n",
    "     [ 2344, 345, 1284, 15, 0],\n",
    "     [ 56, 1285, 15, 181, 545],\n",
    "     [ 87, 600, 0, 0, 0],\n",
    "    ]\n",
    "    \n",
    "Sequences longer than the maximum length of five will be truncated, and zeros will be added to the truncated sequence to achieve uniform length. Similarly, for sequences shorter than the maximum length, zeros will also be added for padding.\n",
    "\n",
    "When pasing these vectors through the attention layers, the zeros will typically disappear  (you will get completely new vectors given the mathematical operations that happen in the attention block). However, you still want the network to attend only to the first few numbers in that vector (given by the sentence length) and this is when a padding mask comes in handy. You will need to define a boolean mask that specifies to which elements you must attend (1) and which elements you must ignore (0) and you do this by looking at all the zeros in the sequence. Then you use the mask to set the values of the vectors (corresponding to the zeros in the initial vector) close to negative infinity (-1e9).\n",
    "\n",
    "Imagine your input vector is `[87, 600, 0, 0, 0]`. This would give you a mask of `[1, 1, 0, 0, 0]`. When your vector passes through the attention mechanism, you get another (randomly looking) vector, let's say `[1, 2, 3, 4, 5]`, which after masking becomes `[1, 2, -1e9, -1e9, -1e9]`, so that when you take the softmax, the last three elements (where there were zeros in the input) don't affect the score.\n",
    "\n",
    "The [MultiheadAttention](https://keras.io/api/layers/attention_layers/multi_head_attention/) layer implemented in Keras, uses this masking logic.\n",
    "\n",
    "**Note:** The below function only creates the mask of an _already padded sequence_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4be6e26",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_padding_mask(decoder_token_ids):\n",
    "    \"\"\"\n",
    "    Creates a matrix mask for the padding cells\n",
    "    \n",
    "    Arguments:\n",
    "        decoder_token_ids (matrix like): matrix of size (n, m)\n",
    "    \n",
    "    Returns:\n",
    "        mask (tf.Tensor): binary tensor of size (n, 1, m)\n",
    "    \"\"\"    \n",
    "    seq = 1 - tf.cast(tf.math.equal(decoder_token_ids, 0), tf.float32)\n",
    "  \n",
    "    # add extra dimensions to add the padding\n",
    "    # to the attention logits. \n",
    "    # this will allow for broadcasting later when comparing sequences\n",
    "    return seq[:, tf.newaxis, :] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13a484c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.constant([[7., 6., 0., 0., 0.], [1., 2., 3., 0., 0.], [3., 0., 0., 0., 0.]])\n",
    "print(create_padding_mask(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce1d7106",
   "metadata": {},
   "source": [
    "If you multiply (1 - mask) by -1e9 and add it to the sample input sequences, the zeros are essentially set to negative infinity. Notice the difference when taking the softmax of the original sequence and the masked sequence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba9a0bdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the mask for x\n",
    "mask = create_padding_mask(x)\n",
    "\n",
    "# Extend the dimension of x to match the dimension of the mask\n",
    "x_extended = x[:, tf.newaxis, :]\n",
    "\n",
    "print(\"Softmax of non-masked vectors:\\n\")\n",
    "print(tf.keras.activations.softmax(x_extended))\n",
    "\n",
    "print(\"\\nSoftmax of masked vectors:\\n\")\n",
    "print(tf.keras.activations.softmax(x_extended + (1 - mask) * -1.0e9))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da92b367",
   "metadata": {},
   "source": [
    "### 1.2 - Look-ahead Mask\n",
    "\n",
    "The look-ahead mask follows similar intuition. In training, you will have access to the complete correct output of your training example. The look-ahead mask helps your model pretend that it correctly predicted a part of the output and see if, *without looking ahead*, it can correctly predict the next output. \n",
    "\n",
    "For example, if the expected correct output is `[1, 2, 3]` and you wanted to see if given that the model correctly predicted the first value it could predict the second value, you would mask out the second and third values. So you would input the masked sequence `[1, -1e9, -1e9]` and see if it could generate `[1, 2, -1e9]`.\n",
    "\n",
    "Just because you've worked so hard, we'll also implement this mask for you 😇😇. Again, take a close look at the code so you can effectively implement it later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "843dd37e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_look_ahead_mask(sequence_length):\n",
    "    \"\"\"\n",
    "    Returns a lower triangular matrix filled with ones\n",
    "    \n",
    "    Arguments:\n",
    "        sequence_length (int): matrix size\n",
    "    \n",
    "    Returns:\n",
    "        mask (tf.Tensor): binary tensor of size (sequence_length, sequence_length)\n",
    "    \"\"\"\n",
    "    mask = tf.linalg.band_part(tf.ones((1, sequence_length, sequence_length)), -1, 0)\n",
    "    return mask "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "393f4398",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.random.uniform((1, 3))\n",
    "temp = create_look_ahead_mask(x.shape[1])\n",
    "temp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50e1114d",
   "metadata": {},
   "source": [
    "**Congratulations on finishing this Lab!** Now you should have a better understanding of the masking in the transformer and this will surely help you with this week's assignment!\n",
    "\n",
    "**Keep it up!**"
   ]
  }
 ],
 "metadata": {
  "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}