Module12/12_seq2seq_attention.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Dataflowr](https://raw.githubusercontent.com/dataflowr/website/master/_assets/dataflowr_logo.png)](https://dataflowr.github.io/website/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attention for seq2seq\n",
    "\n",
    "This notebook has been adapted fom the pytorch tutorial [NLP FROM SCRATCH: TRANSLATION WITH A SEQUENCE TO SEQUENCE NETWORK AND ATTENTION](https://pytorch.org/tutorials/intermediate/seq2seq_translation_tutorial.html) by Sean Robertson and adapted by Marc Lelarge for the [deep learning course](https://dataflowr.github.io/website/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as ticker\n",
    "%matplotlib inline\n",
    "from __future__ import unicode_literals, print_function, division\n",
    "from io import open\n",
    "import unicodedata\n",
    "import string\n",
    "import re\n",
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch import optim\n",
    "import torch.nn.functional as F\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# where data is stored:\n",
    "# if on COLAB, you can comment the following lines \n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "ROOT_DIR = Path.home()\n",
    "data_path = os.path.join(ROOT_DIR,'data/')\n",
    "\n",
    "# and uncomment the following lines if on COLAB, you can download the data with:\n",
    "#!wget https://download.pytorch.org/tutorial/data.zip\n",
    "#!unzip data.zip\n",
    "#data_path = './'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def running_mean(x, N=100):\n",
    "    cumsum = np.cumsum(np.insert(x, 0, 0)) \n",
    "    return (cumsum[N:] - cumsum[:-N]) / float(N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# data preprocessing\n",
    "\n",
    "This code is directly taken from the Pytorch tutorial and creates the corpus of pairs of sentences in french and english as well as the french and englsih tokenizers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SOS_token = 0\n",
    "EOS_token = 1\n",
    "\n",
    "class Lang:\n",
    "    def __init__(self, name):\n",
    "        self.name = name\n",
    "        self.word2index = {}\n",
    "        self.word2count = {}\n",
    "        self.index2word = {0: \"SOS\", 1: \"EOS\"}\n",
    "        self.n_words = 2  # Count SOS and EOS\n",
    "\n",
    "    def addSentence(self, sentence):\n",
    "        for word in sentence.split(' '):\n",
    "            self.addWord(word)\n",
    "\n",
    "    def addWord(self, word):\n",
    "        if word not in self.word2index:\n",
    "            self.word2index[word] = self.n_words\n",
    "            self.word2count[word] = 1\n",
    "            self.index2word[self.n_words] = word\n",
    "            self.n_words += 1\n",
    "        else:\n",
    "            self.word2count[word] += 1\n",
    "\n",
    "# Turn a Unicode string to plain ASCII, thanks to\n",
    "# https://stackoverflow.com/a/518232/2809427\n",
    "def unicodeToAscii(s):\n",
    "    return ''.join(\n",
    "        c for c in unicodedata.normalize('NFD', s)\n",
    "        if unicodedata.category(c) != 'Mn'\n",
    "    )\n",
    "\n",
    "# Lowercase, trim, and remove non-letter characters\n",
    "\n",
    "def normalizeString(s):\n",
    "    s = unicodeToAscii(s.lower().strip())\n",
    "    s = re.sub(r\"([.!?])\", r\" \\1\", s)\n",
    "    s = re.sub(r\"[^a-zA-Z.!?]+\", r\" \", s)\n",
    "    return s\n",
    "\n",
    "def readLangs(lang1, lang2, reverse=False):\n",
    "    print(\"Reading lines...\")\n",
    "\n",
    "    # Read the file and split into lines\n",
    "    lines = open(data_path+'data/%s-%s.txt' % (lang1, lang2), encoding='utf-8').\\\n",
    "        read().strip().split('\\n')\n",
    "\n",
    "    # Split every line into pairs and normalize\n",
    "    pairs = [[normalizeString(s) for s in l.split('\\t')] for l in lines]\n",
    "\n",
    "    # Reverse pairs, make Lang instances\n",
    "    if reverse:\n",
    "        pairs = [list(reversed(p)) for p in pairs]\n",
    "        input_lang = Lang(lang2)\n",
    "        output_lang = Lang(lang1)\n",
    "    else:\n",
    "        input_lang = Lang(lang1)\n",
    "        output_lang = Lang(lang2)\n",
    "\n",
    "    return input_lang, output_lang, pairs\n",
    "\n",
    "MAX_LENGTH = 10\n",
    "\n",
    "eng_prefixes = (\n",
    "    \"i am \", \"i m \",\n",
    "    \"he is\", \"he s \",\n",
    "    \"she is\", \"she s \",\n",
    "    \"you are\", \"you re \",\n",
    "    \"we are\", \"we re \",\n",
    "    \"they are\", \"they re \"\n",
    ")\n",
    "\n",
    "def filterPair(p):\n",
    "    return len(p[0].split(' ')) < MAX_LENGTH and \\\n",
    "        len(p[1].split(' ')) < MAX_LENGTH and \\\n",
    "        p[1].startswith(eng_prefixes)\n",
    "\n",
    "def filterPair_test(p):\n",
    "    return len(p[0].split(' ')) > MAX_LENGTH and \\\n",
    "        len(p[1].split(' ')) > MAX_LENGTH and \\\n",
    "        p[1].startswith(eng_prefixes)\n",
    "\n",
    "def filterPairs(pairs):\n",
    "    return [pair for pair in pairs if filterPair(pair)]\n",
    "\n",
    "def filterPairs_test(pairs):\n",
    "    return [pair for pair in pairs if filterPair_test(pair)]\n",
    "\n",
    "def prepareData(lang1, lang2, reverse=False):\n",
    "    input_lang, output_lang, pairs = readLangs(lang1, lang2, reverse)\n",
    "    print(\"Read %s sentence pairs\" % len(pairs))\n",
    "    pairs = filterPairs(pairs)\n",
    "    print(\"Trimmed to %s sentence pairs\" % len(pairs))\n",
    "    print(\"Counting words...\")\n",
    "    for pair in pairs:\n",
    "        input_lang.addSentence(pair[0])\n",
    "        output_lang.addSentence(pair[1])\n",
    "    print(\"Counted words:\")\n",
    "    print(input_lang.name, input_lang.n_words)\n",
    "    print(output_lang.name, output_lang.n_words)\n",
    "    return input_lang, output_lang, pairs\n",
    "\n",
    "input_lang, output_lang, pairs = prepareData('eng', 'fra', True)\n",
    "print(random.choice(pairs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pairs_train, pairs_val = train_test_split(pairs, test_size=0.2, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def indexesFromSentence(lang, sentence):\n",
    "    return [lang.word2index[word] for word in sentence.split(' ')]\n",
    "\n",
    "def tensorFromSentence(lang, sentence):\n",
    "    indexes = indexesFromSentence(lang, sentence)\n",
    "    indexes.append(EOS_token)\n",
    "    return torch.tensor(indexes, dtype=torch.long, device=device).view(-1, 1)\n",
    "\n",
    "def tensorsFromPair(pair):\n",
    "    input_tensor = tensorFromSentence(input_lang, pair[0])\n",
    "    target_tensor = tensorFromSentence(output_lang, pair[1])\n",
    "    return (input_tensor, target_tensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_pairs = [tensorsFromPair(pv) for pv in pairs_val]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Seq2seq\n",
    "\n",
    "Here we follow the PyTorch tutorial and implement [Sequence to Sequence Learning with Neural Networks](https://arxiv.org/abs/1409.3215v3). The only modification in the code is that the encoder takes a full sentence (and no hidden state) in the forward pass and outputs all the correponding hidden states. As a result, there is no need to make a for loop for the encoder. Still, to get things simple, we do not deal with batches, if you want to deal with batches have a look at [Batches with sequences in Pytorch](https://dataflowr.github.io/website/modules/11c-batches-with-sequences/).\n",
    "\n",
    "We also train on a test set and compute the loss on a validation set (see the split done on the corpus above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderRNN(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size):\n",
    "        super(EncoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.embedding = nn.Embedding(input_size, hidden_size)\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size)\n",
    "\n",
    "    def forward(self, input):\n",
    "        embedded = self.embedding(input)\n",
    "        output, _ = self.gru(embedded, self.initHidden())\n",
    "        return output\n",
    "\n",
    "    def initHidden(self):\n",
    "        return torch.zeros(1, 1, self.hidden_size, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_iters = 4\n",
    "training_pairs = [tensorsFromPair(random.choice(pairs))\n",
    "                      for i in range(n_iters)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_pairs[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_pairs[0][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_size = 256\n",
    "encoder = EncoderRNN(input_lang.n_words, hidden_size).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "one_input = training_pairs[0][0]\n",
    "out = encoder(one_input)\n",
    "# here is the code used in the tutorial:\n",
    "#encoder_hidden = encoder.initHidden()\n",
    "#for c in one_input:\n",
    "#    out, encoder_hidden = encoder(c,encoder_hidden)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderRNN(nn.Module):\n",
    "    def __init__(self, hidden_size, output_size):\n",
    "        super(DecoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        self.embedding = nn.Embedding(output_size, hidden_size)\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size)\n",
    "        self.out = nn.Linear(hidden_size, output_size)\n",
    "        self.softmax = nn.LogSoftmax(dim=1)\n",
    "\n",
    "    def forward(self, input, hidden):\n",
    "        output = self.embedding(input).view(1, 1, -1)\n",
    "        output = F.relu(output)\n",
    "        output, hidden = self.gru(output, hidden)\n",
    "        output = self.softmax(self.out(output[0]))\n",
    "        return output, hidden\n",
    "\n",
    "    def initHidden(self):\n",
    "        return torch.zeros(1, 1, self.hidden_size, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder = DecoderRNN(hidden_size, output_lang.n_words).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder_input = torch.tensor([[SOS_token]], device=device)\n",
    "decoder_hidden = out[-1,:,:].view(1,1,hidden_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder_hidden.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output, hidden = decoder(decoder_input, decoder_hidden)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_onepair(input_tensor, target_tensor, encoder, decoder, encoder_optimizer, decoder_optimizer, criterion, teacher_forcing_ratio = 0.5):\n",
    "    encoder_hidden = encoder.initHidden()\n",
    "    \n",
    "    encoder = encoder.train()\n",
    "    decoder = decoder.train()\n",
    "    encoder_optimizer.zero_grad()\n",
    "    decoder_optimizer.zero_grad()\n",
    "    target_length = target_tensor.size(0)\n",
    "    loss = 0\n",
    "    encoder_outputs = encoder(input_tensor)\n",
    "    decoder_input = torch.tensor([[SOS_token]], device=device)\n",
    "    decoder_hidden = encoder_outputs[-1,:,:].view(1,1,hidden_size)\n",
    "    \n",
    "    use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False\n",
    "\n",
    "    if use_teacher_forcing:\n",
    "        # Teacher forcing: Feed the target as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden = decoder(\n",
    "                decoder_input, decoder_hidden)\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            decoder_input = target_tensor[di]  # Teacher forcing\n",
    "\n",
    "    else:\n",
    "        # Without teacher forcing: use its own predictions as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden = decoder(\n",
    "                decoder_input, decoder_hidden)\n",
    "            topv, topi = decoder_output.topk(1)\n",
    "            decoder_input = topi.squeeze().detach()  # detach from history as input\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            if decoder_input.item() == EOS_token:\n",
    "                break\n",
    "\n",
    "    loss.backward()\n",
    "    encoder_optimizer.step()\n",
    "    decoder_optimizer.step()\n",
    "\n",
    "    return loss.item() / target_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_onepair(input_tensor, target_tensor, encoder, decoder, criterion, teacher_forcing_ratio = 0.5):\n",
    "    encoder_hidden = encoder.initHidden()\n",
    "\n",
    "    encoder = encoder.eval()\n",
    "    decoder = decoder.eval()\n",
    "    target_length = target_tensor.size(0)\n",
    "    loss = 0\n",
    "    encoder_outputs = encoder(input_tensor)\n",
    "    decoder_input = torch.tensor([[SOS_token]], device=device)\n",
    "    decoder_hidden = encoder_outputs[-1,:,:].view(1,1,hidden_size)\n",
    "    \n",
    "    use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False\n",
    "\n",
    "    if use_teacher_forcing:\n",
    "        # Teacher forcing: Feed the target as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden = decoder(\n",
    "                decoder_input, decoder_hidden)\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            decoder_input = target_tensor[di]  # Teacher forcing\n",
    "\n",
    "    else:\n",
    "        # Without teacher forcing: use its own predictions as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden = decoder(\n",
    "                decoder_input, decoder_hidden)\n",
    "            topv, topi = decoder_output.topk(1)\n",
    "            decoder_input = topi.squeeze().detach()  # detach from history as input\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            if decoder_input.item() == EOS_token:\n",
    "                break\n",
    "\n",
    "    return loss.item() / target_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trainIters(encoder, decoder, n_iters, print_every=1000, learning_rate=0.01, teacher_forcing_ratio=0.9):\n",
    "    \n",
    "    plot_losses = []\n",
    "    plot_losses_val = []\n",
    "    print_loss_total = 0  # Reset every print_every\n",
    "\n",
    "    encoder_optimizer = optim.SGD(encoder.parameters(), lr=learning_rate)\n",
    "    decoder_optimizer = optim.SGD(decoder.parameters(), lr=learning_rate)\n",
    "    training_pairs = [tensorsFromPair(random.choice(pairs_train))\n",
    "                      for i in range(n_iters)]\n",
    "    criterion = nn.NLLLoss()\n",
    "\n",
    "    for iter in range(1, n_iters + 1):\n",
    "        training_pair = training_pairs[iter - 1]\n",
    "        input_tensor = training_pair[0]\n",
    "        target_tensor = training_pair[1]\n",
    "\n",
    "        loss = train_onepair(input_tensor, target_tensor, encoder,\n",
    "                     decoder, encoder_optimizer, decoder_optimizer, criterion, teacher_forcing_ratio)\n",
    "        print_loss_total += loss\n",
    "        plot_losses.append(loss)\n",
    "\n",
    "        if iter % print_every == 0:\n",
    "            loss_val = 0\n",
    "            for (input_tensor, target_tensor) in val_pairs:\n",
    "                loss = eval_onepair(input_tensor, target_tensor, encoder,\n",
    "                     decoder, criterion, teacher_forcing_ratio)\n",
    "                loss_val += loss\n",
    "            loss_val = loss_val/len(val_pairs)\n",
    "            plot_losses_val.append(loss_val)\n",
    "            print_loss_avg = print_loss_total / print_every\n",
    "            print_loss_total = 0\n",
    "            print('(%d %d%%) loss train %.4f and val %.4f' % (iter, iter / n_iters * 100, print_loss_avg, loss_val))\n",
    "    return plot_losses, plot_losses_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate=0.01\n",
    "encoder_optimizer = optim.SGD(encoder.parameters(), lr=learning_rate)\n",
    "decoder_optimizer = optim.SGD(decoder.parameters(), lr=learning_rate)\n",
    "criterion = nn.NLLLoss()\n",
    "train_onepair(training_pairs[0][0],training_pairs[0][1],encoder,decoder,encoder_optimizer,decoder_optimizer,criterion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_size = 128\n",
    "n_epochs = 50000\n",
    "encoder = EncoderRNN(input_lang.n_words, hidden_size).to(device)\n",
    "decoder = DecoderRNN(hidden_size, output_lang.n_words).to(device)\n",
    "plot_losses,plot_losses_val = trainIters(encoder,decoder,n_epochs,print_every=2500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(running_mean(plot_losses))\n",
    "plt.plot([2500*i for i in range(int(n_epochs/2500))], plot_losses_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(encoder, decoder, sentence, max_length=MAX_LENGTH):\n",
    "    with torch.no_grad():\n",
    "        input_tensor = tensorFromSentence(input_lang, sentence)\n",
    "        input_length = input_tensor.size()[0]\n",
    "        encoder_hidden = encoder.initHidden()\n",
    "\n",
    "        encoder_outputs = encoder(input_tensor)\n",
    "\n",
    "        decoder_input = torch.tensor([[SOS_token]], device=device)  # SOS\n",
    "\n",
    "        decoder_hidden = encoder_outputs[-1,:,:].view(1,1,hidden_size)\n",
    "\n",
    "        decoded_words = []\n",
    "\n",
    "        for di in range(max_length):\n",
    "            decoder_output, decoder_hidden = decoder(\n",
    "                decoder_input, decoder_hidden)\n",
    "            \n",
    "            topv, topi = decoder_output.data.topk(1)\n",
    "            if topi.item() == EOS_token:\n",
    "                decoded_words.append('<EOS>')\n",
    "                break\n",
    "            else:\n",
    "                decoded_words.append(output_lang.index2word[topi.item()])\n",
    "\n",
    "            decoder_input = topi.squeeze().detach()\n",
    "\n",
    "        return decoded_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluateRandomly(encoder, decoder, n=10):\n",
    "    for i in range(n):\n",
    "        pair = random.choice(pairs_val)\n",
    "        print('>', pair[0])\n",
    "        print('=', pair[1])\n",
    "        output_words = evaluate(encoder, decoder, pair[0])\n",
    "        output_sentence = ' '.join(output_words)\n",
    "        print('<', output_sentence)\n",
    "        print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluateRandomly(encoder, decoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attention mechanism\n",
    "\n",
    "Here we implement the attention mechanism from the paper [Neural Machine Translation by Jointly Learning to Align and Translate](https://arxiv.org/abs/1409.0473). The code below is siginficantly different from the original Pytorch Tutorial..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttnDecoderRNN(nn.Module):\n",
    "    def __init__(self, hidden_size, output_size):\n",
    "        super().__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.output_size = output_size\n",
    "        self.embedding = nn.Embedding(self.output_size, self.hidden_size, padding_idx=0)\n",
    "        self.attn_w = nn.Linear(2 * self.hidden_size, self.hidden_size)\n",
    "        self.attn_v = nn.Linear(self.hidden_size, 1)\n",
    "        self.gru = nn.GRU(self.hidden_size * 2, self.hidden_size)\n",
    "        self.out = nn.Linear(self.hidden_size * 2, self.output_size)\n",
    "\n",
    "    def forward(self, input, hidden, encoder_outputs):\n",
    "        # B = 1 batch size\n",
    "        # encoder_outputs (L,B,H)\n",
    "        seq_len, _, _ = encoder_outputs.shape\n",
    "        # hidden (1,B,H)\n",
    "        #\n",
    "        # your code here\n",
    "        #\n",
    "        return output, hidden, attn_weights\n",
    "\n",
    "    def initHidden(self):\n",
    "        return torch.zeros(1, 1, self.hidden_size, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_size = 256\n",
    "encoder = EncoderRNN(input_lang.n_words, hidden_size).to(device)\n",
    "one_input = training_pairs[0][0]\n",
    "out = encoder(one_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder = AttnDecoderRNN(hidden_size, output_lang.n_words).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder_input = torch.tensor([[SOS_token]], device=device)\n",
    "decoder_hidden = decoder.initHidden()\n",
    "encoder_outputs = out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder_input.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output, hidden, attn_weights = decoder(decoder_input, decoder_hidden, encoder_outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attn_weights.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.sum(attn_weights,dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_onepair(input_tensor, target_tensor, encoder, decoder, encoder_optimizer, decoder_optimizer, criterion, teacher_forcing_ratio = 0.5):\n",
    "    encoder_hidden = encoder.initHidden()\n",
    "    \n",
    "    encoder = encoder.train()\n",
    "    decoder = decoder.train()\n",
    "    encoder_optimizer.zero_grad()\n",
    "    decoder_optimizer.zero_grad()\n",
    "    target_length = target_tensor.size(0)\n",
    "    loss = 0\n",
    "    encoder_outputs = encoder(input_tensor)\n",
    "    decoder_input = torch.tensor([[SOS_token]], device=device)\n",
    "    decoder_hidden = decoder.initHidden()\n",
    "    \n",
    "    use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False\n",
    "\n",
    "    if use_teacher_forcing:\n",
    "        # Teacher forcing: Feed the target as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden, _ = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs)\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            decoder_input = target_tensor[di]  # Teacher forcing\n",
    "\n",
    "    else:\n",
    "        # Without teacher forcing: use its own predictions as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden, _ = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs)\n",
    "            topv, topi = decoder_output.topk(1)\n",
    "            decoder_input = topi.squeeze().detach()  # detach from history as input\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            if decoder_input.item() == EOS_token:\n",
    "                break\n",
    "\n",
    "    loss.backward()\n",
    "    encoder_optimizer.step()\n",
    "    decoder_optimizer.step()\n",
    "\n",
    "    return loss.item() / target_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_onepair(input_tensor, target_tensor, encoder, decoder, criterion, teacher_forcing_ratio = 0.5):\n",
    "    encoder_hidden = encoder.initHidden()\n",
    "\n",
    "    encoder = encoder.eval()\n",
    "    decoder = decoder.eval()\n",
    "    target_length = target_tensor.size(0)\n",
    "    loss = 0\n",
    "    encoder_outputs = encoder(input_tensor)\n",
    "    decoder_input = torch.tensor([[SOS_token]], device=device)\n",
    "    decoder_hidden = decoder.initHidden()\n",
    "    \n",
    "    use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False\n",
    "\n",
    "    if use_teacher_forcing:\n",
    "        # Teacher forcing: Feed the target as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden, _ = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs)\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            decoder_input = target_tensor[di]  # Teacher forcing\n",
    "\n",
    "    else:\n",
    "        # Without teacher forcing: use its own predictions as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden, _ = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs)\n",
    "            topv, topi = decoder_output.topk(1)\n",
    "            decoder_input = topi.squeeze().detach()  # detach from history as input\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            if decoder_input.item() == EOS_token:\n",
    "                break\n",
    "\n",
    "    return loss.item() / target_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate=0.01\n",
    "encoder_optimizer = optim.SGD(encoder.parameters(), lr=learning_rate)\n",
    "decoder_optimizer = optim.SGD(decoder.parameters(), lr=learning_rate)\n",
    "criterion = nn.NLLLoss()\n",
    "train_onepair(training_pairs[0][0],training_pairs[0][1],encoder,decoder,encoder_optimizer,decoder_optimizer,criterion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_size = 128\n",
    "n_epochs = 50000\n",
    "encoder = EncoderRNN(input_lang.n_words, hidden_size).to(device)\n",
    "decoder = AttnDecoderRNN(hidden_size, output_lang.n_words).to(device)\n",
    "plot_losses,plot_losses_val = trainIters(encoder,decoder,n_epochs,print_every=2500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(running_mean(plot_losses))\n",
    "plt.plot([2500*i for i in range(int(n_epochs/2500))], plot_losses_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(encoder, decoder, sentence, max_length=MAX_LENGTH):\n",
    "    with torch.no_grad():\n",
    "        input_tensor = tensorFromSentence(input_lang, sentence)\n",
    "        input_length = input_tensor.size()[0]\n",
    "        encoder_hidden = encoder.initHidden()\n",
    "\n",
    "        encoder_outputs = encoder(input_tensor)\n",
    "\n",
    "        decoder_input = torch.tensor([[SOS_token]], device=device)  # SOS\n",
    "\n",
    "        decoder_hidden = decoder.initHidden()\n",
    "\n",
    "        decoded_words = []\n",
    "        decoder_attentions = torch.zeros(max_length, max_length)\n",
    "\n",
    "        for di in range(max_length):\n",
    "            decoder_output, decoder_hidden, decoder_attention = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs)\n",
    "            decoder_attentions[di,:input_length] = decoder_attention[:,0,0].data\n",
    "            topv, topi = decoder_output.data.topk(1)\n",
    "            if topi.item() == EOS_token:\n",
    "                decoded_words.append('<EOS>')\n",
    "                break\n",
    "            else:\n",
    "                decoded_words.append(output_lang.index2word[topi.item()])\n",
    "\n",
    "            decoder_input = topi.squeeze().detach()\n",
    "\n",
    "        return decoded_words, decoder_attentions[:di + 1,:input_length]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "francais = \"c est un jeune directeur plein de talent .\"#\"elle a cinq ans de moins que moi .\"\n",
    "output_words, attentions = evaluate(encoder, decoder, francais)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attentions.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.sum(attentions,dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.matshow(attentions.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "indexesFromSentence(input_lang, francais)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorFromSentence(input_lang, francais)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluateRandomly(encoder, decoder, n=10):\n",
    "    for i in range(n):\n",
    "        pair = random.choice(pairs_val)\n",
    "        print('>', pair[0])\n",
    "        print('=', pair[1])\n",
    "        output_words, attentions = evaluate(encoder, decoder, pair[0])\n",
    "        output_sentence = ' '.join(output_words)\n",
    "        print('<', output_sentence)\n",
    "        print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluateRandomly(encoder, decoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "def showAttention(input_sentence, output_words, attentions):\n",
    "    # Set up figure with colorbar\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    cax = ax.matshow(attentions.numpy(), cmap='bone')\n",
    "    fig.colorbar(cax)\n",
    "\n",
    "    # Set up axes\n",
    "    ax.set_xticklabels([''] + input_sentence.split(' ') +\n",
    "                       ['<EOS>'],rotation =90)\n",
    "    ax.set_yticklabels([''] + output_words)\n",
    "\n",
    "    # Show label at every tick\n",
    "    ax.xaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def evaluateAndShowAttention(input_sentence):\n",
    "    output_words, attentions = evaluate(\n",
    "        encoder, decoder, input_sentence)\n",
    "    print('input =', input_sentence)\n",
    "    print('output =', ' '.join(output_words))\n",
    "    showAttention(input_sentence, output_words, attentions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluateAndShowAttention(\"elle a cinq ans de moins que moi .\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluateAndShowAttention(\"elle est trop petite .\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluateAndShowAttention(\"je ne crains pas de mourir .\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluateAndShowAttention(\"c est un jeune directeur plein de talent .\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Dataflowr](https://raw.githubusercontent.com/dataflowr/website/master/_assets/dataflowr_logo.png)](https://dataflowr.github.io/website/)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dldiy",
   "language": "python",
   "name": "dldiy"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}