notebooks/Trainer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "importing Jupyter notebook from MoveData.ipynb\n",
      "importing Jupyter notebook from EncoderDecoder.ipynb\n",
      "importing Jupyter notebook from Elements.ipynb\n"
     ]
    }
   ],
   "source": [
    "import time, sys\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Variable\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import import_ipynb\n",
    "from MoveData import Options, json2datatools, num_batches, nopeak_mask, create_masks\n",
    "from EncoderDecoder import Encoder, Decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are not using this notebook to learn, change the below variable `teaching` to `False` so that other notebooks can import the functions defined in this notebook without running all the examples, if you are here to learn and interact with the notebook, change it to `True`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "teaching = True #False # True "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Big Picture\n",
    "\n",
    "I'm glad you decided to dive deeper. If we want to understand where deep learning is going, what new frontiers the experts are \"eyeing\", how to modify deep learning pipelines for our own goals, or to create new progress in AI, we need strive for a fundamental understanding.  \n",
    "\n",
    "There is soo much information out there on the internet about the basic feed forward neural network. If you are a beginner to neural networks I do suggest [reading this](https://colah.github.io/posts/2014-03-NN-Manifolds-Topology/) or [watching this](https://youtu.be/kNPGXgzxoHw) as a starter into the world of WHY / HOW neural networks work, in a mathematical and geometric sense.  \n",
    "\n",
    "I will not try to reinvent the wheel, instead I want to do what I feel is not done enough of, namely bringing these concepts into a fun, relatable, tangible robot while showing the actual code. This section on Training will be by far the most reading heavy, if you can get through this one, the rest will be mostly toy coding examples and diagrams. \n",
    "\n",
    "In Deep Reinforcement Learning, Neural Machine Translation, Convolutional Neural Networks, training chloe to talk, etc, and all Deep Learning, we are adjusting the weights / parameters / synapses of our neural network to either maximize some reward or minimize some loss function. Heres a seed I want to plant in your head that we will elaborate on later: In the diagram below, what determines the colored landscape, is our training data, that comes out of `data_iter` that you used in the introduction section START_HERE.ipynb\n",
    "\n",
    "<img src=\"../saved/images/LossTopology.png\">\n",
    "\n",
    "All those red dots you see are the initial starting points to the weights / parameters / synapses of our neural network. Usually the weights are randomly initialized as small numbers close to 0.0, there are tricks for sampling those weights from an optimal distribution, but thats for another lesson. As we get into the details, I think it is useful to keep this big picture in mind. \n",
    "\n",
    "OK lets get started, first, whether you are running this on a GPU ready environment or not, the below `if else` statement will set `device` appropriately to leverage the GPU or CPU. We are working on such a small dataset that for demonstration purposes, a CPU will do just as well and GPU is not usefull unless you have augmented the data to be significantly larger."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using CPU for evaluation and training\n"
     ]
    }
   ],
   "source": [
    "if teaching:\n",
    "    if torch.cuda.is_available():\n",
    "        print(\"you have\", torch.cuda.device_count(), \"GPUs\")\n",
    "        device = torch.device(\"cuda:0\")\n",
    "    else:\n",
    "        print('using CPU for evaluation and training')\n",
    "        device = torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing the Data\n",
    "in START_HERE.ipynb you have already seen the `Options` class and `json2datatools` function. As a reminder, `opt` is just a way to keep and pass all your preferences in a single input, rather than type out every hyperparameter again and again. we will need the input and output vocabulary `infield, outfield` for our demonstration of how a sequence of words is represented by the transformer and we will use `data_iter` to show how data flows from the dataset though our transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input vocab size 38 output vocab size 42\n"
     ]
    }
   ],
   "source": [
    "if teaching:\n",
    "    \n",
    "    opt = Options(batchsize=2, device=torch.device(\"cpu\"), epochs=25, \n",
    "                  lr=0.01, beam_width=3, max_len = 25, save_path = '../saved/weights/transformer_example_weights')\n",
    "\n",
    "    data_iter, infield, outfield, opt = json2datatools(path='../saved/examplepairs.json', opt=opt)\n",
    "    \n",
    "    print('input vocab size', len(infield.vocab), 'output vocab size', len(outfield.vocab))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Transformer and Optimizer \n",
    "Lets take a look inside the Transformer, like older sequence to sequence models, Transformers also encode the sentence into vector representations and pass those representations along to the decoder to generate the response/reply/output/translation/etc. The Encoder and Decoder have subcomponents which is discussed in Elements, Encoder and Decoder. For now, just define the Transformer class, instantiate a model as chloe and define the optimizer and scheduler by running the cell below "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self, in_vocab_size, out_vocab_size, emb_dim, n_layers, heads, dropout):\n",
    "        super().__init__()\n",
    "        self.encoder = Encoder(in_vocab_size, emb_dim, n_layers, heads, dropout)\n",
    "        self.decoder = Decoder(out_vocab_size, emb_dim, n_layers, heads, dropout)\n",
    "        self.out = nn.Linear(emb_dim, out_vocab_size)\n",
    "    def forward(self, src_seq, src_mask, trg_seq,  trg_mask):\n",
    "        e_output = self.encoder(src_seq, src_mask)\n",
    "        d_output = self.decoder(trg_seq, trg_mask, e_output, src_mask)\n",
    "        output = self.out(d_output)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the scheduler in 'min' mode, learning rate will be reduced by a factor of 0.5 when the quantity monitored (epoch_loss) has stopped decreasing. In the training loop, you will see this quantity being monitored in `scheduler.step(epoch_loss)` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Transformer(\n",
       "  (encoder): Encoder(\n",
       "    (embed): Embedder(\n",
       "      (embed): Embedding(38, 32)\n",
       "    )\n",
       "    (pe): PositionalEncoder(\n",
       "      (dropout): Dropout(p=0.01, inplace=False)\n",
       "    )\n",
       "    (layers): ModuleList(\n",
       "      (0): EncoderLayer(\n",
       "        (norm_1): Norm()\n",
       "        (dropout_1): Dropout(p=0.01, inplace=False)\n",
       "        (attn): MultiHeadAttention(\n",
       "          (q_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (k_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (v_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (dropout): Dropout(p=0.01, inplace=False)\n",
       "          (out): Linear(in_features=32, out_features=32, bias=True)\n",
       "        )\n",
       "        (norm_2): Norm()\n",
       "        (ff): FeedForward(\n",
       "          (linear_1): Linear(in_features=32, out_features=2048, bias=True)\n",
       "          (dropout): Dropout(p=0.01, inplace=False)\n",
       "          (linear_2): Linear(in_features=2048, out_features=32, bias=True)\n",
       "        )\n",
       "        (dropout_2): Dropout(p=0.01, inplace=False)\n",
       "      )\n",
       "      (1): EncoderLayer(\n",
       "        (norm_1): Norm()\n",
       "        (dropout_1): Dropout(p=0.01, inplace=False)\n",
       "        (attn): MultiHeadAttention(\n",
       "          (q_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (k_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (v_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (dropout): Dropout(p=0.01, inplace=False)\n",
       "          (out): Linear(in_features=32, out_features=32, bias=True)\n",
       "        )\n",
       "        (norm_2): Norm()\n",
       "        (ff): FeedForward(\n",
       "          (linear_1): Linear(in_features=32, out_features=2048, bias=True)\n",
       "          (dropout): Dropout(p=0.01, inplace=False)\n",
       "          (linear_2): Linear(in_features=2048, out_features=32, bias=True)\n",
       "        )\n",
       "        (dropout_2): Dropout(p=0.01, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm): Norm()\n",
       "  )\n",
       "  (decoder): Decoder(\n",
       "    (embed): Embedder(\n",
       "      (embed): Embedding(42, 32)\n",
       "    )\n",
       "    (pe): PositionalEncoder(\n",
       "      (dropout): Dropout(p=0.01, inplace=False)\n",
       "    )\n",
       "    (layers): ModuleList(\n",
       "      (0): DecoderLayer(\n",
       "        (norm_1): Norm()\n",
       "        (norm_2): Norm()\n",
       "        (norm_3): Norm()\n",
       "        (dropout_1): Dropout(p=0.01, inplace=False)\n",
       "        (dropout_2): Dropout(p=0.01, inplace=False)\n",
       "        (dropout_3): Dropout(p=0.01, inplace=False)\n",
       "        (attn_1): MultiHeadAttention(\n",
       "          (q_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (k_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (v_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (dropout): Dropout(p=0.01, inplace=False)\n",
       "          (out): Linear(in_features=32, out_features=32, bias=True)\n",
       "        )\n",
       "        (attn_2): MultiHeadAttention(\n",
       "          (q_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (k_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (v_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (dropout): Dropout(p=0.01, inplace=False)\n",
       "          (out): Linear(in_features=32, out_features=32, bias=True)\n",
       "        )\n",
       "        (ff): FeedForward(\n",
       "          (linear_1): Linear(in_features=32, out_features=2048, bias=True)\n",
       "          (dropout): Dropout(p=0.01, inplace=False)\n",
       "          (linear_2): Linear(in_features=2048, out_features=32, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (1): DecoderLayer(\n",
       "        (norm_1): Norm()\n",
       "        (norm_2): Norm()\n",
       "        (norm_3): Norm()\n",
       "        (dropout_1): Dropout(p=0.01, inplace=False)\n",
       "        (dropout_2): Dropout(p=0.01, inplace=False)\n",
       "        (dropout_3): Dropout(p=0.01, inplace=False)\n",
       "        (attn_1): MultiHeadAttention(\n",
       "          (q_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (k_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (v_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (dropout): Dropout(p=0.01, inplace=False)\n",
       "          (out): Linear(in_features=32, out_features=32, bias=True)\n",
       "        )\n",
       "        (attn_2): MultiHeadAttention(\n",
       "          (q_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (k_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (v_linear): Linear(in_features=32, out_features=32, bias=True)\n",
       "          (dropout): Dropout(p=0.01, inplace=False)\n",
       "          (out): Linear(in_features=32, out_features=32, bias=True)\n",
       "        )\n",
       "        (ff): FeedForward(\n",
       "          (linear_1): Linear(in_features=32, out_features=2048, bias=True)\n",
       "          (dropout): Dropout(p=0.01, inplace=False)\n",
       "          (linear_2): Linear(in_features=2048, out_features=32, bias=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (norm): Norm()\n",
       "  )\n",
       "  (out): Linear(in_features=32, out_features=42, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if teaching:\n",
    "    emb_dim, n_layers, heads, dropout = 32, 2, 8, 0.01 \n",
    "    chloe = Transformer(len(infield.vocab), len(outfield.vocab), emb_dim, n_layers, heads, dropout)\n",
    "    chloe.load_state_dict(torch.load(opt.save_path))\n",
    "\n",
    "    if opt.device == torch.device(\"cuda:0\"):\n",
    "        chloe =  chloe.cuda()\n",
    "\n",
    "    optimizer = torch.optim.Adam(chloe.parameters(), lr=opt.lr, betas=(0.9, 0.98), eps=1e-9)\n",
    "    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.5, patience=5)\n",
    "    \n",
    "chloe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Training Batch\n",
    "The line `batch = next(iter(data_iter))` extracts one batch of data into the variable `batch`. The batch has two parts to it, the `listen` part that is the input to the model, and the `reply` part that is the preferred output response to the corresponding `listen` input. This is a flawed assumption for conversational modeling since there are many valid responses to anything one listens to or hears. But for now we will use this method to teach chloe to reply in some linguistically coherent manner. \n",
    "\n",
    "`.transpose(0,1)` is used to flip the orientation of the data contained in batch.listen\n",
    "\n",
    "the first dimension of `listen.shape` and `reply.shape` should match whatever you put into `opt = Options(batchsize=2` above, it is the number of samples in each batch. Each sample is one line in our pairs.json dataset. The next dimension is the sequence length. \n",
    "\n",
    "source input sequence shape **(batch_size, sequence_length)**\n",
    "\n",
    "Although it seems messy, I have printed out each piece of data from the batch and it's shape. Each is a 2-dimensional tensor, or matrix. Above the listen and reply tensors I have printed the entire vocabulary so you can resolve that the integers represent coherent sentences. \n",
    "\n",
    "Notice that the reply always starts with an integer that represents the \"Start Of Sentence\" or `<sos>` token, and ends with the \"End Of Sentence\"or `<eos>` token, so it's sequence length is the length of the reply sentence + 2. Since not all reply sentences are the same length, we use `<pad>` as empty tokens to fill up space and make the sentences the same length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "defaultdict(<bound method Vocab._default_unk_index of <torchtext.vocab.Vocab object at 0x15ec2ef10>>, {'<unk>': 0, '<pad>': 1, 'chloe': 2, 'you': 3, '?': 4, 'hi': 5, 'i': 6, 'are': 7, 'haha': 8, 'lol': 9, 'ok': 10, 'bye': 11, 'hello': 12, 'how': 13, 'later': 14, 'sure': 15, 'who': 16, '!': 17, 'a': 18, 'alive': 19, 'am': 20, 'any': 21, 'dont': 22, 'dunno': 23, 'go': 24, 'gotta': 25, 'ill': 26, 'joke': 27, 'know': 28, 'me': 29, 'more': 30, 'see': 31, 'talk': 32, 'tell': 33, 'to': 34, 'true': 35, 'ttyl': 36, 'yes': 37})\n",
      " ------------------------------------------------------ \n",
      "tensor([[10, 15],\n",
      "        [ 5,  2]]) torch.Size([2, 2])\n",
      " ------------------------------------------------------ \n",
      "defaultdict(<bound method Vocab._default_unk_index of <torchtext.vocab.Vocab object at 0x15ec2eb80>>, {'<unk>': 0, '<pad>': 1, '<sos>': 2, '<eos>': 3, '?': 4, 'you': 5, 'i': 6, 'joke': 7, '!': 8, ',': 9, 'a': 10, 'can': 11, 'cats': 12, 'do': 13, 'french': 14, 'hi': 15, 'how': 16, 'say': 17, 'tell': 18, 'thank': 19, 'thanks': 20, 'at': 21, 'bye': 22, 'for': 23, 'laughing': 24, 'my': 25, 'ttyl': 26, 'am': 27, 'beaucoup': 28, 'chloe': 29, 'dont': 30, 'either': 31, 'know': 32, 'meowci': 33, 'alive': 34, 'are': 35, 'have': 36, 'me': 37, 'teach': 38, 'to': 39, 'viruses': 40, 'will': 41})\n",
      " ------------------------------------------------------ \n",
      "tensor([[ 2, 16, 13, 14, 12, 17, 19,  5,  4,  3,  1,  1],\n",
      "        [ 2, 15,  8,  9, 11,  6, 18,  5, 10,  7,  4,  3]]) torch.Size([2, 12])\n"
     ]
    }
   ],
   "source": [
    "if teaching:\n",
    "    batch = next(iter(data_iter))\n",
    "    listen = batch.listen.transpose(0,1)\n",
    "    reply = batch.reply.transpose(0,1)\n",
    "    reply_input = reply[:, :-1]\n",
    "\n",
    "    print(infield.vocab.stoi)\n",
    "    print(\" ------------------------------------------------------ \")\n",
    "    print(listen, listen.shape)\n",
    "    print(\" ------------------------------------------------------ \")\n",
    "    print(outfield.vocab.stoi)\n",
    "    print(\" ------------------------------------------------------ \")\n",
    "    print(reply, reply.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Padding and Masks\n",
    "Tensors have to be single values (scalars), lines (vectors), rectangles (matrices) or boxes. This means that if I want my batch to have sentence of unequal length, the shorter sentences have to be padded with the `<pad>` token so that every sentence is the same length, the example input batch here has been padded into a 3x3 tensor \n",
    "\n",
    "`i`      `love`   `you`\n",
    "\n",
    "`hello`  `<pad>` `<pad>`\n",
    "\n",
    "`see`  `ya` `<pad>`\n",
    "\n",
    "\n",
    "However, the padding does not carry useful information, so it must be masked away. Masks are tensors that are the same shape as the tensors they are a mask for, here we show a mask for the listen tensor and for the reply tensor. Think of a real mask, usually there will be two holes poked out for the eyes and one hole for the mouth, but otherwise the mask is meant to cover the rest of my face. The mask for the input listen tensor covers the padded elements. Seems tedious, but this is an artifact of the way we need to prepare data to train efficiently. The mask tensors printed below have a `True` or `1` in the positions that are open for use, and a `False` or `0` in the positions that are meant to be covered or not used. \n",
    "\n",
    "The shape of the listen mask is **(batch_size, 1, input_sequence_length)**, why is there a 1 in there? We will explain later. \n",
    "\n",
    "The Mask for the output sequence, or reply tensor, has an additional meaning. Transformers use something called \"Attention\", which we will talk about in the Elements.ipynb, Encoder.ipynb and Decoder.ipynb sections. At a high level, attention is a good word to describe what is occuring computationally. Suppose you are replying to this phrase:\n",
    "\n",
    "\"action potentials flow from dendrite to soma to where ?\"\n",
    "\n",
    "you write:\n",
    "\n",
    "\"they from from soma to axon\"\n",
    "\n",
    "As you, the decoder, intend to write the words \"they flow from\", you are paying attention to the words \"action potentials flow from\" from the encoder, this is called encoder attention. In addition you see that instead of flow I wrote \"from\" twice becaue I didnt pay enough attention to what I had already written, This is called self attention. \n",
    "\n",
    "As the decoder writes each word in responce, one word at a time, the word it just produced and it's previous words are fed back into the decoder to get the next word so it can pay attention to what has already been written. \n",
    "\n",
    "During training, we judge it based on whether it produced the right word or not, if it does not, we insert the right word at each position and hope that the next position is better predicted. This is called teacher forcing. during training, it can only pay attention to words in previous earlier positions, thus, even though we have the entire correct sequence given to us with each training batch, we will hide the words in later positions through the mask to simulate real self attention (you cant pay self attention to words you havent said yet). In addition to suppressing/covering the pad tokens, the reply mask also covers sequence positions in the future. \n",
    "\n",
    "This is why the shape of the reply mask is **(batch_size, output_sequence_length, output_sequence_length)**, having 2 dimensions named the same might be confusing, so i will rewrite this as \n",
    "\n",
    "**reply mask (batch_index, time_step, output_mask_position)**\n",
    "\n",
    "In the print statement below, I have shown you the mask `reply_mask[0,0,:]` indexed as `[0,0` for the first sample and first time step, `:` means all , it prints `True, False, . . . False` because at the first time step you can pay attention to the `<sos>` token only, not any future words in the \"correct reply\", also called the target sequence. At the next time step you can pay attention to both `<sos>` and also the first correct token, thus the mask will have the first two positions as True and the rest False `True, True, False . . . False`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[True]],\n",
      "\n",
      "        [[True]]]) torch.Size([2, 1, 1])\n",
      " ------------------------------------------------------ \n",
      "torch.Size([2, 7, 7])\n",
      " ------------------------------------------------------ \n",
      "first time step tensor([ True, False, False, False, False, False, False])\n",
      " ------------------------------------------------------ \n",
      "second time step tensor([ True,  True, False, False, False, False, False])\n"
     ]
    }
   ],
   "source": [
    "if teaching:\n",
    "    listen_mask, reply_mask = create_masks(listen, reply_input, opt)\n",
    "    print(listen_mask, listen_mask.shape)\n",
    "    print(\" ------------------------------------------------------ \")\n",
    "    print( reply_mask.shape)\n",
    "    print(\" ------------------------------------------------------ \")\n",
    "    print(\"first time step\", reply_mask[0,0,:])\n",
    "    print(\" ------------------------------------------------------ \")\n",
    "    print(\"second time step\", reply_mask[0,1,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What the Model Output Tensor Means\n",
    "The shape of the model (chloe's) output is **(batch_size, output_sequence_length, output_vocab_size)**\n",
    "For every reply in batch_size, chloe says output_sequence_length number of words, for each output position, `len(outfield.vocab)` numbers are given. `len(outfield.vocab)` is the number of words that are in chloe's output vocabulary. Yes, chloe has a separate vocab size for what she can hear and what she can say. suppose output_vocab_size = 100, then chloe only knows to output 100 different tokens. I say token, because this includes not only words, but also `<eos>`, `?`, `ummm`, `,` and other discreet symbols. \n",
    "\n",
    "Each of these numbers in `chloes_reply[0,-1,:]` represents one token, when I ran this cell the response included \n",
    "\n",
    "`'<unk>': 0, '<pad>': 1, '<sos>': 2, '<eos>': 3, 'you': 4, '?': 5,`\n",
    "\n",
    "`tensor([-1.7353, -0.5510, -1.4945,  7.2842,  2.6917,` \n",
    "\n",
    "and this makes alot of sense because `[0,-1,:]` indexes to the first example and the last output token. It includes all `:` the vocabulary. The tensor has in position 0, `-1.7353`, in position 3 it has `7.2842`, this was the highest number in the vector, it means that of all the tokens in the vocabulary, the word chloe things should come next is the `<eos>` token, 3 is the index for the end of sentence token  `'<eos>': 3,`, meaning that she is done talking. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input shape torch.Size([2, 1])\n",
      " ------------------------------------------------------ \n",
      "input_vocab_size = 38 , output_vocab_size = 42\n",
      " ------------------------------------------------------ \n",
      "defaultdict(<bound method Vocab._default_unk_index of <torchtext.vocab.Vocab object at 0x15ec2eb80>>, {'<unk>': 0, '<pad>': 1, '<sos>': 2, '<eos>': 3, '?': 4, 'you': 5, 'i': 6, 'joke': 7, '!': 8, ',': 9, 'a': 10, 'can': 11, 'cats': 12, 'do': 13, 'french': 14, 'hi': 15, 'how': 16, 'say': 17, 'tell': 18, 'thank': 19, 'thanks': 20, 'at': 21, 'bye': 22, 'for': 23, 'laughing': 24, 'my': 25, 'ttyl': 26, 'am': 27, 'beaucoup': 28, 'chloe': 29, 'dont': 30, 'either': 31, 'know': 32, 'meowci': 33, 'alive': 34, 'are': 35, 'have': 36, 'me': 37, 'teach': 38, 'to': 39, 'viruses': 40, 'will': 41})\n",
      " ------------------------------------------------------ \n",
      "tensor([-1.2105, -1.0750, -1.9493,  1.5802, -1.2343,  2.6195,  2.1463, -5.8859,\n",
      "         0.7584, -4.6275, -3.2901, -3.7808,  0.1250,  5.7570, -2.0828, -5.9735,\n",
      "         1.8321, -0.4483, -3.0120,  1.5837,  1.5458, -1.5513, -0.9909,  2.9960,\n",
      "         0.8230,  0.5954, -2.0812, -0.5969,  1.9854, -2.6555,  2.4549, -1.8326,\n",
      "        -2.1731,  5.1114, -1.4931,  0.1341, -3.2389,  0.6031, -2.9125,  1.2774,\n",
      "        -0.6778, -0.4346], grad_fn=<SliceBackward>) torch.Size([2, 7, 42])\n"
     ]
    }
   ],
   "source": [
    "if teaching:\n",
    "    print('input shape', listen.shape)\n",
    "    print(\" ------------------------------------------------------ \")\n",
    "    chloes_reply = chloe(listen, listen_mask, reply_input, reply_mask)\n",
    "    print('input_vocab_size =', len(infield.vocab), ', output_vocab_size =', len(outfield.vocab))\n",
    "    print(\" ------------------------------------------------------ \")\n",
    "    print(outfield.vocab.stoi)\n",
    "    print(\" ------------------------------------------------------ \")\n",
    "    print(chloes_reply[0,-1,:], chloes_reply.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Supervised Learning\n",
    "\n",
    "supervised Learning is the kind of learning that started the AI renaissance of the 2010s, but it is widely agreed by experts that it is not the type of learning that will get us to human level AI, at least not by itself. \n",
    "\n",
    "suppose the correct reply to listen = `bye`  `chloe` is reply = `see`  `ya` `<eos>`\n",
    "\n",
    "the batch we get gives is listen = `bye`  `chloe` and reply = `<sos>` `see`  `ya` `<eos>`\n",
    "\n",
    "we input `bye`  `chloe` into chloe, give her `<sos>` `see`  `ya`, plus the masks to prevent peaking into the future, and we compare whatever she replies against `see`  `ya` `<eos>`. I have been referring to `see`  `ya` `<eos>` as the \"correct response\", but this is also called the \"target\" abbreviated trg. In addition, what I have been referring to as \"listen\" is also called the \"source\" abbreviated src. \n",
    "\n",
    "This kind of training, in which we repeatedly challenge chloe by giving her an input sentence (listen tensor), see what she reponds (`chloes_reply`), compare `chloes_reply` to the right answer `reply`, aka \"the target\", then modify chloe's weights and biases (her synaptic connections) in order to get her closer to matching her `chloes_reply` to `reply`, aka \"the target\", is called supervised learning. \n",
    "\n",
    "When I say \"compare\" `chloes_reply` to the right answer `reply`, what do i mean? I mean that we have to give her a grade, or a numerical score, that tells us how well she is doing, then once we know that number we can change chloe \n",
    "in such a way that increases or decreases the numerical score. In machine learning, we typically choose to tie improvement to lower scores rather than in school where we tie better performance to higher scores. One way to rationalize this to yourself is that in supervised learning, there is one correct answer, and if you get that correct answer, your score is zero, there is no extra credit or creativity points here, just more and more punishment the further you are away from correct answer, what a terrible class right? \n",
    "\n",
    "This numerical score is the Loss, and the algorithm used to \"compare\" `chloes_reply` to the right answer `reply` and generate that score is the loss function. The way that we change chloe to make her get a better Loss is using one of the gradient decent optimizers such as the Adam Optimizer we mentioned in the intro START HERE, we will go more into this later. If you dont know, gradient decent algorithms are fancier version of the same concept: if you weight of synaptic connection has strength `W`, your input is `x`, your algorithm is `Wx`, the target is `y` and your loss function is `Loss = (y - Wx)^2`, then the derivative of the loss with respect to W `dL/dW` tells you whether to increase or decrease W in order to increase or decrease Loss. The derivative is also known as gradient when `W` is not one number but several numbers. So changing chloe's `W` is how we get her to say what we want. We will explain where `W` is in chloe's neural network in the Elements section. \n",
    "\n",
    "## Loss function\n",
    "\n",
    "The Loss function outputs a scalar, calculated from the output of your neural network and the ground truth labels, that requires the neural network to do what you want it to do in order to be minimized. This means if you change your parameters in the opposite direction of the gradient with respect to the loss, the neural network gets better at classification.\n",
    "\n",
    "A common loss function is the Cross Entropy Loss function, I have written a simplified version for you here\n",
    "\n",
    "$$Cross Entropy(chloes reply, target) = [i == target] \\cdot - ln(chloes reply_i)$$\n",
    "\n",
    "Let me explain this with a toy example. Suppose chloe only has 2 tokens in her vocabulary, `'cat': 0, 'french': 1` and the correct answer to input source \"what is chloe?\" is the output target \"french cat\", the target in token indices is therefore `[1, 0]`. Supposed your chloe model has never been trained before, so her output is `[0.0, 0.0]` for the first token and `[0.0, 0.0]` for the second token, meaning, she has no idea what to say. Run the cell below to see what loss score chloe will get for this \"unsure\" response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.6931, grad_fn=<NllLossBackward>)\n"
     ]
    }
   ],
   "source": [
    "if teaching:\n",
    "    \n",
    "    source = np.asarray([[0.0, 0.0], [0.0, 0.0]]) \n",
    "    target = np.asarray([1,0])\n",
    "\n",
    "    source = Variable(torch.from_numpy(source).float(), requires_grad=True)\n",
    "    target = Variable(torch.from_numpy(target))\n",
    "\n",
    "    print(F.cross_entropy(source, target)) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output you should get is `tensor(0.6931, grad_fn=<NllLossBackward>)`. PyTorch took your input `[0.0, 0.0]` and put it through a softmax function, which essentially forces your vector of numbers to sum (add up to) to 1.0 while remaining in proportion with one another. So `[0.0, 0.0]` was changed to `[0.5, 0.5]`, similarly `[1.0, 1.0]` is changed to `[0.5, 0.5]`,  `[-1.0, -1.0]` is changed to `[0.5, 0.5]` and `[1.0, 2.0]` is changed to `[0.33, 0.77]`. \n",
    "\n",
    "in the simplified cross entropy equation above \n",
    "\n",
    "$$[𝑖==𝑡𝑎𝑟𝑔𝑒𝑡]$$ \n",
    "\n",
    "means to only pay attention to the target index. Since the taget was `[1,0]`, only the bolded postions in chloe's output [0.5 , **0.5**], [**0.5**, 0.0 ] are considered.\n",
    "\n",
    "now open https://www.wolframalpha.com/ and paste in -ln(0.5) to the calculator, the answer you should get is 0.6931471805599453.  \n",
    "\n",
    "In the cell below, you will see that 2 things have changed. Chloe has now learned to favor the word cat in the second position, however she favors the word cat even more so in the first position. We have also added the argument \n",
    "`ignore_index = 1` into our loss function and so this mistake of failing to use the token for `french` will not count against her. In the cell below, you will see that the softmax function has normalized the vector `[2.0, 1.0]` to `[0.7311, 0.2689]` . Follow along in the equation for cross entropy and you will see that the loss -ln(0.731) is equal to 0.3133418192323586"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.9933, 0.0067],\n",
      "        [0.7311, 0.2689]], grad_fn=<SoftmaxBackward>)\n",
      " ------------------------------------------------------ \n",
      "tensor(0.3133, grad_fn=<NllLossBackward>)\n"
     ]
    }
   ],
   "source": [
    "if teaching:\n",
    "    source = np.asarray([[4.0, -1.0], [2.0, 1.0]]) \n",
    "    target = np.asarray([1,0])\n",
    "\n",
    "    source = Variable(torch.from_numpy(source).float(), requires_grad=True)\n",
    "    print(F.softmax(source, dim=1))\n",
    "    print(\" ------------------------------------------------------ \")\n",
    "    target = Variable(torch.from_numpy(target))\n",
    "\n",
    "    print(F.cross_entropy(source, target, ignore_index = 1)) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Loop\n",
    "\n",
    "The function below is our training function, it set the model into training mode, it runs a prespecified number of epochs, each epoch feeding the model batches of source target pairs until the entire dataset is seen. \n",
    "\n",
    "The loss is calculated for each source target pair.\n",
    "\n",
    "`batch_loss = F.cross_entropy(preds.view(-1, preds.size(-1)), ys, ignore_index = options.trg_pad)`\n",
    "\n",
    "The gradients with respect to the model weights, aka parameters, are calculated with the backpropagation algorithm and used to update the weights.\n",
    "\n",
    "`batch_loss.backward()`\n",
    "\n",
    "The optimizer and scheduler need to be updated of our progress to make it's adjustments to the weights and update rules\n",
    "\n",
    "`optimizer.step()`\n",
    "`scheduler.step(epoch_loss)`\n",
    "\n",
    "Whenever we reach a new best (lowest) loss value, we save the model\n",
    "\n",
    "`if epoch_loss < best_loss:`\n",
    "\n",
    "`torch.save(model.state_dict(), options.save_path)`\n",
    "\n",
    "Like we said in the intro, Optimization is a very large field, [beginners video](https://youtu.be/vMh0zPT0tLI). Before going into neural networks and all the components, or modules, of our model, I think it helps to have a sense of how all of it is supposed to work together so you keep in mind the big picture, even while we dive into the details.\n",
    "\n",
    "in the training loop below you will notice something abit strange\n",
    "\n",
    "`trg_input = trg[:, :-1]`\n",
    "\n",
    "`ys = trg[:, 1:]`\n",
    "\n",
    "thats because in addition to providing the input for encoding and encoder attention, we also train the neural network to apply self attention, which is why we also provide the transformer with a target input that is shifted -1 from the target. If the target is  `see`  `ya` `<eos>`, the target input is `<sos>` `see`  `ya` so that we can train to neural network to recognize that if prompted with  `<sos>`, the appropriate initial token in  `see`, and if the neural network has already spoken  `see`, the next appropriate token is `ya`, and if the neural network has already spoken  `see`  `ya`, then the sentence is finished and the next appropriate token is `<eos>`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trainer(model, data_iterator, options, optimizer, scheduler):\n",
    "\n",
    "    if torch.cuda.is_available() and options.device == torch.device(\"cuda:0\"):\n",
    "        print(\"a GPU was detected, model will be trained on GPU\")\n",
    "        model = model.cuda()\n",
    "    else:\n",
    "        print(\"training on cpu\")\n",
    "\n",
    "    model.train()\n",
    "    start = time.time()\n",
    "    best_loss = 100\n",
    "    for epoch in range(options.epochs):\n",
    "        total_loss = 0\n",
    "        for i, batch in enumerate(data_iterator): \n",
    "            src = batch.listen.transpose(0,1) # get an input ready\n",
    "            trg = batch.reply.transpose(0,1)  # get a target ready\n",
    "            trg_input = trg[:, :-1] # trg_input is exactly 1 timestep behind target (ys)\n",
    "            src_mask, trg_mask = create_masks(src, trg_input, options) # get the mask ready \n",
    "            preds = model(src, src_mask, trg_input, trg_mask)\n",
    "            ys = trg[:, 1:].contiguous().view(-1) # target. \n",
    "            optimizer.zero_grad() # clean out the gradients\n",
    "            batch_loss = F.cross_entropy(preds.view(-1, preds.size(-1)), \n",
    "                                         ys, ignore_index = options.trg_pad)\n",
    "            batch_loss.backward() # calculate the gradient\n",
    "            optimizer.step() # update the weights\n",
    "            total_loss += batch_loss.item()\n",
    "\n",
    "        epoch_loss = total_loss/(num_batches(data_iterator)+1)\n",
    "        scheduler.step(epoch_loss)\n",
    "\n",
    "        if epoch_loss < best_loss:\n",
    "            best_loss = epoch_loss\n",
    "            torch.save(model.state_dict(), options.save_path)\n",
    "        print(\"%dm: epoch %d loss = %.3f\" %((time.time() - start)//60, epoch, epoch_loss))\n",
    "        total_loss = 0\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take a quick look at the `notebooks/MoveData.ipynb` just to examine the data processing functions incase they might be helpful in debugging and also to examine and understand the source and target masks. After that, the next lesson is about the building blocks of our neural network, open `notebooks/Elements.ipynb` for the next part in our adventure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}