assignments/C3_W4/C3_W4_Assignment.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-Jv7Y4hXwt0j"
   },
   "source": [
    "# Assignment 4:  Question duplicates\n",
    "\n",
    "Welcome to the fourth assignment of chapter 3. In this assignment you will explore Siamese networks applied to natural language processing. You will further explore the fundamentals of Trax and you will be able to implement a more complicated structure using it. By completing this assignment, you will learn how to implement models with different architectures. \n",
    "\n",
    "## Important Note on Submission to the AutoGrader\n",
    "\n",
    "Before submitting your assignment to the AutoGrader, please make sure you are not doing the following:\n",
    "\n",
    "1. You have not added any _extra_ `print` statement(s) in the assignment.\n",
    "2. You have not added any _extra_ code cell(s) in the assignment.\n",
    "3. You have not changed any of the function parameters.\n",
    "4. You are not using any global variables inside your graded exercises. Unless specifically instructed to do so, please refrain from it and use the local variables instead.\n",
    "5. You are not changing the assignment code where it is not required, like creating _extra_ variables.\n",
    "\n",
    "If you do any of the following, you will get something like, `Grader Error: Grader feedback not found` (or similarly unexpected) error upon submitting your assignment. Before asking for help/debugging the errors in your assignment, check for these first. If this is the case, and you don't remember the changes you have made."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Table of Contents\n",
    "\n",
    "- [Overview](#0)\n",
    "- [1 - Importing the Data](#1)\n",
    "    - [1.1 - Loading in the Data](#1-1)\n",
    "    - [1.2 - Converting a Question to a Tensor](#1-2)\n",
    "    - [1.3 - Understanding the Iterator](#1-3)\n",
    "        - [Exercise 1 - data_generator (UNQ_C1)](#ex-1)\n",
    "- [2 - Defining the Siamese Model](#2)\n",
    "    - [2.1 - Understanding Siamese Network](#2-1)\n",
    "        - [Exercise 2 - Siamese (UNQ_C2)](#ex-2)\n",
    "    - [2.2 - Hard Negative Mining](#2-2)\n",
    "        - [Exercise 3 - TripletLossFn (UNQ_C3)](#ex-3)\n",
    "- [3 - Training](#3)\n",
    "    - [3.1 - Training the Model](#3-1)\n",
    "        - [Exercise 4 - train_model (UNQ_C4)](#ex-4)\n",
    "- [4 - Evaluation](#4)\n",
    "    - [4.1 - Evaluating your Siamese Network](#4-1)\n",
    "    - [4.2 - Classify](#4-2)\n",
    "        - [Exercise 5 - classify (UNQ_C5)](#ex-5)\n",
    "- [5 -Testing with your Own Questions](#5)\n",
    "    - [Exercise 6 - predict (UNQ_C6)](#ex-6)\n",
    "- [On Siamese Networks](#6)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='0'></a>\n",
    "## Overview\n",
    "In this assignment, concretely you will: \n",
    "\n",
    "- Learn about Siamese networks\n",
    "- Understand how the triplet loss works\n",
    "- Understand how to evaluate accuracy\n",
    "- Use cosine similarity between the model's outputted vectors\n",
    "- Use the data generator to get batches of questions\n",
    "- Predict using your own model\n",
    "\n",
    "By now, you are familiar with trax and know how to make use of classes to define your model. We will start this homework by asking you to preprocess the data the same way you did in the previous assignments. After processing the data you will build a classifier that will allow you to identify whether two questions are the same or not.\n",
    "\n",
    "<img src = \"images/meme.png\" style=\"width:550px;height:300px;\"/>\n",
    "\n",
    "\n",
    "You will process the data first and then pad in a similar way you have done in the previous assignment. Your model will take in the two question embeddings, run them through an LSTM, and then compare the outputs of the two sub networks using cosine similarity. Before taking a deep dive into the model, start by importing the data set."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4sF9Hqzgwt0l"
   },
   "source": [
    "<a name='1'></a>\n",
    "## 1 - Importing the Data\n",
    "\n",
    "<a name='1-1'></a>\n",
    "### 1.1 - Loading in the Data\n",
    "\n",
    "You will be using the Quora question answer dataset to build a model that could identify similar questions. This is a useful task because you don't want to have several versions of the same question posted. Several times when teaching I end up responding to similar questions on piazza, or on other community forums. This data set has been labeled for you. Run the cell below to import some of the packages you will be using. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "zdACgs491cs2",
    "outputId": "b31042ef-845b-46b8-c783-185e96b135f7"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import nltk\n",
    "import trax\n",
    "from trax import layers as tl\n",
    "from trax.supervised import training\n",
    "from trax.fastmath import numpy as fastnp\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import random as rnd\n",
    "from trax import shapes\n",
    "\n",
    "import w4_unittest\n",
    "\n",
    "# set nltk path\n",
    "nltk.data.path.append('./nltk_data')\n",
    "\n",
    "# set random seeds\n",
    "rnd.seed(34)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3GYhQRMspitx"
   },
   "source": [
    "**Notice that for this assignment Trax's numpy is referred to as `fastnp`, while regular numpy is referred to as `np`.**\n",
    "\n",
    "You will now load in the data set. We have done some preprocessing for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 528
    },
    "colab_type": "code",
    "id": "sXWBVGWnpity",
    "outputId": "afa90d4d-fed7-43b8-bcba-48c95d600ad5"
   },
   "outputs": [],
   "source": [
    "data = pd.read_csv(\"data/questions.csv\")\n",
    "n=len(data)\n",
    "print('Number of question pairs: ', n)\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gkSQTu7Ypit0"
   },
   "source": [
    "We first split the data into a train and test set. The test set will be used later to evaluate our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "z00A7vEMpit1",
    "outputId": "c12ae7e8-a959-4f56-aa29-6ad34abc1c81"
   },
   "outputs": [],
   "source": [
    "N_train = 300000\n",
    "N_test  = 10*1024\n",
    "data_train = data[:N_train]\n",
    "data_test  = data[N_train:N_train+N_test]\n",
    "print(\"Train set:\", len(data_train), \"Test set:\", len(data_test))\n",
    "del(data) # remove to free memory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FbqIRRyEpit4"
   },
   "source": [
    "As explained in the lectures, we select only the question pairs that are duplicate to train the model. <br>\n",
    "We build two batches as input for the Siamese network and we assume that question $q1_i$ (question $i$ in the first batch) is a duplicate of $q2_i$ (question $i$ in the second batch), but all other questions in the second batch are not duplicates of $q1_i$.  \n",
    "The test set uses the original pairs of questions and the status describing if the questions are duplicates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "Xi_TwXxxpit4",
    "outputId": "f146046f-9c0d-4d8a-ecf8-8d6a4a5371f7"
   },
   "outputs": [],
   "source": [
    "td_index = (data_train['is_duplicate'] == 1).to_numpy()\n",
    "td_index = [i for i, x in enumerate(td_index) if x] \n",
    "print('number of duplicate questions: ', len(td_index))\n",
    "print('indexes of first ten duplicate questions:', td_index[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "id": "3I9oXSsKpit7",
    "outputId": "6f6bd3a1-219f-4fb3-a524-450c38bf44ba"
   },
   "outputs": [],
   "source": [
    "print(data_train['question1'][5])  #  Example of question duplicates (first one in data)\n",
    "print(data_train['question2'][5])\n",
    "print('is_duplicate: ', data_train['is_duplicate'][5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XHpZO58Dss_v"
   },
   "outputs": [],
   "source": [
    "Q1_train_words = np.array(data_train['question1'][td_index])\n",
    "Q2_train_words = np.array(data_train['question2'][td_index])\n",
    "\n",
    "Q1_test_words = np.array(data_test['question1'])\n",
    "Q2_test_words = np.array(data_test['question2'])\n",
    "y_test  = np.array(data_test['is_duplicate'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "P5vBkxunpiuB"
   },
   "source": [
    "Above, you have seen that you only took the duplicated questions for training our model. <br>You did so on purpose, because the data generator will produce batches $([q1_1, q1_2, q1_3, ...]$, $[q2_1, q2_2,q2_3, ...])$  where $q1_i$ and $q2_k$ are duplicate if and only if $i = k$.\n",
    "\n",
    "<br>Let's print to see what your data looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 170
    },
    "colab_type": "code",
    "id": "joyrS1XEpLWn",
    "outputId": "3257cde7-3164-40d9-910e-fa91eae917a0"
   },
   "outputs": [],
   "source": [
    "print('TRAINING QUESTIONS:\\n')\n",
    "print('Question 1: ', Q1_train_words[0])\n",
    "print('Question 2: ', Q2_train_words[0], '\\n')\n",
    "print('Question 1: ', Q1_train_words[5])\n",
    "print('Question 2: ', Q2_train_words[5], '\\n')\n",
    "\n",
    "print('TESTING QUESTIONS:\\n')\n",
    "print('Question 1: ', Q1_test_words[0])\n",
    "print('Question 2: ', Q2_test_words[0], '\\n')\n",
    "print('is_duplicate =', y_test[0], '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "WC_BZU3XpiuF"
   },
   "source": [
    "You will now encode each word of the selected duplicate pairs with an index. <br> Given a question, you can then just encode it as a list of numbers.  \n",
    "\n",
    "First you tokenize the questions using `nltk.word_tokenize`. <br>\n",
    "You need a python default dictionary which later, during inference, assigns the values $0$ to all Out Of Vocabulary (OOV) words.<br>\n",
    "Then you encode each word of the selected duplicate pairs with an index. Given a question, you can then just encode it as a list of numbers. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QbCoIgLQpiuF"
   },
   "outputs": [],
   "source": [
    "#create arrays\n",
    "Q1_train = np.empty_like(Q1_train_words)\n",
    "Q2_train = np.empty_like(Q2_train_words)\n",
    "\n",
    "Q1_test = np.empty_like(Q1_test_words)\n",
    "Q2_test = np.empty_like(Q2_test_words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "id": "m9ZmfpGWpiuI",
    "outputId": "d2995c9a-92b4-4892-d34b-c77b94b27134"
   },
   "outputs": [],
   "source": [
    "# Building the vocabulary with the train set         (this might take a minute)\n",
    "from collections import defaultdict\n",
    "\n",
    "vocab = defaultdict(lambda: 0)\n",
    "vocab['<PAD>'] = 1\n",
    "\n",
    "for idx in range(len(Q1_train_words)):\n",
    "    Q1_train[idx] = nltk.word_tokenize(Q1_train_words[idx])\n",
    "    Q2_train[idx] = nltk.word_tokenize(Q2_train_words[idx])\n",
    "    q = Q1_train[idx] + Q2_train[idx]\n",
    "    for word in q:\n",
    "        if word not in vocab:\n",
    "            vocab[word] = len(vocab) + 1\n",
    "print('The length of the vocabulary is: ', len(vocab))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "id": "TTMRF8eZpiuK",
    "outputId": "f81d4dc1-7cf9-4476-a454-467b54fe4dc4"
   },
   "outputs": [],
   "source": [
    "print(vocab['<PAD>'])\n",
    "print(vocab['Astrology'])\n",
    "print(vocab['Astronomy'])  #not in vocabulary, returns 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5sDs36m81g6f"
   },
   "outputs": [],
   "source": [
    "for idx in range(len(Q1_test_words)): \n",
    "    Q1_test[idx] = nltk.word_tokenize(Q1_test_words[idx])\n",
    "    Q2_test[idx] = nltk.word_tokenize(Q2_test_words[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "3QgGE9KlpiuP",
    "outputId": "19c3cf93-cf0d-4f8f-da99-e75481f16599"
   },
   "outputs": [],
   "source": [
    "print('Train set has reduced to: ', len(Q1_train) ) \n",
    "print('Test set length: ', len(Q1_test) ) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BDcxEmX31y3d"
   },
   "source": [
    "<a name='1-2'></a>\n",
    "### 1.2 - Converting a Question to a Tensor\n",
    "\n",
    "You will now convert every question to a tensor, or an array of numbers, using your vocabulary built above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zOhNa-sapiuS"
   },
   "outputs": [],
   "source": [
    "# Converting questions to array of integers\n",
    "for i in range(len(Q1_train)):\n",
    "    Q1_train[i] = [vocab[word] for word in Q1_train[i]]\n",
    "    Q2_train[i] = [vocab[word] for word in Q2_train[i]]\n",
    "\n",
    "        \n",
    "for i in range(len(Q1_test)):\n",
    "    Q1_test[i] = [vocab[word] for word in Q1_test[i]]\n",
    "    Q2_test[i] = [vocab[word] for word in Q2_test[i]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 102
    },
    "colab_type": "code",
    "id": "Dpawm38dpiuU",
    "outputId": "ef1aa65b-c89b-46f9-a9cf-f73748f1ee56"
   },
   "outputs": [],
   "source": [
    "print('first question in the train set:\\n')\n",
    "print(Q1_train_words[0], '\\n') \n",
    "print('encoded version:')\n",
    "print(Q1_train[0],'\\n')\n",
    "\n",
    "print('first question in the test set:\\n')\n",
    "print(Q1_test_words[0], '\\n')\n",
    "print('encoded version:')\n",
    "print(Q1_test[0]) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SuggGPaQpiuY"
   },
   "source": [
    "You will now split your train set into a training/validation set so that you can use it to train and evaluate your Siamese model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "id": "BmhrWPtgpiuY",
    "outputId": "7272fb74-79e6-499a-ce95-d11b9edcd64a"
   },
   "outputs": [],
   "source": [
    "# Splitting the data\n",
    "cut_off = int(len(Q1_train)*.8)\n",
    "train_Q1, train_Q2 = Q1_train[:cut_off], Q2_train[:cut_off]\n",
    "val_Q1, val_Q2 = Q1_train[cut_off: ], Q2_train[cut_off:]\n",
    "print('Number of duplicate questions: ', len(Q1_train))\n",
    "print(\"The length of the training set is:  \", len(train_Q1))\n",
    "print(\"The length of the validation set is: \", len(val_Q1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "iFOR19cX2TQs"
   },
   "source": [
    "<a name='1-3'></a>\n",
    "### 1.3 - Understanding the Iterator \n",
    "\n",
    "Most of the time in Natural Language Processing, and AI in general we use batches when training our data sets. If you were to use stochastic gradient descent with one example at a time, it will take you forever to build a model. In this example, we show you how you can build a data generator that takes in $Q1$ and $Q2$ and returns a batch of size `batch_size`  in the following format $([q1_1, q1_2, q1_3, ...]$, $[q2_1, q2_2,q2_3, ...])$. The tuple consists of two arrays and each array has `batch_size` questions. Again, $q1_i$ and $q2_i$ are duplicates, but they are not duplicates with any other elements in the batch. \n",
    "\n",
    "<br>\n",
    "\n",
    "The command ```next(data_generator)```returns the next batch. This iterator returns the data in a format that you could directly use in your model when computing the feed-forward of your algorithm. This iterator returns a pair of arrays of questions. \n",
    "\n",
    "<a name='ex-1'></a>\n",
    "### Exercise 1 - data_generator\n",
    "\n",
    "**Instructions:**  \n",
    "Implement the data generator below. Here are some things you will need. \n",
    "\n",
    "- While true loop.\n",
    "- if `index >= len_Q1`, set the `idx` to $0$.\n",
    "- The generator should return shuffled batches of data. To achieve this without modifying the actual question lists, a list containing the indexes of the questions is created. This list can be shuffled and used to get random batches everytime the index is reset.\n",
    "- Append elements of $Q1$ and $Q2$ to `input1` and `input2` respectively.\n",
    "- if `len(input1) == batch_size`, determine `max_len` as the longest question in `input1` and `input2`. Ceil `max_len` to a power of $2$ (for computation purposes) using the following command:  `max_len = 2**int(np.ceil(np.log2(max_len)))`.\n",
    "- Pad every question by `vocab['<PAD>']` until you get the length `max_len`.\n",
    "- Use yield to return `input1, input2`. \n",
    "- Don't forget to reset `input1, input2`  to empty arrays at the end (data generator resumes from where it last left)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ibchgos48MtA"
   },
   "outputs": [],
   "source": [
    "# UNQ_C1 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)\n",
    "# GRADED FUNCTION: data_generator\n",
    "def data_generator(Q1, Q2, batch_size, pad=1, shuffle=True):\n",
    "    \"\"\"Generator function that yields batches of data\n",
    "\n",
    "    Args:\n",
    "        Q1 (list): List of transformed (to tensor) questions.\n",
    "        Q2 (list): List of transformed (to tensor) questions.\n",
    "        batch_size (int): Number of elements per batch.\n",
    "        pad (int, optional): Pad character from the vocab. Defaults to 1.\n",
    "        shuffle (bool, optional): If the batches should be randomnized or not. Defaults to True.\n",
    "    Yields:\n",
    "        tuple: Of the form (input1, input2) with types (numpy.ndarray, numpy.ndarray)\n",
    "        NOTE: input1: inputs to your model [q1a, q2a, q3a, ...] i.e. (q1a,q1b) are duplicates\n",
    "              input2: targets to your model [q1b, q2b,q3b, ...] i.e. (q1a,q2i) i!=a are not duplicates\n",
    "    \"\"\"\n",
    "\n",
    "    input1 = []\n",
    "    input2 = []\n",
    "    idx = 0\n",
    "    len_q = len(Q1)\n",
    "    question_indexes = [*range(len_q)]\n",
    "    \n",
    "    if shuffle:\n",
    "        rnd.shuffle(question_indexes)\n",
    "    \n",
    "    ### START CODE HERE (Replace instances of 'None' with your code) ###\n",
    "    while True:\n",
    "        if idx >= len_q:\n",
    "            # if idx is greater than or equal to len_q, set idx accordingly \n",
    "            # (Hint: look at the instructions above)\n",
    "            idx = None\n",
    "            # shuffle to get random batches if shuffle is set to True\n",
    "            if shuffle:\n",
    "                rnd.shuffle(question_indexes) \n",
    "        \n",
    "        # get questions at the `question_indexes[idx]` position in Q1 and Q2\n",
    "        q1 = None\n",
    "        q2 = None\n",
    "        \n",
    "        # increment idx by 1\n",
    "        idx += None\n",
    "        # append q1\n",
    "        input1.append(None)\n",
    "        # append q2\n",
    "        input2.append(None)\n",
    "        if len(input1) == batch_size:\n",
    "            # determine max_len as the longest question in input1 & input 2\n",
    "            # Hint: use the `max` function. \n",
    "            # take max of input1 & input2 and then max out of the two of them.\n",
    "            max_len = None\n",
    "            # pad to power-of-2 (Hint: look at the instructions above)\n",
    "            max_len = None\n",
    "            b1 = [] \n",
    "            b2 = [] \n",
    "            for q1, q2 in zip(input1, input2):\n",
    "                # add [pad] to q1 until it reaches max_len\n",
    "                q1 = None\n",
    "                # add [pad] to q2 until it reaches max_len\n",
    "                q2 = None                \n",
    "                # append q1\n",
    "                b1.append(None)\n",
    "                # append q2\n",
    "                b2.append(None)\n",
    "            # use b1 and b2\n",
    "            yield np.array(None), np.array(None)\n",
    "    ### END CODE HERE ###\n",
    "            # reset the batches\n",
    "            input1, input2 = [], []  # reset the batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 204
    },
    "colab_type": "code",
    "id": "ZFZeBPnW8Mlb",
    "outputId": "7a31cd19-55dc-4b97-f288-6c59c6a34b53"
   },
   "outputs": [],
   "source": [
    "batch_size = 2\n",
    "res1, res2 = next(data_generator(train_Q1, train_Q2, batch_size))\n",
    "print(\"First questions  : \",'\\n', res1, '\\n')\n",
    "print(\"Second questions : \",'\\n', res2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tWJ1L9m2piui"
   },
   "source": [
    "**Note**: The following expected output is valid only if you run the above test cell **_once_** (first time). The output will change on each execution.\n",
    "\n",
    "If you think your implementation is correct and it is not matching the output, make sure to restart the kernel and run all the cells from the top again. \n",
    "\n",
    "**Expected Output:**\n",
    "```CPP\n",
    "First questions  :  \n",
    " [[  30   87   78  134 2132 1981   28   78  594   21    1    1    1    1\n",
    "     1    1]\n",
    " [  30   55   78 3541 1460   28   56  253   21    1    1    1    1    1\n",
    "     1    1]] \n",
    "\n",
    "Second questions :  \n",
    " [[  30  156   78  134 2132 9508   21    1    1    1    1    1    1    1\n",
    "     1    1]\n",
    " [  30  156   78 3541 1460  131   56  253   21    1    1    1    1    1\n",
    "     1    1]]\n",
    "```\n",
    "Now that you have your generator, you can just call it and it will return tensors which correspond to your questions in the Quora data set.<br>Now you can go ahead and start building your neural network. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test your function\n",
    "w4_unittest.test_data_generator(data_generator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KmZRBoaMwt0w"
   },
   "source": [
    "<a name='2'></a>\n",
    "## 2 - Defining the Siamese Model\n",
    "\n",
    "<a name='2-1'></a>\n",
    "### 2.1 - Understanding Siamese Network \n",
    "A Siamese network is a neural network which uses the same weights while working in tandem on two different input vectors to compute comparable output vectors.The Siamese network you are about to implement looks like this:\n",
    "\n",
    "<img src = \"images/siamese.png\" style=\"width:600px;height:300px;\"/>\n",
    "\n",
    "You get the question embedding, run it through an LSTM layer, normalize $v_1$ and $v_2$, and finally use a triplet loss (explained below) to get the corresponding cosine similarity for each pair of questions. As usual, you will start by importing the data set. The triplet loss makes use of a baseline (anchor) input that is compared to a positive (truthy) input and a negative (falsy) input. The distance from the baseline (anchor) input to the positive (truthy) input is minimized, and the distance from the baseline (anchor) input to the negative (falsy) input is maximized. In math equations, you are trying to maximize the following.\n",
    "\n",
    "$$\\mathcal{L}(A, P, N)=\\max \\left(\\|\\mathrm{f}(A)-\\mathrm{f}(P)\\|^{2}-\\|\\mathrm{f}(A)-\\mathrm{f}(N)\\|^{2}+\\alpha, 0\\right)$$\n",
    "\n",
    "$A$ is the anchor input, for example $q1_1$, $P$ the duplicate input, for example, $q2_1$, and $N$ the negative input (the non duplicate question), for example $q2_2$.<br>\n",
    "$\\alpha$ is a margin; you can think about it as a safety net, or by how much you want to push the duplicates from the non duplicates. \n",
    "<br>\n",
    "\n",
    "<a name='ex-2'></a>\n",
    "### Exercise 2 - Siamese\n",
    "\n",
    "**Instructions:** Implement the `Siamese` function below. You should be using all the objects explained below. \n",
    "\n",
    "To implement this model, you will be using `trax`. Concretely, you will be using the following functions.\n",
    "\n",
    "\n",
    "- `tl.Serial`: Combinator that applies layers serially (by function composition) allows you set up the overall structure of the feedforward. [docs](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.combinators.Serial) / [source code](https://github.com/google/trax/blob/1372b903bb66b0daccee19fd0b1fdf44f659330b/trax/layers/combinators.py#L26)\n",
    "    - You can pass in the layers as arguments to `Serial`, separated by commas. \n",
    "    - For example: `tl.Serial(tl.Embeddings(...), tl.Mean(...), tl.Dense(...), tl.LogSoftmax(...))` \n",
    "\n",
    "\n",
    "-  `tl.Embedding`: Maps discrete tokens to vectors. It will have shape (vocabulary length X dimension of output vectors). The dimension of output vectors (also called d_feature) is the number of elements in the word embedding. [docs](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Embedding) / [source code](https://github.com/google/trax/blob/1372b903bb66b0daccee19fd0b1fdf44f659330b/trax/layers/core.py#L113)\n",
    "    - `tl.Embedding(vocab_size, d_feature)`.\n",
    "    - `vocab_size` is the number of unique words in the given vocabulary.\n",
    "    - `d_feature` is the number of elements in the word embedding (some choices for a word embedding size range from 150 to 300, for example).\n",
    "\n",
    "\n",
    "-  `tl.LSTM` The LSTM layer. It leverages another Trax layer called [`LSTMCell`](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.rnn.LSTMCell). The number of units should be specified and should match the number of elements in the word embedding. [docs](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.rnn.LSTM) / [source code](https://github.com/google/trax/blob/1372b903bb66b0daccee19fd0b1fdf44f659330b/trax/layers/rnn.py#L87)\n",
    "    - `tl.LSTM(n_units)` Builds an LSTM layer of n_units.\n",
    "    \n",
    "    \n",
    "- `tl.Mean`: Computes the mean across a desired axis. Mean uses one tensor axis to form groups of values and replaces each group with the mean value of that group. [docs](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Mean) / [source code](https://github.com/google/trax/blob/1372b903bb66b0daccee19fd0b1fdf44f659330b/trax/layers/core.py#L276)\n",
    "    - `tl.Mean(axis=1)` mean over columns.\n",
    "\n",
    "\n",
    "- `tl.Fn` Layer with no weights that applies the function f, which should be specified using a lambda syntax. [docs](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.base.Fn) / [source doce](https://github.com/google/trax/blob/70f5364dcaf6ec11aabbd918e5f5e4b0f5bfb995/trax/layers/base.py#L576)\n",
    "    - $x$ -> This is used for cosine similarity.\n",
    "    - `tl.Fn('Normalize', lambda x: normalize(x))` Returns a layer with no weights that applies the function `f`\n",
    "    \n",
    "    \n",
    "- `tl.parallel`: It is a combinator layer (like `Serial`) that applies a list of layers in parallel to its inputs. [docs](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.combinators.Parallel) / [source code](https://github.com/google/trax/blob/37aba571a89a8ad86be76a569d0ec4a46bdd8642/trax/layers/combinators.py#L152)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hww76f8_wt0x"
   },
   "outputs": [],
   "source": [
    "# UNQ_C2 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)\n",
    "# GRADED FUNCTION: Siamese\n",
    "def Siamese(vocab_size=41699, d_model=128, mode='train'):\n",
    "    \"\"\"Returns a Siamese model.\n",
    "\n",
    "    Args:\n",
    "        vocab_size (int, optional): Length of the vocabulary. Defaults to len(vocab).\n",
    "        d_model (int, optional): Depth of the model. Defaults to 128.\n",
    "        mode (str, optional): 'train', 'eval' or 'predict', predict mode is for fast inference. Defaults to 'train'.\n",
    "\n",
    "    Returns:\n",
    "        trax.layers.combinators.Parallel: A Siamese model. \n",
    "    \"\"\"\n",
    "\n",
    "    def normalize(x):  # normalizes the vectors to have L2 norm 1\n",
    "        return x / fastnp.sqrt(fastnp.sum(x * x, axis=-1, keepdims=True))\n",
    "    \n",
    "    ### START CODE HERE (Replace instances of 'None' with your code) ###\n",
    "    q_processor = tl.Serial( # Processor will run on Q1 and Q2. \n",
    "        None, # Embedding layer\n",
    "        None, # LSTM layer\n",
    "        None, # Mean over columns\n",
    "        None, # Apply normalize function\n",
    "    )  # Returns one vector of shape [batch_size, d_model]. \n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "    \n",
    "    # Run on Q1 and Q2 in parallel.\n",
    "    model = tl.Parallel(q_processor, q_processor)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "es2gfwZypiul"
   },
   "source": [
    "Setup the Siamese network model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 255
    },
    "colab_type": "code",
    "id": "kvQ_jf52-JAn",
    "outputId": "d409460d-2ffb-4ae6-8745-ddcfa1d892ad"
   },
   "outputs": [],
   "source": [
    "# check your model\n",
    "model = Siamese()\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "LMK9zqhHpiuo"
   },
   "source": [
    "**Expected output:**  \n",
    "\n",
    "```CPP\n",
    "Parallel_in2_out2[\n",
    "  Serial[\n",
    "    Embedding_41699_128\n",
    "    LSTM_128\n",
    "    Mean\n",
    "    Normalize\n",
    "  ]\n",
    "  Serial[\n",
    "    Embedding_41699_128\n",
    "    LSTM_128\n",
    "    Mean\n",
    "    Normalize\n",
    "  ]\n",
    "]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test your function\n",
    "w4_unittest.test_Siamese(Siamese)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KVo1Gvripiuo"
   },
   "source": [
    "<a name='2-2'></a>\n",
    "### 2.2 - Hard Negative Mining\n",
    "\n",
    "\n",
    "You will now implement the `TripletLoss`.<br>\n",
    "As explained in the lecture, loss is composed of two terms. One term utilizes the mean of all the non duplicates, the second utilizes the *closest negative*. Our loss expression is then:\n",
    " \n",
    "\\begin{align}\n",
    " \\mathcal{Loss_{1}(A,P,N)} &=\\max \\left( -cos(A,P)  + mean_{neg} +\\alpha, 0\\right) \\\\\n",
    " \\mathcal{Loss_{2}(A,P,N)} &=\\max \\left( -cos(A,P)  + closest_{neg} +\\alpha, 0\\right) \\\\\n",
    "\\mathcal{Loss(A,P,N)} &= mean(Loss_1 + Loss_2) \\\\\n",
    "\\end{align}\n",
    "\n",
    "\n",
    "Further, two sets of instructions are provided. The first set provides a brief description of the task. If that set proves insufficient, a more detailed set can be displayed.  \n",
    "\n",
    "<a name='ex-3'></a>\n",
    "### Exercise 3 - TripletLossFn\n",
    "\n",
    "**Instructions (Brief):** Here is a list of things you should do: <br>\n",
    "\n",
    "- As this will be run inside trax, use `fastnp.xyz` when using any `xyz` numpy function\n",
    "- Use `fastnp.dot` to calculate the similarity matrix $v_1v_2^T$ of dimension `batch_size` x `batch_size`\n",
    "- Take the score of the duplicates on the diagonal `fastnp.diagonal`\n",
    "- Use the `trax` functions `fastnp.eye` and `fastnp.maximum` for the identity matrix and the maximum."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GWsX-Wz3piup"
   },
   "source": [
    "<details>  \n",
    "<summary>\n",
    "    <font size=\"3\" color=\"darkgreen\"><b>More Detailed Instructions </b></font>\n",
    "</summary>\n",
    "We'll describe the algorithm using a detailed example. Below, V1, V2 are the output of the normalization blocks in our model. Here we will use a batch_size of 4 and a d_model of 3. As explained in lecture, the inputs, Q1, Q2 are arranged so that corresponding inputs are duplicates while non-corresponding entries are not. The outputs will have the same pattern.\n",
    "<img src = \"images/C3_W4_triploss1.png\" style=\"width:1021px;height:229px;\"/>\n",
    "This testcase arranges the outputs, v1,v2, to highlight different scenarios. Here, the first two outputs V1[0], V2[0] match exactly - so the model is generating the same vector for Q1[0] and Q2[0] inputs. The second outputs differ, circled in orange, we set, V2[1] is set to match V2[2], simulating a model which is generating very poor results. V1[3] and V2[3] match exactly again while V1[4] and V2[4] are set to be exactly wrong - 180 degrees from each other, circled in blue. \n",
    "\n",
    "The first step is to compute the cosine similarity matrix or `score` in the code. As explained in lecture, this is $$V_1 V_2^T$$ This is generated with `fastnp.dot`.\n",
    "<img src = \"images/C3_W4_triploss2.png\" style=\"width:959px;height:236px;\"/>\n",
    "The clever arrangement of inputs creates the data needed for positive *and* negative examples without having to run all pair-wise combinations. Because Q1[n] is a duplicate of only Q2[n], other combinations are explicitly created negative examples or *Hard Negative* examples. The matrix multiplication efficiently produces the cosine similarity of all positive/negative combinations as shown above on the left side of the diagram. 'Positive' are the results of duplicate examples and 'negative' are the results of explicitly created negative examples. The results for our test case are as expected, V1[0]V2[0] match producing '1' while our other 'positive' cases (in green) don't match well, as was arranged. The V2[2] was set to match V1[3] producing a poor match at `score[2,2]` and an undesired 'negative' case of a '1' shown in grey. \n",
    "\n",
    "With the similarity matrix (`score`) we can begin to implement the loss equations. First, we can extract $$cos(A,P)$$ by utilizing `fastnp.diagonal`. The goal is to grab all the green entries in the diagram above. This is `positive` in the code.\n",
    "\n",
    "Next, we will create the *closest_negative*. This is the nonduplicate entry in V2 that is closest (has largest cosine similarity) to an entry in V1. Each row, n, of `score` represents all comparisons of the results of Q1[n] vs Q2[x] within a batch. A specific example in our testcase is row `score[2,:]`. It has the cosine similarity of V1[2] and V2[x]. The *closest_negative*, as was arranged, is V2[2] which has a score of 1. This is the maximum value of the 'negative' entries (blue entries in the diagram).\n",
    "\n",
    "To implement this, we need to pick the maximum entry on a row of `score`, ignoring the 'positive'/green entries. To avoid selecting the 'positive'/green entries, we can make them larger negative numbers. Multiply `fastnp.eye(batch_size)` with 2.0 and subtract it out of `scores`. The result is `negative_without_positive`. Now we can use `fastnp.max`, row by row (axis=1), to select the maximum which is `closest_negative`.\n",
    "\n",
    "Next, we'll create *mean_negative*. As the name suggests, this is the mean of all the 'negative'/blue values in `score` on a row by row basis. We can use `fastnp.eye(batch_size)` and a constant, this time to create a mask with zeros on the diagonal. Element-wise multiply this with `score` to get just the 'negative values. This is `negative_zero_on_duplicate` in the code. Compute the mean by using `fastnp.sum` on `negative_zero_on_duplicate` for `axis=1` and divide it by `(batch_size - 1)` . This is `mean_negative`.\n",
    "\n",
    "Now, we can compute loss using the two equations above and `fastnp.maximum`. This will form `triplet_loss1` and `triplet_loss2`. \n",
    "\n",
    "`triple_loss` is the `fastnp.sum` of the sum of the two individual losses.\n",
    "\n",
    "Once you have this code matching the expected results, you can clip out the section between `### START CODE HERE` and `### END CODE HERE` it out and insert it into TripletLoss below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# UNQ_C3 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)\n",
    "# GRADED FUNCTION: TripletLossFn\n",
    "def TripletLossFn(v1, v2, margin=0.25):\n",
    "    \"\"\"Custom Loss function.\n",
    "\n",
    "    Args:\n",
    "        v1 (numpy.ndarray): Array with dimension (batch_size, model_dimension) associated to Q1.\n",
    "        v2 (numpy.ndarray): Array with dimension (batch_size, model_dimension) associated to Q2.\n",
    "        margin (float, optional): Desired margin. Defaults to 0.25.\n",
    "\n",
    "    Returns:\n",
    "        jax.interpreters.xla.DeviceArray: Triplet Loss.\n",
    "    \"\"\"\n",
    "    ### START CODE HERE (Replace instances of 'None' with your code) ###\n",
    "    # use fastnp to take the dot product of the two batches (don't forget to transpose the second argument)\n",
    "    scores = None # pairwise cosine sim    \n",
    "    # calculate new batch size\n",
    "    batch_size = len(scores)\n",
    "    # use fastnp to grab all postive `diagonal` entries in `scores`\n",
    "    positive = None  # the positive ones (duplicates)\n",
    "    # subtract `fastnp.eye(batch_size)` out of 1.0 and do element-wise multiplication with `scores`\n",
    "    negative_zero_on_duplicate = None\n",
    "    # use `fastnp.sum` on `negative_zero_on_duplicate` for `axis=1` and divide it by `(batch_size - 1)`\n",
    "    mean_negative = None\n",
    "    # create a composition of two masks: \n",
    "    # the first mask to extract the diagonal elements, \n",
    "    # the second mask to extract elements in the negative_zero_on_duplicate matrix that are larger than the elements in the diagonal \n",
    "    mask_exclude_positives = (None)|(None)\n",
    "    # multiply `mask_exclude_positives` with 2.0 and subtract it out of `negative_zero_on_duplicate`\n",
    "    negative_without_positive = None\n",
    "    # take the row by row `max` of `negative_without_positive`. \n",
    "    # Hint: negative_without_positive.max(axis = [?])  \n",
    "    closest_negative = None\n",
    "    # compute `fastnp.maximum` among 0.0 and `A`\n",
    "    # where A = subtract `positive` from `margin` and add `closest_negative`\n",
    "    # IMPORTANT: DO NOT create an extra variable 'A'\n",
    "    triplet_loss1 = None\n",
    "    # compute `fastnp.maximum` among 0.0 and `B`\n",
    "    # where B = ubtract `positive` from `margin` and add `mean_negative`\n",
    "    # IMPORTANT: DO NOT create an extra variable 'B'\n",
    "    triplet_loss2 = None\n",
    "    # add the two losses together and take the `fastnp.sum` of it    \n",
    "    triplet_loss = None    \n",
    "    ### END CODE HERE ###\n",
    "    \n",
    "    return triplet_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v1 = np.array([[ 0.26726124,  0.53452248,  0.80178373],[-0.5178918 , -0.57543534, -0.63297887]])\n",
    "v2 = np.array([[0.26726124, 0.53452248, 0.80178373],[0.5178918 , 0.57543534, 0.63297887]])\n",
    "print(\"Triplet Loss:\", TripletLossFn(v1,v2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```CPP\n",
    "Triplet Loss: ~ 0.70\n",
    "```   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test your function\n",
    "w4_unittest.test_TripletLossFn(TripletLossFn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "r974ozuHYAom"
   },
   "source": [
    "To make a layer out of a function with no trainable variables, use `tl.Fn`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "def TripletLoss(margin=0.25):\n",
    "    triplet_loss_fn = partial(TripletLossFn, margin=margin)\n",
    "    return tl.Fn('TripletLoss', triplet_loss_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lsvjaCQ6wt02"
   },
   "source": [
    "<a name='3'></a>\n",
    "## 3 - Training\n",
    "\n",
    "Now you are going to train your model. As usual, you have to define the cost function and the optimizer. You also have to feed in the built model. Before, going into the training, we will use a special data set up. We will define the inputs using the data generator we built above. The lambda function acts as a seed to remember the last batch that was given. Run the cell below to get the question pairs inputs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "iPk7gh-nzCBg",
    "outputId": "a2e8525d-f89a-4d9d-c0d6-bd7406f0246a"
   },
   "outputs": [],
   "source": [
    "batch_size = 256\n",
    "train_generator = data_generator(train_Q1, train_Q2, batch_size, vocab['<PAD>'])\n",
    "val_generator = data_generator(val_Q1, val_Q2, batch_size, vocab['<PAD>'])\n",
    "print('train_Q1.shape ', train_Q1.shape)\n",
    "print('val_Q1.shape   ', val_Q1.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IgFMfH5awt07"
   },
   "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='ex-4'></a>\n",
    "### Exercise 4 - train_model\n",
    "\n",
    "**Instructions:** Implement the `train_model` below to train the neural network above. Here is a list of things you should do, as already shown in lecture 7: \n",
    "\n",
    "- Create `TrainTask` and `EvalTask`\n",
    "- Create the training loop `trax.supervised.training.Loop`\n",
    "- Pass in the following depending on the context (train_task or eval_task):\n",
    "    - `labeled_data=generator`\n",
    "    - `metrics=[TripletLoss()]`,\n",
    "    - `loss_layer=TripletLoss()`\n",
    "    - `optimizer=trax.optimizers.Adam` with learning rate of 0.01\n",
    "    - `lr_schedule=trax.lr.warmup_and_rsqrt_decay(400, 0.01)`,\n",
    "    - `output_dir=output_dir`\n",
    "\n",
    "\n",
    "You will be using your triplet loss function with Adam optimizer. Please read the [trax](https://trax-ml.readthedocs.io/en/latest/trax.optimizers.html?highlight=adam#trax.optimizers.adam.Adam) documentation to get a full understanding. \n",
    "\n",
    "This function should return a `training.Loop` object. To read more about this check the [docs](https://trax-ml.readthedocs.io/en/latest/trax.supervised.html?highlight=loop#trax.supervised.training.Loop)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_kbtfz4T_m7x"
   },
   "outputs": [],
   "source": [
    "# UNQ_C4 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)\n",
    "# GRADED FUNCTION: train_model\n",
    "def train_model(Siamese, TripletLoss\n",
    "                , train_generator, val_generator, output_dir='model/'):\n",
    "    \"\"\"Training the Siamese Model\n",
    "\n",
    "    Args:\n",
    "        Siamese (function): Function that returns the Siamese model.\n",
    "        TripletLoss (function): Function that defines the TripletLoss loss function.\n",
    "        train_generator (generator, optional): Training generator. Defaults to train_generator.\n",
    "        val_generator (generator, optional): Validation generator. Defaults to val_generator.\n",
    "        output_dir (str, optional): Path to save model to. Defaults to 'model/'.\n",
    "\n",
    "    Returns:\n",
    "        trax.supervised.training.Loop: Training loop for the model.\n",
    "    \"\"\"\n",
    "    output_dir = os.path.expanduser(output_dir)\n",
    "\n",
    "    ### START CODE HERE (Replace instances of 'None' with your code) ###\n",
    "\n",
    "    train_task = training.TrainTask( \n",
    "        labeled_data=None,      # Use generator (train)\n",
    "        loss_layer=None,        # Use triplet loss. Don't forget to instantiate this object\n",
    "        optimizer=None,         # Don't forget to add the learning rate parameter\n",
    "        lr_schedule=None # Use Trax multifactor schedule function\n",
    "    )\n",
    "\n",
    "    eval_task = training.EvalTask(\n",
    "        labeled_data=None,      # Use generator (val)\n",
    "        metrics=[None],         # Use triplet loss. Don't forget to instantiate this object\n",
    "    )\n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "\n",
    "    training_loop = training.Loop(Siamese(),\n",
    "                                  train_task,\n",
    "                                  eval_tasks=[eval_task],\n",
    "                                  output_dir=output_dir)\n",
    "\n",
    "    return training_loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 391
    },
    "colab_type": "code",
    "id": "-3KXjmBo_6Xa",
    "outputId": "9d57f731-1534-4218-e744-783359d5cd19"
   },
   "outputs": [],
   "source": [
    "train_steps = 5\n",
    "training_loop = train_model(Siamese, TripletLoss, train_generator, val_generator)\n",
    "training_loop.run(train_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test your function\n",
    "w4_unittest.test_train_model(train_model, Siamese, TripletLoss, data_generator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model was only trained for 5 steps due to the constraints of this environment. For the rest of the assignment you will be using a pretrained model but now you should understand how the training can be done using Trax."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "abKPe7d4wt1C"
   },
   "source": [
    "<a name='4'></a>\n",
    "## 4 - Evaluation  \n",
    "\n",
    "<a name='4-1'></a>\n",
    "### 4.1 - Evaluating your Siamese Network\n",
    "\n",
    "In this section you will learn how to evaluate a Siamese network. You will first start by loading a pretrained model and then you will use it to predict. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Siamese()\n",
    "model.init_from_file(file_name='model.pkl.gz', weights_only=True, input_signature=shapes.signature(next(train_generator)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QDi4MBiKpivF"
   },
   "source": [
    "<a name='4-2'></a>\n",
    "### 4.2 - Classify\n",
    "To determine the accuracy of the model, we will utilize the test set that was configured earlier. While in training we used only positive examples, the test data, Q1_test, Q2_test and y_test, is setup as pairs of questions, some of which are duplicates some are not. \n",
    "This routine will run all the test question pairs through the model, compute the cosine simlarity of each pair, threshold it and compare the result to  y_test - the correct response from the data set. The results are accumulated to produce an accuracy.\n",
    "\n",
    "\n",
    "<a name='ex-5'></a>\n",
    "### Exercise 5 - classify\n",
    "\n",
    "**Instructions**  \n",
    " - Loop through the incoming data in `batch_size` chunks\n",
    " - Use the data generator to load `q1`, `q2` a batch at a time. **Don't forget to set `shuffle=False`!**\n",
    " - copy a `batch_size` chunk of `y` into `y_test`\n",
    " - compute `v1`, `v2` using the model\n",
    " - for each element of the batch\n",
    "        - compute the cos similarity of each pair of entries, `v1[j]`,`v2[j]`\n",
    "        - determine if `d` > threshold\n",
    "        - increment accuracy if that result matches the expected results (`y_test[j]`)\n",
    " - compute the final accuracy and return\n",
    " \n",
    "Due to some limitations of this environment, running classify multiple times may result in the kernel failing. If that happens *Restart Kernal & clear output* and then run from the top. During development, consider using a smaller set of data to reduce the number of calls to model(). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "K-h6ZH507fUm"
   },
   "outputs": [],
   "source": [
    "# UNQ_C5 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)\n",
    "# GRADED FUNCTION: classify\n",
    "def classify(test_Q1, test_Q2, y, threshold, model, vocab, data_generator=data_generator, batch_size=64):\n",
    "    \"\"\"Function to test the accuracy of the model.\n",
    "\n",
    "    Args:\n",
    "        test_Q1 (numpy.ndarray): Array of Q1 questions.\n",
    "        test_Q2 (numpy.ndarray): Array of Q2 questions.\n",
    "        y (numpy.ndarray): Array of actual target.\n",
    "        threshold (float): Desired threshold.\n",
    "        model (trax.layers.combinators.Parallel): The Siamese model.\n",
    "        vocab (collections.defaultdict): The vocabulary used.\n",
    "        data_generator (function): Data generator function. Defaults to data_generator.\n",
    "        batch_size (int, optional): Size of the batches. Defaults to 64.\n",
    "\n",
    "    Returns:\n",
    "        float: Accuracy of the model.\n",
    "    \"\"\"    \n",
    "    \n",
    "    \n",
    "    accuracy = 0\n",
    "    ### START CODE HERE (Replace instances of 'None' with your code) ###\n",
    "    for i in range(0, len(test_Q1), batch_size):\n",
    "        # Call the data generator (built in Ex 01) with shuffle= None\n",
    "        # use batch size chuncks of questions as Q1 & Q2 arguments of the data generator. e.g x[i:i + batch_size]\n",
    "        # Hint: use `vocab['<PAD>']` for the `pad` argument of the data generator\n",
    "        q1, q2 = None\n",
    "        # use batch size chuncks of actual output targets (same syntax as example above)\n",
    "        y_test = None\n",
    "        # Call the model    \n",
    "        v1, v2 = None\n",
    "\n",
    "        for j in range(batch_size):\n",
    "            # take dot product to compute cos similarity of each pair of entries, v1[j], v2[j]\n",
    "            # don't forget to transpose the second argument\n",
    "            d = None\n",
    "            # is d greater than the threshold?\n",
    "            res = None\n",
    "            # increment accurancy if y_test is equal `res`\n",
    "            accuracy += None\n",
    "    # compute accuracy using accuracy and total length of test questions\n",
    "    accuracy = None\n",
    "    ### END CODE HERE ###\n",
    "    \n",
    "    return accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "yeQjHxkfpivH",
    "outputId": "103b8449-896f-403d-f011-583df70afdae"
   },
   "outputs": [],
   "source": [
    "# this takes around 1 minute\n",
    "accuracy = classify(Q1_test,Q2_test, y_test, 0.7, model, vocab, batch_size = 512) \n",
    "print(\"Accuracy\", accuracy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "CsokYZwhpivJ"
   },
   "source": [
    "**Expected Result**  \n",
    "Accuracy ~0.74"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test your function\n",
    "w4_unittest.test_classify(classify, vocab, data_generator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4-STC44Ywt1I"
   },
   "source": [
    "<a name='5'></a>\n",
    "## 5 -Testing with your Own Questions\n",
    "\n",
    "In this section you will test the model with your own questions. You will write a function `predict` which takes two questions as input and returns $1$ or $0$ depending on whether the question pair is a duplicate or not.   \n",
    "\n",
    "But first, we build a reverse vocabulary that allows to map encoded questions back to words: "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "21h3Y0FNpivK"
   },
   "source": [
    "Write a function `predict`that takes in two questions, the model, and the vocabulary and returns whether the questions are duplicates ($1$) or not duplicates ($0$) given a similarity threshold. \n",
    "\n",
    "<a name='ex-6'></a>\n",
    "### Exercise 6 - predict\n",
    "\n",
    "\n",
    "**Instructions:** \n",
    "- Tokenize your question using `nltk.word_tokenize` \n",
    "- Create Q1,Q2 by encoding your questions as a list of numbers using vocab\n",
    "- pad Q1,Q2 with next(data_generator([Q1], [Q2],1,vocab['<PAD>']))\n",
    "- use model() to create v1, v2\n",
    "- compute the cosine similarity (dot product) of v1, v2\n",
    "- compute res by comparing d to the threshold\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# UNQ_C6 (UNIQUE CELL IDENTIFIER, DO NOT EDIT)\n",
    "# GRADED FUNCTION: predict\n",
    "def predict(question1, question2, threshold, model, vocab, data_generator=data_generator, verbose=False):\n",
    "    \"\"\"Function for predicting if two questions are duplicates.\n",
    "\n",
    "    Args:\n",
    "        question1 (str): First question.\n",
    "        question2 (str): Second question.\n",
    "        threshold (float): Desired threshold.\n",
    "        model (trax.layers.combinators.Parallel): The Siamese model.\n",
    "        vocab (collections.defaultdict): The vocabulary used.\n",
    "        data_generator (function): Data generator function. Defaults to data_generator.\n",
    "        verbose (bool, optional): If the results should be printed out. Defaults to False.\n",
    "\n",
    "    Returns:\n",
    "        bool: True if the questions are duplicates, False otherwise.\n",
    "    \"\"\"\n",
    "    ### START CODE HERE (Replace instances of 'None' with your code) ###\n",
    "    # use `nltk` word tokenize function to tokenize\n",
    "    q1 = None  # tokenize\n",
    "    q2 = None  # tokenize\n",
    "    Q1, Q2 = [], []\n",
    "    for word in q1:  # encode q1\n",
    "        # append the 'word' index in `vocab`\n",
    "        Q1.append(None)\n",
    "    for word in q2:  # encode q2\n",
    "        # append the 'word' index in `vocab`\n",
    "        Q2.append(None)\n",
    "        \n",
    "    # Call the data generator (built in Ex 01) using next()\n",
    "    # pass [Q1] & [Q2] as Q1 & Q2 arguments of the data generator. Set batch size as 1\n",
    "    # Hint: use `vocab['<PAD>']` for the `pad` argument of the data generator\n",
    "    Q1, Q2 = None\n",
    "    # Call the model\n",
    "    v1, v2 = None\n",
    "    # take dot product to compute cos similarity of each pair of entries, v1, v2\n",
    "    # don't forget to transpose the second argument\n",
    "    d = None\n",
    "    # is d greater than the threshold?\n",
    "    res = None\n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "    \n",
    "    if(verbose):\n",
    "        print(\"Q1  = \", Q1, \"\\nQ2  = \", Q2)\n",
    "        print(\"d   = \", d)\n",
    "        print(\"res = \", res)\n",
    "\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 102
    },
    "colab_type": "code",
    "id": "Raojyhw3z7HE",
    "outputId": "b0907aaf-63c0-448d-99b0-012359381a97"
   },
   "outputs": [],
   "source": [
    "# Feel free to try with your own questions\n",
    "question1 = \"When will I see you?\"\n",
    "question2 = \"When can I see you again?\"\n",
    "# 1/True means it is duplicated, 0/False otherwise\n",
    "predict(question1 , question2, 0.7, model, vocab, verbose = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7OEKCa_hpivP"
   },
   "source": [
    "##### Expected Output\n",
    "If input is:\n",
    "```CPP\n",
    "question1 = \"When will I see you?\"\n",
    "question2 = \"When can I see you again?\"\n",
    "```\n",
    "\n",
    "Output is (d may vary a bit):\n",
    "```CPP\n",
    "Q1  =  [[585  76   4  46  53  21   1   1]] \n",
    "Q2  =  [[ 585   33    4   46   53 7280   21    1]]\n",
    "d   =  0.8585515\n",
    "res =  True\n",
    "True\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 102
    },
    "colab_type": "code",
    "id": "DZccIQ_lpivQ",
    "outputId": "3ed0af7e-5d44-4eb3-cebe-d6f74abe3e41"
   },
   "outputs": [],
   "source": [
    "# Feel free to try with your own questions\n",
    "question1 = \"Do they enjoy eating the dessert?\"\n",
    "question2 = \"Do they like hiking in the desert?\"\n",
    "# 1/True means it is duplicated, 0/False otherwise\n",
    "predict(question1 , question2, 0.7, model, vocab, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lWrt-yCMpivS"
   },
   "source": [
    "##### Expected output\n",
    "\n",
    "If input is:\n",
    "```CPP\n",
    "question1 = \"Do they enjoy eating the dessert?\"\n",
    "question2 = \"Do they like hiking in the desert?\"\n",
    "```\n",
    "\n",
    "Output  (d may vary a bit):\n",
    "\n",
    "```CPP\n",
    "Q1  =  [[  443  1145  3159  1169    78 29017    21     1]] \n",
    "Q2  =  [[  443  1145    60 15302    28    78  7431    21]]\n",
    "d   =  0.53644466\n",
    "res =  False\n",
    "False\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NAfV3l5Zwt1L"
   },
   "source": [
    "You can see that the Siamese network is capable of catching complicated structures. Concretely it can identify question duplicates although the questions do not have many words in common. \n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test your function\n",
    "w4_unittest.test_predict(predict, vocab, data_generator)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FsE8tdTLwt1M"
   },
   "source": [
    "<a name='6'></a>\n",
    "###  <span style=\"color:blue\"> On Siamese Networks </span>\n",
    "\n",
    "Siamese networks are important and useful. Many times there are several questions that are already asked in quora, or other platforms and you can use Siamese networks to avoid question duplicates. \n",
    "\n",
    "Congratulations, you have now built a powerful system that can recognize question duplicates. In the next chapter we will use transformers for machine translation, summarization, question answering, and chatbots. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "encoding": "# -*- coding: utf-8 -*-"
  },
  "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "56d44d6a8424451b5ce45d1ae0b0b7865dc60710e7f74571dd51dd80d7829ee9"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}