ch2en_gru.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xGw8DDZT0d5X"
   },
   "outputs": [],
   "source": [
    "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 operator\n",
    "from torch.utils.data import Dataset\n",
    "from collections import Counter\n",
    "import matplotlib.pyplot as plt\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 sacrebleu import corpus_bleu, TOKENIZERS, DEFAULT_TOKENIZER\n",
    "#from masked_cross_entropy import *\n",
    "from torch.optim.lr_scheduler import ExponentialLR\n",
    "from torch.autograd import Variable\n",
    "import random\n",
    "from torch.nn import functional"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pre-processing Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_pairs(pairs):\n",
    "    filtered_pairs = []\n",
    "    for pair in pairs:\n",
    "        if len(pair[0].split()) >= MIN_LENGTH and len(pair[0].split()) <= MAX_LENGTH \\\n",
    "            and len(pair[1].split()) >= MIN_LENGTH and len(pair[1].split()) <= MAX_LENGTH:\n",
    "            filtered_pairs.append(pair)\n",
    "    return filtered_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0-6BiXr1MaHT"
   },
   "outputs": [],
   "source": [
    "def sequence_mask(sequence_length, max_len=None):\n",
    "    \"\"\"\n",
    "    Code paraphrased from \n",
    "    https://github.com/spro/practical-pytorch/blob/master/seq2seq-translation/masked_cross_entropy.py\n",
    "    \"\"\"\n",
    "    if max_len is None:\n",
    "        max_len = sequence_length.data.max()\n",
    "    batch_size = sequence_length.size(0)\n",
    "    seq_range = torch.arange(0, max_len).long()\n",
    "    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len).contiguous()\n",
    "    seq_range_expand = seq_range_expand.to(device)\n",
    "    seq_length_expand = (sequence_length.unsqueeze(1)\n",
    "                         .expand_as(seq_range_expand))\n",
    "    return seq_range_expand < seq_length_expand\n",
    "\n",
    "\n",
    "\n",
    "def masked_cross_entropy(logits, target, length):\n",
    "    length = torch.LongTensor(length).to(device)\n",
    "    \"\"\"\n",
    "    Code paraphrased from \n",
    "    https://github.com/spro/practical-pytorch/blob/master/seq2seq-translation/masked_cross_entropy.py\n",
    "    \"\"\"\n",
    "    \n",
    "    \n",
    "    \"\"\"\n",
    "    Args:\n",
    "        logits: A Variable containing a FloatTensor of size\n",
    "            (batch, max_len, num_classes) which contains the\n",
    "            unnormalized probability for each class.\n",
    "        target: A Variable containing a LongTensor of size\n",
    "            (batch, max_len) which contains the index of the true\n",
    "            class for each corresponding step.\n",
    "        length: A Variable containing a LongTensor of size (batch,)\n",
    "            which contains the length of each data in a batch.\n",
    "\n",
    "    Returns:\n",
    "        loss: An average loss value masked by the length.\n",
    "    \"\"\"\n",
    "\n",
    "    # logits_flat: (batch * max_len, num_classes)\n",
    "    logits_flat = logits.view(-1, logits.size(-1))\n",
    "    # log_probs_flat: (batch * max_len, num_classes)\n",
    "    log_probs_flat = functional.log_softmax(logits_flat, dim=1)\n",
    "    # target_flat: (batch * max_len, 1)\n",
    "    target_flat = target.view(-1, 1)\n",
    "    # losses_flat: (batch * max_len, 1)\n",
    "    losses_flat = -torch.gather(log_probs_flat, dim=1, index=target_flat)\n",
    "    # losses: (batch, max_len)\n",
    "    losses = losses_flat.view(*target.size())\n",
    "    # mask: (batch, max_len)\n",
    "    mask = sequence_mask(sequence_length=length, max_len=target.size(1))\n",
    "    losses = losses * mask.float()\n",
    "    loss = losses.sum() / length.float().sum()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EYDlc9JJ0d5Z"
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "PAD_token = 0\n",
    "SOS_token = 1\n",
    "EOS_token = 2\n",
    "hidden_size = 512\n",
    "dropout_p = 0.1\n",
    "teacher_forcing_ratio = 1\n",
    "BATCH_SIZE = 64\n",
    "MIN_LENGTH = 3\n",
    "MAX_LENGTH = 60\n",
    "source_vocab_size = 4800\n",
    "target_vocab_size = 50000\n",
    "n_layers = 4\n",
    "lr_rate_en = 0.0001\n",
    "lr_rate_de = 0.0005\n",
    "lr_decay = True\n",
    "gamma_encoder = 0.9\n",
    "gamma_decoder = 0.9\n",
    "n_epochs = 20\n",
    "plot_every = 100\n",
    "print_every = 100\n",
    "evaluate_every = 100\n",
    "attn_model = 'dot'\n",
    "Attention = True\n",
    "search_method = 'greedy'\n",
    "beam_size = 10\n",
    "n_best = 5\n",
    "sentence_ratio = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OESV2zjg0d5c"
   },
   "source": [
    "## Loading Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "X0G2eDBP0d5c"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'感 应 电 机 可 以 加 热 东 西 ， 尤 其 擅 长 加 热 钢 铁'"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Lang:\n",
    "    '''\n",
    "    Some codes are paraphrased from\n",
    "    https://pytorch.org/tutorials/intermediate/seq2seq_translation_tutorial.html\n",
    "    '''\n",
    "    def __init__(self, name):\n",
    "        self.name = name\n",
    "        self.word2index = {}\n",
    "        self.word2count = {}\n",
    "        self.index2word = {0: \"<PAD>\", 1: \"<SOS>\", 2:\"<EOS>\", 3:\"<UNK>\"}\n",
    "        self.n_words = 4  # 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",
    "# def unicodeToAscii(s):\n",
    "#     \"\"\" Turn a Unicode string to plain ASCII, \n",
    "#     thanks to http://stackoverflow.com/a/518232/2809427\"\"\"\n",
    "\n",
    "#     return ''.join(c for c in unicodedata.normalize('NFD', s)\n",
    "#         if unicodedata.category(c) != 'Mn')\n",
    "\n",
    "def normalizeEnString(s):\n",
    "#     s = unicodeToAscii(s.strip())\n",
    "#     s = s.replace(\"&apos;\", \"&apos\").replace(\"&quot\",\"\")\n",
    "   \n",
    "    s = re.sub(r\"([.!?])\", r\" \\1\", s)\n",
    "    s = re.sub(r\"[^a-zA-Z.!?0-9]+\", r\" \", s)\n",
    "#     s = s.replace(\"apos\", \"'\")\n",
    "    return s\n",
    "\n",
    "def normalizeChString(s):\n",
    "    s = \"\".join(s.split())\n",
    "    s = ' '.join(str(x) for x in s)\n",
    "    return s.strip() \n",
    "\n",
    "\n",
    "normalizeEnString(\"It &apos;s very pretty , and it has rapidly started to overgrow the \\\n",
    "                  once very rich biodiversity of the northwestern Mediterranean .\")\n",
    "normalizeChString('感 应 电 机 可      以 加 热 东    西     ，     尤 其 擅 长 加 热 钢 铁')               "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "def readLangs(lang1, lang2, data='train'):\n",
    "    '''\n",
    "    Some codes are paraphrased from\n",
    "    https://pytorch.org/tutorials/intermediate/seq2seq_translation_tutorial.html\n",
    "    '''\n",
    "    print(\"Reading lines...\")\n",
    "    # Read the file and split into lines\n",
    "    zh_lines = open('data/iwslt-zh-en/{}.tok.zh'.format(data)).read().split('\\n')\n",
    "    en_lines = open('data/iwslt-zh-en/{}.tok.en'.format(data)).read().split('\\n')\n",
    "    # Split every line into pairs and normalize\n",
    "    pairs = [[normalizeChString(element[0]), normalizeEnString(element[1])] for element in zip(zh_lines, en_lines)]\n",
    "\n",
    "    input_lang = Lang(lang1)\n",
    "    output_lang = Lang(lang2)\n",
    "\n",
    "    return input_lang, output_lang, pairs\n",
    "\n",
    "def prepareData(lang1, lang2):\n",
    "    input_lang, output_lang, pairs = readLangs(lang1, lang2)\n",
    "    print(\"Read %s sentence pairs\" % len(pairs))\n",
    "    pairs = filter_pairs(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",
    "def build_topwordVocab(lang, vocab_size):\n",
    "    print(\"Build vocabulary by top {} frequent word...\".format(vocab_size))\n",
    "    sorted_word2Count = sorted(lang.word2count.items(),\n",
    "        key=operator.itemgetter(1),\n",
    "        reverse=True)\n",
    "    sorted_words = [x[0] for x in sorted_word2Count[:vocab_size]]\n",
    "    \n",
    "    lang.word2index = {}\n",
    "\n",
    "    for ind, word in enumerate(sorted_words):\n",
    "            lang.word2index[word] = ind + 4\n",
    "    lang.index2word = {}\n",
    "    lang.index2word[0] = \"<PAD>\"\n",
    "    lang.index2word[1] = \"<SOS>\"\n",
    "    lang.index2word[2] = \"<EOS>\"\n",
    "    lang.index2word[3] = \"<UNK>\"\n",
    "\n",
    "    for ind, word in enumerate(sorted_words):\n",
    "        lang.index2word[ind + 4] = word\n",
    "    \n",
    "    lang.n_words = len(lang.index2word)\n",
    "    \n",
    "    print(lang.name, lang.n_words)\n",
    "    return lang\n",
    "\n",
    "input_lang, output_lang, pairs = prepareData('ch', 'eng')\n",
    "\n",
    "input_lang = build_topwordVocab(input_lang,vocab_size=source_vocab_size)\n",
    "output_lang = build_topwordVocab(output_lang, vocab_size=target_vocab_size)\n",
    "print(random.choice(pairs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "I_PL2_xh0d5y",
    "outputId": "2dac260b-34df-44ab-a8b0-19a84845b789"
   },
   "outputs": [],
   "source": [
    "_, _, val_pairs = readLangs('ch', 'eng', 'dev')\n",
    "val_pairs = filter_pairs(val_pairs[:-1])\n",
    "_, _, test_pairs = readLangs('ch', 'eng', 'test')\n",
    "test_pairs = filter_pairs(test_pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 54
    },
    "colab_type": "code",
    "id": "FJv4sN4W0d52",
    "outputId": "17fafc52-c26d-44c8-b68c-613e77c2cea3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['他 们 会 把 尸 体 上 的 肉 吃 的 干 干 净 净 只 剩 白 骨 。', 'They clean up carcasses right to the bone .']\n"
     ]
    }
   ],
   "source": [
    "print(random.choice(val_pairs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vmMewXau0d6E"
   },
   "source": [
    "# Preparing Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SPnqKF_F0d6F"
   },
   "outputs": [],
   "source": [
    "def indexesFromSentence(lang, sentence):\n",
    "    idxs = []\n",
    "    for word in sentence.split(' '):\n",
    "        try:\n",
    "            idxs.append(lang.word2index[word])\n",
    "        except KeyError:\n",
    "            idxs.append(3)  # 3 is the id of 'UNK'\n",
    "    idxs.append(EOS_token)\n",
    "    return idxs\n",
    "\n",
    "\n",
    "class VocabDataset(Dataset):\n",
    "    def __init__(self, pairs):\n",
    "        \n",
    "        self.source_sent_list = [indexesFromSentence(input_lang,pair[0]) for pair in pairs]\n",
    "        self.target_sent_list = [indexesFromSentence(output_lang,pair[1]) for pair in pairs]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.source_sent_list)\n",
    "        \n",
    "    def __getitem__(self, key):\n",
    "        token1_idx = self.source_sent_list[key]\n",
    "        token2_idx = self.target_sent_list[key]\n",
    "        return [token1_idx,token2_idx, len(token1_idx), len(token2_idx)]\n",
    "\n",
    "    \n",
    "def Vocab_collate_func(batch):\n",
    "    source_sent_list = []\n",
    "    target_sent_list = []\n",
    "    source_len_list = []\n",
    "    target_len_list = []\n",
    "\n",
    "    for datum in batch:   ### batch = sample\n",
    "        source_len_list.append(datum[2])\n",
    "        target_len_list.append(datum[3])\n",
    "    \n",
    "    max_len_src = max(source_len_list)\n",
    "    max_len_trg = max(target_len_list)\n",
    "    \n",
    "    # padding\n",
    "    for datum in batch:\n",
    "        \n",
    "        # source sentence processing\n",
    "        padded_source = np.pad(np.array(datum[0]), \n",
    "                                pad_width=((0,max_len_src-datum[2])),          ### 0代表左边没有pad,右边的值代表右边pad的个数\n",
    "                                mode=\"constant\", constant_values=PAD_token)\n",
    "        source_sent_list.append(padded_source)\n",
    "        \n",
    "        # target sentence processing\n",
    "        padded_target = np.pad(np.array(datum[1]), \n",
    "                                pad_width=((0,max_len_trg-datum[3])),          ### 0代表左边没有pad,右边的值代表右边pad的个数\n",
    "                                mode=\"constant\", constant_values=PAD_token)\n",
    "        target_sent_list.append(padded_target)\n",
    "        \n",
    "    #sort sentences for the batch\n",
    "    sort_idx = sorted(range(len(source_len_list)), key=source_len_list.__getitem__, reverse=True)\n",
    "    source_sent_list = np.array(source_sent_list)[sort_idx]\n",
    "    target_sent_list = np.array(target_sent_list)[sort_idx]\n",
    "    source_len_list = np.array(source_len_list)[sort_idx]\n",
    "    target_len_list = np.array(target_len_list)[sort_idx]\n",
    "        \n",
    "    return [torch.tensor(source_sent_list).to(device), \n",
    "            torch.tensor(target_sent_list).to(device),\n",
    "            torch.LongTensor(source_len_list), \n",
    "            torch.LongTensor(target_len_list)]\n",
    "\n",
    "train_dataset = VocabDataset(pairs)\n",
    "train_loader = torch.utils.data.DataLoader(dataset=train_dataset, \n",
    "                                           batch_size=BATCH_SIZE,\n",
    "                                           collate_fn=Vocab_collate_func,\n",
    "                                           shuffle=True)\n",
    "\n",
    "val_dataset = VocabDataset(val_pairs)\n",
    "val_loader = torch.utils.data.DataLoader(dataset=val_dataset,\n",
    "                                        batch_size=BATCH_SIZE,\n",
    "                                        collate_fn=Vocab_collate_func,\n",
    "                                        shuffle=False)\n",
    "\n",
    "\n",
    "test_dataset = VocabDataset(test_pairs)\n",
    "test_loader = torch.utils.data.DataLoader(dataset=test_dataset,\n",
    "                                        batch_size=BATCH_SIZE,\n",
    "                                        collate_fn=Vocab_collate_func,\n",
    "                                        shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "PhwZoZz_0d6Q"
   },
   "source": [
    "# Build Encoder-Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "E4lH5BjW10Jj"
   },
   "outputs": [],
   "source": [
    "class EncoderRNN(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, n_layers=1, dropout=0.1):\n",
    "        super(EncoderRNN, self).__init__()\n",
    "        \n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.n_layers = n_layers\n",
    "        self.dropout = dropout\n",
    "        self.embedding = nn.Embedding(input_size, hidden_size)\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size, n_layers, dropout=self.dropout, bidirectional=True)\n",
    "        \n",
    "    def forward(self, input_seqs, input_lengths, hidden):\n",
    "        # Note: we run this all at once (over multiple batches of multiple sequences)\n",
    "        embedded = self.embedding(input_seqs)\n",
    "        packed = torch.nn.utils.rnn.pack_padded_sequence(embedded, input_lengths)\n",
    "        outputs, hidden = self.gru(packed, hidden)\n",
    "        outputs, output_lengths = torch.nn.utils.rnn.pad_packed_sequence(outputs) # unpack (back to padded)\n",
    "        outputs = outputs[:, :, :self.hidden_size] + outputs[:, : ,self.hidden_size:] # Sum bidirectional outputs\n",
    "        return outputs, hidden, output_lengths\n",
    "      \n",
    "    def initHidden(self,batch_size):\n",
    "        return torch.zeros(2 * self.n_layers, batch_size, self.hidden_size, device=device)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EF4Fj5_pN9Lv"
   },
   "outputs": [],
   "source": [
    "class DecoderRNN(nn.Module):\n",
    "    def __init__(self, hidden_size, output_size,dropout_p=0.1, n_layers=1, max_length=MAX_LENGTH):\n",
    "        super(DecoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.n_layers = n_layers\n",
    "        self.dropout_p = dropout_p\n",
    "        self.output_size = output_size\n",
    "        \n",
    "        \n",
    "        self.embedding = nn.Embedding(output_size, hidden_size,padding_idx=PAD_token)\n",
    "        self.gru = nn.GRU(self.hidden_size, self.hidden_size, n_layers, dropout=dropout_p)\n",
    "        self.out = nn.Linear(self.hidden_size, self.output_size)\n",
    "        self.dropout = nn.Dropout(self.dropout_p)\n",
    "        self.softmax = nn.LogSoftmax(dim=1)\n",
    "\n",
    "    def forward(self, input_seqs, hidden, batch_size):\n",
    "        embedded = self.embedding(input_seqs).view(1, batch_size, -1)\n",
    "        embedded = self.dropout(embedded)\n",
    "        output = F.relu(embedded)\n",
    "        output, hidden = self.gru(output, hidden)\n",
    "        output = self.softmax(self.out(output[0]))\n",
    "        return output, hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SbhB0-fFEd8Q"
   },
   "outputs": [],
   "source": [
    "class Attn(nn.Module):\n",
    "    \"\"\"\n",
    "    Some code is paraphrased from\n",
    "    https://github.com/spro/practical-pytorch/blob/master/seq2seq-translation/seq2seq-translation-batched.ipynb\n",
    "    \"\"\"\n",
    "    def __init__(self, method, hidden_size):\n",
    "        super(Attn, self).__init__()\n",
    "        \n",
    "        self.method = method\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        if self.method == 'concat':\n",
    "            self.attn = nn.Linear(self.hidden_size * 2, hidden_size)\n",
    "            self.v = nn.Parameter(torch.FloatTensor(hidden_size))\n",
    "\n",
    "    def forward(self, hidden, encoder_outputs):\n",
    "        max_len = encoder_outputs.size(0)\n",
    "        this_batch_size = encoder_outputs.size(1)\n",
    "\n",
    "        # Create variable to store attention energies\n",
    "        attn_energies = Variable(torch.zeros(this_batch_size, max_len)).to(device) # B x S\n",
    "        \n",
    "        if self.method == 'dot':\n",
    "\n",
    "            attn_energies = torch.matmul(encoder_outputs.permute(1,0,2), hidden.permute(1,2,0)).squeeze()\n",
    "            \n",
    "        if self.method == 'concat':\n",
    "            hidden_expand = hidden.expand(max_len, -1, -1).permute(1, 0, 2)  # shape of (B, S, N)\n",
    "            enc_cat_hid = torch.cat([encoder_outputs.permute(1,0,2), hidden_expand], dim=-1)  # shape of (B, S, 2*N)\n",
    "            # After nn.Linear(2*N, N), enc_cat_hid with shape (B, S, N)\n",
    "            # v is shape of (N)\n",
    "            attn_energies = torch.matmul(self.attn(enc_cat_hid), self.v)  # shape of (B, S)\n",
    "        # Normalize energies to weights in range 0 to 1, resize to 1 x B x S\n",
    "        # Dangerous\n",
    "        if attn_energies.dim() == 1:\n",
    "            attn_energies = attn_energies.unsqueeze(0)\n",
    "        return F.softmax(attn_energies, dim=1).unsqueeze(1)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LqklCWyTEnpG"
   },
   "outputs": [],
   "source": [
    "class LuongAttnDecoderRNN(nn.Module):\n",
    "    \"\"\"\n",
    "    Some code is paraphrased from\n",
    "    https://github.com/spro/practical-pytorch/blob/master/seq2seq-translation/seq2seq-translation-batched.ipynb\n",
    "    \"\"\"\n",
    "    def __init__(self, attn_model, hidden_size, output_size, n_layers=1, dropout=0.1):\n",
    "        super(LuongAttnDecoderRNN, self).__init__()\n",
    "\n",
    "        # Keep for reference\n",
    "        self.attn_model = attn_model\n",
    "        self.hidden_size = hidden_size\n",
    "        self.output_size = output_size\n",
    "        self.n_layers = n_layers\n",
    "        self.dropout = dropout\n",
    "\n",
    "        # Define layers\n",
    "        self.embedding = nn.Embedding(output_size, hidden_size).to(device)\n",
    "        self.embedding_dropout = nn.Dropout(dropout)\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size, n_layers, dropout=dropout)\n",
    "        self.concat = nn.Linear(hidden_size * 2, hidden_size)\n",
    "        self.out = nn.Linear(hidden_size, output_size)\n",
    "        \n",
    "        # Choose attention model\n",
    "        if attn_model != 'none':\n",
    "            self.attn = Attn(attn_model, hidden_size)\n",
    "\n",
    "    def forward(self, input_seq, last_hidden, encoder_outputs):\n",
    "        # Note: we run this one step at a time\n",
    "\n",
    "        # Get the embedding of the current input word (last output word)\n",
    "        batch_size = input_seq.size(0)\n",
    "        embedded = self.embedding(input_seq)\n",
    "        embedded = self.embedding_dropout(embedded)\n",
    "        embedded = embedded.view(1, batch_size, self.hidden_size) # S=1 x B x N\n",
    "\n",
    "        # Get current hidden state from input word and last hidden state\n",
    "        rnn_output, hidden = self.gru(embedded, last_hidden)\n",
    "\n",
    "        # Calculate attention from current RNN state and all encoder outputs;\n",
    "        # apply to encoder outputs to get weighted average\n",
    "        attn_weights = self.attn(rnn_output, encoder_outputs)\n",
    "        context = attn_weights.bmm(encoder_outputs.transpose(0, 1)) # B x S=1 x N\n",
    "\n",
    "        # Attentional vector using the RNN hidden state and context vector\n",
    "        # concatenated together (Luong eq. 5)\n",
    "        rnn_output = rnn_output.squeeze(0) # S=1 x B x N -> B x N\n",
    "        context = context.squeeze(1)       # B x S=1 x N -> B x N\n",
    "        concat_input = torch.cat((rnn_output, context), 1)\n",
    "        concat_output = torch.tanh(self.concat(concat_input))\n",
    "\n",
    "        # Finally predict next token (Luong eq. 6, without softmax)\n",
    "        output = self.out(concat_output)\n",
    "\n",
    "        # Return final output, hidden state, and attention weights (for visualization)\n",
    "        return output, hidden, attn_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "d1GxY0qw0d6V"
   },
   "source": [
    "# Training Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4ICvoIhH0d6W"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import math\n",
    "\n",
    "\n",
    "def asMinutes(s):\n",
    "    m = math.floor(s / 60)\n",
    "    s -= m * 60\n",
    "    return '%dm %ds' % (m, s)\n",
    "\n",
    "\n",
    "def timeSince(since, percent):\n",
    "    now = time.time()\n",
    "    s = now - since\n",
    "    es = s / (percent)\n",
    "    rs = es - s\n",
    "    return '%s (- %s)' % (asMinutes(s), asMinutes(rs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qiiPjqOy0d6Y"
   },
   "outputs": [],
   "source": [
    "def train(input_tensor, target_tensor, input_lengths, target_lengths, encoder, decoder, \n",
    "          encoder_optimizer, decoder_optimizer, clip=10.0):\n",
    "    encoder_optimizer.zero_grad()  # zero out the accumulated gradient over mini-batch\n",
    "    decoder_optimizer.zero_grad()\n",
    "    \n",
    "    \n",
    "    batch_size = input_tensor.size(1)\n",
    "    input_tensor = input_tensor.to(device)\n",
    "    target_tensor = target_tensor.to(device)\n",
    "\n",
    "    encoder_hidden = encoder.initHidden(batch_size)\n",
    "    encoder_outputs = torch.zeros(input_lengths.max(), batch_size, encoder.hidden_size, device=device) \n",
    " \n",
    "\n",
    "    encoder_outputs, encoder_hidden, encoder_output_lengths = encoder(input_tensor, input_lengths, encoder_hidden)\n",
    "    #encoder_outputs:  # max_len x batch_size x hidden_size\n",
    "    #hidden: n_layers * 2 x batch_size x hidden_size\n",
    "    loss = 0\n",
    "\n",
    "    \n",
    "    decoder_input = torch.tensor([SOS_token]*batch_size).to(device)  # decoder_input: torch.Size([1, 32])\n",
    "    decoder_hidden = encoder_hidden[:decoder.n_layers] # Use last (forward) hidden state from encoder\n",
    "    all_decoder_outputs = Variable(torch.zeros(target_lengths.max(), batch_size, decoder.output_size)).to(device)\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_lengths.max()):\n",
    "            if Attention:\n",
    "                decoder_output, decoder_hidden, decoder_attention = decoder(\n",
    "                    decoder_input, decoder_hidden, encoder_outputs)\n",
    "     \n",
    "            else:\n",
    "                decoder_output, decoder_hidden = decoder(\n",
    "                    decoder_input, decoder_hidden, batch_size)\n",
    "            \n",
    "            decoder_input = target_tensor[di]  # Teacher forcing\n",
    "            all_decoder_outputs[di] = decoder_output\n",
    "            \n",
    "    else:\n",
    "        # Without teacher forcing: use its own predictions as the next input\n",
    "        for di in range(target_lengths.max()):\n",
    "          \n",
    "            if Attention:\n",
    "                decoder_output, decoder_hidden, decoder_attention = decoder(\n",
    "                    decoder_input, decoder_hidden, encoder_outputs)\n",
    "     \n",
    "            else:\n",
    "                decoder_output, decoder_hidden = decoder(\n",
    "                    decoder_input, decoder_hidden, batch_size)\n",
    "            \n",
    "            topv, topi = decoder_output.topk(1)\n",
    "\n",
    "            decoder_input = topi.squeeze().detach()  # detach from history as input\n",
    "            all_decoder_outputs[di] = decoder_output\n",
    "           \n",
    "            \n",
    "    # Loss calculation and backpropagation\n",
    "\n",
    "    loss = masked_cross_entropy(\n",
    "        all_decoder_outputs.transpose(0, 1).contiguous(), # -> batch x seq\n",
    "        target_tensor.transpose(0, 1).contiguous(), # -> batch x seq\n",
    "        target_lengths\n",
    "    )\n",
    "#     loss = loss.sum()/batch_size \n",
    "    loss.backward()\n",
    "    #    ave_loss.backward()\n",
    "    \n",
    "    # Clip gradient norms\n",
    "    torch.nn.utils.clip_grad_norm_(encoder.parameters(), clip)\n",
    "    torch.nn.utils.clip_grad_norm_(decoder.parameters(), clip)\n",
    "    \n",
    "    encoder_optimizer.step()   # update parameters\n",
    "    decoder_optimizer.step()\n",
    "\n",
    "    return loss.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HmHvipVe0d6a"
   },
   "outputs": [],
   "source": [
    "def trainIters(encoder, decoder, n_iters, lr_decay=True, gamma_encoder=0.9, gamma_decoder=0.9, print_every=100, plot_every=100, learning_rate_encoder=0.0005, learning_rate_decoder=0.002,evaluate_every=3000):\n",
    "    start = time.time()\n",
    "\n",
    "    encoder_optimizer = optim.Adam(encoder.parameters(), lr=learning_rate_encoder)\n",
    "    decoder_optimizer = optim.Adam(decoder.parameters(), lr=learning_rate_decoder)\n",
    "    \n",
    "    scheduler_encoder = ExponentialLR(encoder_optimizer, gamma_encoder, last_epoch=-1) \n",
    "    scheduler_decoder = ExponentialLR(decoder_optimizer, gamma_decoder, last_epoch=-1) \n",
    "    \n",
    "    encoder.to(device)\n",
    "    decoder.to(device)\n",
    "    score_max = 0\n",
    "    plot_losses = []\n",
    "    validation_scores = []\n",
    "    \n",
    "    for epoch in range(1, n_iters + 1):\n",
    "        print_loss_total = 0  # Reset every print_every\n",
    "        plot_loss_total = 0  # Reset every plot_every\n",
    "        if lr_decay:\n",
    "            scheduler_encoder.step()\n",
    "            scheduler_decoder.step()\n",
    "        \n",
    "        for i, (input_sentences, target_sentences,len1,len2) in enumerate(train_loader): \n",
    "            encoder.train()\n",
    "            decoder.train()\n",
    "            \n",
    "            input_tensor = input_sentences.transpose(0,1)   # 13*100 to 100*13\n",
    "            target_tensor = target_sentences.transpose(0,1)\n",
    "            loss = train(input_tensor, target_tensor, len1, len2, encoder,\n",
    "                         decoder, encoder_optimizer, decoder_optimizer)\n",
    "            print_loss_total += loss\n",
    "            plot_loss_total += loss\n",
    "            \n",
    "            if i > 0 and i % evaluate_every == 0:\n",
    "                bleu_score, (src_sents, sys_sents, ref_sents) = test_model(encoder, decoder, val_loader)\n",
    "                print('Validation Score: {} \\n source sentence {} \\n predicted sentence {} \\n Reference sentence: {}'.format(bleu_score,src_sents, sys_sents, ref_sents))\n",
    "                validation_scores.append(bleu_score)\n",
    "                \n",
    "                if bleu_score > score_max:\n",
    "                    score_max = bleu_score\n",
    "                \n",
    "                    torch.save({\n",
    "                                'epoch': epoch,\n",
    "                                'encoder': encoder.state_dict(),\n",
    "                                'encoder_optimizer': encoder_optimizer.state_dict(),\n",
    "                                'decoder': decoder.state_dict(),\n",
    "                                'decoder_optimizer': decoder_optimizer.state_dict()\n",
    "                                }, \"saved_model/attnIs{}_hiddenSize{}_nLayer{}_batchSize{}_epoch{}_srcVocSize{}_tgtVocSize{}_lrDecay{}_teacherF{}\"\\\n",
    "                        .format(Attention,hidden_size,n_layers,BATCH_SIZE,n_iters,source_vocab_size,\n",
    "                                target_vocab_size,lr_decay,teacher_forcing_ratio))   \n",
    "                    \n",
    "            if i > 0 and i % print_every == 0:\n",
    "                print_loss_avg = print_loss_total / print_every\n",
    "                print_loss_total = 0\n",
    "#                bleu_score, (sys_sents, ref_sents) = test_model(encoder, decoder, val_loader)\n",
    "                print('Time: {}, Epoch: [{}/{}], Step: [{}/{}], Train Loss: {}'.format(\n",
    "                    timeSince(start, i + 1/len(train_loader)), epoch, n_iters, i, \n",
    "                    len(train_loader),print_loss_avg))\n",
    "\n",
    "            if i > 0 and i % plot_every == 0:\n",
    "                plot_loss_avg = plot_loss_total / plot_every\n",
    "                plot_losses.append(plot_loss_avg)\n",
    "                plot_loss_total = 0\n",
    "                torch.save({\n",
    "                            'plot_losses': plot_losses,\n",
    "                            'validation_scores': validation_scores\n",
    "                            }, \"saved_scores/attnIs{}_hiddenSize{}_nLayer{}_batchSize{}_epoch{}_srcVocSize{}_tgtVocSize{}_lrDecay{}_teacherF{}\"\\\n",
    "                    .format(Attention,hidden_size,n_layers,BATCH_SIZE,n_iters,source_vocab_size,\n",
    "                            target_vocab_size,lr_decay,teacher_forcing_ratio))   \n",
    "                \n",
    "            torch.cuda.empty_cache()    \n",
    "        print(\"plot_losses:\",plot_losses)\n",
    "        print(\"validation_scores:\",validation_scores)\n",
    "    showPlot(plot_losses)\n",
    "    showPlot(validation_scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mj8OwEU_0d6d"
   },
   "source": [
    "# Plotting results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bDqy1HHH0d6e"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.switch_backend('agg')\n",
    "import matplotlib.ticker as ticker\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def showPlot(points):\n",
    "    plt.figure()\n",
    "    fig, ax = plt.subplots()\n",
    "    # this locator puts ticks at regular intervals\n",
    "    loc = ticker.MultipleLocator(base=0.2)\n",
    "    ax.yaxis.set_major_locator(loc)\n",
    "    plt.plot(points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "L_-Y9XII0d6g"
   },
   "source": [
    "# Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "class beam_search(object):\n",
    "    \"\"\"\n",
    "    Some code is paraphrased from\n",
    "    https://github.com/OpenNMT/OpenNMT-py/blob/master/onmt/translate/beam.py\n",
    "    \"\"\"\n",
    "    def __init__(self, encoder, decoder, max_length, beam_size, attention = True,sentence_ratio = False): \n",
    "        super(beam_search, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.attention = attention\n",
    "        self.max_length = max_length\n",
    "        self.beam_size = beam_size\n",
    "        self.sentence_ratio = sentence_ratio\n",
    "        \n",
    "    def search(self, encoder_outputs, decoder_input, decoder_hidden, src_len):\n",
    "\n",
    "        prob = {k:0 for k in range(self.beam_size)}\n",
    "        bestSent = []\n",
    "        bestScore = []       \n",
    "        decoder_word_choices = {k:[] for k in range(self.beam_size)}\n",
    "        decoder_hidden_choices = {}\n",
    "        decoder_input_choices = {}\n",
    "        decoder_output_choices = {}\n",
    "        \n",
    "        # Initialize beam serach\n",
    "        if self.attention == True:\n",
    "            decoder_output, decoder_hidden, decoder_attention = self.decoder(decoder_input.contiguous(), decoder_hidden.contiguous(), encoder_outputs)\n",
    "            decoder_output = F.log_softmax(decoder_output, dim=1)\n",
    "            topv, topi = decoder_output.data.topk(self.beam_size)\n",
    "        else: \n",
    "            print(\"Only available when attention = True\")\n",
    "        \n",
    "        # Initialize beam candidates \n",
    "        for i in range(self.beam_size):\n",
    "            decoder_word_choices[i].append(topi.squeeze()[i].item())\n",
    "            decoder_input_choices[i] = topi.squeeze()[i].detach()\n",
    "            decoder_hidden_choices[i] = decoder_hidden\n",
    "            prob[i] += topv.squeeze()[i].detach()\n",
    "            \n",
    "        ## running beam search\n",
    "        cur_len = 0\n",
    "        max_length = 2*src_len if self.sentence_ratio else self.max_length\n",
    "        # delete\n",
    "#         print(self.sentence_ratio)\n",
    "#         print(src_len)\n",
    "#         print(max_length)\n",
    "        \n",
    "        while decoder_hidden_choices and cur_len <= max_length:\n",
    "            cur_len += 1\n",
    "            topi = {}\n",
    "            key_list = list(decoder_hidden_choices.keys())\n",
    "            scores = []\n",
    "            for key in key_list:\n",
    "                    \n",
    "                decoder_output, decoder_hidden_choices[key],decoder_attn  = self.decoder(decoder_input_choices[key].unsqueeze(0), decoder_hidden_choices[key],encoder_outputs)\n",
    "                decoder_output_choices[key] = F.log_softmax(decoder_output, dim=1)\n",
    "                topv, topi[key] = decoder_output_choices[key].data.topk(len(decoder_hidden_choices))\n",
    "                scores.extend((topv+prob[key]).tolist()[0])\n",
    "                \n",
    "            scores = np.array(scores)   \n",
    "            max_candidate_score = scores.argsort()[-len(decoder_hidden_choices):][::-1]\n",
    "            decoded_sent_score = scores[max_candidate_score]\n",
    "\n",
    "            choice_sentence = {}\n",
    "            choiceHidden = {}\n",
    "            trashOfKeys = []\n",
    "            \n",
    "            for j in range(len(max_candidate_score)):\n",
    "                prev_choice_idx = key_list[int(np.floor(max_candidate_score[j]/len(decoder_hidden_choices)))]\n",
    "                if topi[prev_choice_idx].squeeze().dim() == 0:\n",
    "                    next_idx = topi[prev_choice_idx].squeeze()\n",
    "                else:\n",
    "                    next_idx = topi[prev_choice_idx].squeeze()[max_candidate_score[j] % len(decoder_hidden_choices)]\n",
    "                \n",
    "                s_choice = decoder_word_choices[prev_choice_idx].copy()\n",
    "                s_choice.append(next_idx.item())\n",
    "                choice_sentence[j] = s_choice\n",
    "                h_choice = decoder_hidden_choices[prev_choice_idx]\n",
    "                choiceHidden[j] = h_choice\n",
    "                decoder_input_choices[j] = next_idx.detach()   \n",
    "                prob[j] = decoded_sent_score[j] \n",
    "   \n",
    "            decoder_word_choices = choice_sentence\n",
    "            decoder_hidden_choices = choiceHidden\n",
    "            \n",
    "            for key, s in decoder_word_choices.items():\n",
    "                if EOS_token in s:\n",
    "                    bestSent.append(s)\n",
    "                    bestScore.append(prob[key]) \n",
    "                    trashOfKeys.append(key)\n",
    "                    \n",
    "            for k in trashOfKeys:\n",
    "                decoder_hidden_choices.pop(k)\n",
    "                decoder_word_choices.pop(k)\n",
    "\n",
    "        if len(bestScore) == 0:\n",
    "            max_prob = prob[0]\n",
    "            max_prob_idx = 0\n",
    "            for k in prob.keys():\n",
    "                if prob[k] > max_prob: \n",
    "                    max_prob_idx = k\n",
    "                    max_prob = prob[k]\n",
    "            bestScore.append(max_prob)\n",
    "            bestSent.append(decoder_word_choices[max_prob_idx])\n",
    "                \n",
    "        return bestSent, bestScore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "muc68_Oz0d6h"
   },
   "outputs": [],
   "source": [
    "def get_batch_outputs(encoder, decoder, input_sentences, input_lengths, output_lengths): \n",
    "    with torch.no_grad():\n",
    "        input_tensor = input_sentences.transpose(0,1).to(device)   # 32*100 to 100*32\n",
    "        batch_size = input_tensor.size(1)\n",
    "        encoder_hidden = encoder.initHidden(batch_size)\n",
    "        encoder_outputs, encoder_hidden,encoder_output_lengths = encoder(input_tensor, input_lengths, encoder_hidden)\n",
    "       \n",
    "        decoder_hidden = encoder_hidden[:decoder.n_layers]\n",
    "\n",
    "        decoder_input = Variable(torch.tensor([SOS_token]*batch_size)).to(device)  # decoder_input: torch.Size([1, 32])\n",
    "        decoded_words = np.empty((output_lengths.max(), batch_size), dtype=object)\n",
    "\n",
    "        for di in range(output_lengths.max()):\n",
    "            if Attention:\n",
    "                decoder_output, decoder_hidden, decoder_attention = decoder(\n",
    "                    decoder_input, decoder_hidden, encoder_outputs)\n",
    "            else:\n",
    "                decoder_output, decoder_hidden = decoder(\n",
    "                    decoder_input, decoder_hidden, batch_size)\n",
    "\n",
    "            topv, topi = decoder_output.topk(1)\n",
    "            decoder_input = topi.squeeze().detach().to(device)  # detach from history as input\n",
    "\n",
    "            decoded_words[di:] = np.array(['<EOS>' if idx==EOS_token else output_lang.index2word[idx] for idx in decoder_input.tolist()])\n",
    "\n",
    "        return decoded_words.transpose()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_beam_batch_outputs(encoder, decoder, input_sentences, input_lengths): #####\n",
    "    with torch.no_grad():\n",
    "        input_tensor = input_sentences.transpose(0,1).to(device)   # 32*100 to 100*32\n",
    "        batch_size = input_tensor.size(1)\n",
    "        encoder_hidden = encoder.initHidden(batch_size).to(device)\n",
    "        encoder_outputs, encoder_hidden,encoder_output_lengths = encoder(input_tensor, input_lengths, encoder_hidden)\n",
    "        \n",
    "        decoder_hidden = encoder_hidden[:decoder.n_layers].to(device)   \n",
    "        my_beam_search = beam_search(encoder, decoder,input_sentences.max().item(), beam_size, True, sentence_ratio)\n",
    "        beam_search_result = []\n",
    "        for i in range(batch_size):\n",
    "            decoder_input = torch.tensor([SOS_token], device=device, requires_grad=False).unsqueeze(0)#.view(1,-1) # take care of different input shape\n",
    "            sentences, probs = my_beam_search.search(encoder_outputs[:,i,:].unsqueeze(1), decoder_input, \n",
    "                                                     decoder_hidden[:,i,:].unsqueeze(1), input_lengths[i].item())\n",
    "            \n",
    "            beam_search_result.append(sentences[probs.index(max(probs))])\n",
    "\n",
    "        padded_beam_search_result = []\n",
    "\n",
    "        max_length = 0\n",
    "        for sentence in beam_search_result:\n",
    "            if len(sentence) > max_length:\n",
    "                max_length = len(sentence)\n",
    "\n",
    "        for sentence in beam_search_result:\n",
    "            while len(sentence) < max_length + 2:\n",
    "                sentence.append(PAD_token)\n",
    "            padded_beam_search_result.append(sentence)\n",
    "\n",
    "        batch_sentences = []\n",
    "        \n",
    "        for sentence in padded_beam_search_result:\n",
    "            sentence = [output_lang.index2word[k] for k in sentence]\n",
    "            try:\n",
    "                end_idx = sentence.index('<EOS>')\n",
    "                batch_sentences.append(' '.join(sentence[:end_idx]))\n",
    "            except ValueError:\n",
    "                batch_sentences.append(' '.join(sentence))\n",
    "\n",
    "    return batch_sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ce-QyODD0d6l"
   },
   "outputs": [],
   "source": [
    "def test_model(encoder, decoder, loader, search_method = 'greedy'):\n",
    "    \n",
    "    encoder.eval()\n",
    "    decoder.eval()\n",
    "    \n",
    "    score = []\n",
    "    src_sentences = []\n",
    "    sys_sentences = []\n",
    "    ref_sentences = []\n",
    "    encoder.train(False)\n",
    "    decoder.train(False)\n",
    "    for i, (input_sentences, target_sentences, len1, len2) in enumerate(loader):\n",
    "        for sentence in target_sentences:\n",
    "            trg_list = []\n",
    "            for idx in sentence:\n",
    "                if idx.item() == EOS_token:\n",
    "                    break\n",
    "                else:\n",
    "                    trg_list.append(output_lang.index2word[idx.item()])\n",
    "            ref_sentences.append(' '.join(trg_list))\n",
    "        for sentence in input_sentences:\n",
    "            src_list = []\n",
    "            for idx in sentence:\n",
    "                if idx.item() == EOS_token:\n",
    "                    break\n",
    "                else:\n",
    "                    src_list.append(input_lang.index2word[idx.item()])\n",
    "            src_sentences.append(' '.join(src_list))\n",
    "\n",
    "        #ref_sentences.append(' '.join(sent) for sent in target_sentences)\n",
    "        #src_sentences.append(' '.join(sent) for sent in input_sentences)\n",
    "        batch_size = input_sentences.size(0)\n",
    "        if search_method == 'greedy':\n",
    "            for sentence in get_batch_outputs(encoder, decoder, input_sentences, len1, len2):\n",
    "                try:\n",
    "                    end_idx = sentence.tolist().index('<EOS>')\n",
    "                    sys_sentences.append(' '.join(sentence[:end_idx]))\n",
    "                except ValueError:\n",
    "                    sys_sentences.append(' '.join(sentence))\n",
    "                    \n",
    "        elif search_method == 'beam':\n",
    "            translation_output = get_beam_batch_outputs(encoder, decoder, input_sentences, len1)\n",
    "            sys_sentences.extend(translation_output)\n",
    "            \n",
    "    encoder.train(True)\n",
    "    decoder.train(True) \n",
    "    \n",
    "    score = corpus_bleu(sys_sentences,[ref_sentences], smooth=\"floor\", smooth_floor=0.01, lowercase=False, use_effective_order=True, tokenize=DEFAULT_TOKENIZER).score\n",
    "    return score, (src_sentences[0:1], sys_sentences[0:1], ref_sentences[0:1])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hjafDuaU0d6s"
   },
   "source": [
    "# TRAINING AND EVALUATING"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "T3iWCiIQ0d6t",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "encoder1 = EncoderRNN(input_lang.n_words, hidden_size, n_layers=n_layers).to(device)\n",
    "if Attention:\n",
    "    decoder1 = LuongAttnDecoderRNN(attn_model, hidden_size, output_lang.n_words, dropout=dropout_p, n_layers=n_layers).to(device)\n",
    "else:\n",
    "    decoder1 = DecoderRNN(hidden_size, output_lang.n_words, dropout_p=0)\n",
    "\n",
    "trainIters(encoder1, decoder1, n_iters=n_epochs, print_every=print_every, plot_every=plot_every, evaluate_every=evaluate_every, learning_rate_encoder=lr_rate_en, learning_rate_decoder=lr_rate_de, lr_decay=lr_decay, gamma_encoder=gamma_decoder,\n",
    "          gamma_decoder=gamma_decoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Resume Model and Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder2 = EncoderRNN(input_lang.n_words, hidden_size,n_layers=n_layers).to(device)\n",
    "if not Attention:\n",
    "    decoder2 = DecoderRNN(hidden_size, output_lang.n_words, dropout_p = 0).to(device)\n",
    "else:\n",
    "    decoder2 = LuongAttnDecoderRNN(attn_model, hidden_size, output_lang.n_words, dropout=dropout_p, n_layers=n_layers).to(device)\n",
    "    \n",
    "    \n",
    "# # encoder_optimizer = optim.Adam(encoder1.parameters(), lr=0.003)\n",
    "# # decoder_optimizer = optim.Adam(attn_decoder1.parameters(), lr=0.003)\n",
    "\n",
    "checkpoint = torch.load('saved_model/attnIsTrue_hiddenSize512_nLayer4_batchSize64_epoch20_srcVocSize4300_tgtVocSize39000_lrDecayTrue_teacherF1')\n",
    "encoder2.load_state_dict(checkpoint['encoder'])\n",
    "decoder2.load_state_dict(checkpoint['decoder'])\n",
    "# # encoder_optimizer.load_state_dict(checkpoint['encoder_optimizer'])\n",
    "# # decoder_optimizer.load_state_dict(checkpoint['decoder_optimizer'])\n",
    "# # epoch = checkpoint['epoch']\n",
    "\n",
    "encoder2.eval()\n",
    "decoder2.eval()\n",
    "# # encoder1.train()\n",
    "# # attn_decoder1.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# greedy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "search_method = 'greedy'\n",
    "print(test_model(encoder2, decoder2, val_loader))\n",
    "# print(test_model(encoder2, decoder2, test_loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Beam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "search_method = 'beam'\n",
    "for beam_size in range(2,15):\n",
    "    print(\"beam_size: \",beam_size)\n",
    "    sentence_comparison = test_model(encoder2, decoder2, val_loader,search_method)\n",
    "#     print(sentence_comparison)\n",
    "#     print(test_model(encoder2, decoder2, test_loader,search_method))\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_num = random.choice(range(len(sentence_comparison[1][0])))\n",
    "[i[random_num] for i in sentence_comparison[1]]"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "ch2en_LuongAttnDecoderRNN.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}