ch2en_lstm.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ch2en_lstm.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "id": "-C8HSFRaoGZo",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive/')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "deEZ1dYkoGZw",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "!pip3 install sacrebleu"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "5nbB0_wEoGZ0",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from os import path\n",
        "from wheel.pep425tags import get_abbr_impl, get_impl_ver, get_abi_tag\n",
        "platform = '{}{}-{}'.format(get_abbr_impl(), get_impl_ver(), get_abi_tag())\n",
        "\n",
        "accelerator = 'cu80' if path.exists('/opt/bin/nvidia-smi') else 'cpu'\n",
        "\n",
        "!pip install -q http://download.pytorch.org/whl/{accelerator}/torch-0.4.1-{platform}-linux_x86_64.whl torchvision\n",
        "import torch\n",
        "print(torch.__version__)\n",
        "print(torch.cuda.is_available())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "BkxJAAtaoGZ3",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "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"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "i78NZKtLoGZ6",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Pre-processing Data"
      ]
    },
    {
      "metadata": {
        "id": "KRHuEMYioGZ9",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "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"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "0-6BiXr1MaHT",
        "colab": {}
      },
      "cell_type": "code",
      "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"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "eq1zX9UtoGaD",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Variables"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "EYDlc9JJ0d5Z",
        "colab": {}
      },
      "cell_type": "code",
      "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 = 55000\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"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "OESV2zjg0d5c"
      },
      "cell_type": "markdown",
      "source": [
        "## Loading Data"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "X0G2eDBP0d5c",
        "colab": {}
      },
      "cell_type": "code",
      "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",
        "\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('感 应 电 机 可      以 加 热 东    西     ，     尤 其 擅 长 加 热 钢 铁')           "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "3JI2aRvToGaO",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def readLangs(lang1, lang2, data='train'):\n",
        "  \n",
        "    '''\n",
        "    Some codes are paraphrased from\n",
        "    https://pytorch.org/tutorials/intermediate/seq2seq_translation_tutorial.html\n",
        "    '''\n",
        "    \n",
        "    \n",
        "    #data: train/dev/test\n",
        "    print(\"Reading lines...\")\n",
        "    # Read the file and split into lines\n",
        "    zh_lines = open('/content/drive/My Drive/Colab Notebooks/Neural-Machine-Translation/iwslt-zh-en-sn/{}.tok.zh'.format(data)).read().split('\\n')\n",
        "    en_lines = open('/content/drive/My Drive/Colab Notebooks/Neural-Machine-Translation/iwslt-zh-en-sn/{}.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))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "I_PL2_xh0d5y",
        "colab": {}
      },
      "cell_type": "code",
      "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)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "FJv4sN4W0d52",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print(random.choice(val_pairs))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "vmMewXau0d6E"
      },
      "cell_type": "markdown",
      "source": [
        "# Preparing Data"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "SPnqKF_F0d6F",
        "colab": {}
      },
      "cell_type": "code",
      "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)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "PhwZoZz_0d6Q"
      },
      "cell_type": "markdown",
      "source": [
        "# Build Encoder-Decoder"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "E4lH5BjW10Jj",
        "colab": {}
      },
      "cell_type": "code",
      "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",
        "       \n",
        "        self.lstm = nn.LSTM(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, cell) = self.lstm(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),torch.zeros(2*self.n_layers, batch_size, self.hidden_size,device=device)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "EF4Fj5_pN9Lv",
        "colab": {}
      },
      "cell_type": "code",
      "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.lstm = nn.LSTM(self.hidden_size, self.hidden_size, n_layers,dropout=dropout_p)\n",
        "        \n",
        "        \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",
        "        print(\"#####\",hidden)\n",
        "        output, (hidden,cell) = self.lstm(output, hidden)\n",
        "        output = self.softmax(self.out(output[0]))\n",
        "        return output, hidden\n",
        "      \n",
        "    def initHidden(self):\n",
        "        return torch.zeros(2, batch_size, self.hidden_size, device=device)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "SbhB0-fFEd8Q",
        "colab": {}
      },
      "cell_type": "code",
      "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",
        "    "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "LqklCWyTEnpG",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class LuongAttnDecoderRNN(nn.Module):\n",
        "  \n",
        "  \n",
        "    \"\"\"\n",
        "    Some code is paraphrased from\n",
        "    https://github.com/spro/practical-pytorch/blob/master/seq2seq-translation/seq2seq-translation-batched.ipynb\n",
        "    \"\"\"\n",
        "    \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",
        "        \n",
        "        self.lstm = nn.LSTM(hidden_size, hidden_size, n_layers,dropout=dropout)\n",
        "        \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",
        "        #last_hidden=[last_hidden,last_hidden]\n",
        "        \n",
        "        rnn_output, (hidden,cell) = self.lstm(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"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "d1GxY0qw0d6V"
      },
      "cell_type": "markdown",
      "source": [
        "# Training Model"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "4ICvoIhH0d6W",
        "colab": {}
      },
      "cell_type": "code",
      "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))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "qiiPjqOy0d6Y",
        "colab": {}
      },
      "cell_type": "code",
      "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] \n",
        "    # 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",
        "                #print(decoder_hidden[0].size())\n",
        "                decoder_output, decoder_hidden, decoder_attention = decoder(\n",
        "                    decoder_input, [decoder_hidden,decoder_hidden], encoder_outputs)\n",
        "     \n",
        "            else:\n",
        "                decoder_output, decoder_hidden = decoder(\n",
        "                    decoder_input, [decoder_hidden,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,decoder_hidden], encoder_outputs)\n",
        "     \n",
        "            else:\n",
        "                decoder_output, decoder_hidden = decoder(\n",
        "                    decoder_input, [decoder_hidden,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()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "HmHvipVe0d6a",
        "colab": {}
      },
      "cell_type": "code",
      "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",
        "                                }, \"/content/drive/My Drive/Colab Notebooks/saved_model/aaa\"\\\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",
        "#                             }, \"/content/drive/My Drive/Colab Notebooks/saved_scores/attnIs{}_hiddenSize{}_nLayer{}_batchSize{}_epoch{}_srcVocSize{}_tgtVocSize{}_lrDecay{}\"\\\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)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "mj8OwEU_0d6d"
      },
      "cell_type": "markdown",
      "source": [
        "# Plotting results"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "bDqy1HHH0d6e",
        "colab": {}
      },
      "cell_type": "code",
      "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)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "L_-Y9XII0d6g"
      },
      "cell_type": "markdown",
      "source": [
        "# Evaluation"
      ]
    },
    {
      "metadata": {
        "id": "5M29sEy0oGbE",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "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(),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],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"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "muc68_Oz0d6h",
        "colab": {}
      },
      "cell_type": "code",
      "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,decoder_hidden], encoder_outputs)\n",
        "            else:\n",
        "                decoder_output, decoder_hidden = decoder(\n",
        "                    decoder_input, [decoder_hidden,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()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "2CjcWcuxoGbO",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "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)\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",
        "        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"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "ce-QyODD0d6l",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def test_model(encoder, decoder, loader, search_method = 'greedy'):\n",
        "  \n",
        "    \"\"\"\n",
        "    Thanks for the sacrebleu.py in \n",
        "    github https://github.com/mjpost/sacreBLEU/blob/master/sacrebleu.py#L1022-L1080\n",
        "    \"\"\"\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",
        "    \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"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "hjafDuaU0d6s"
      },
      "cell_type": "markdown",
      "source": [
        "# TRAINING AND EVALUATING"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "T3iWCiIQ0d6t",
        "scrolled": true,
        "colab": {}
      },
      "cell_type": "code",
      "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)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "SgZYXF1GoGbX",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Resume Model and Test"
      ]
    },
    {
      "metadata": {
        "id": "SYg-uwjAoGbY",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "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('/content/drive/My Drive/Colab Notebooks/saved_model/aaa')\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()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "su3ahIn-oGba",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# greedy"
      ]
    },
    {
      "metadata": {
        "id": "JjN-_ikuoGbb",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "search_method = 'greedy'\n",
        "print(test_model(encoder2, decoder2, val_loader))\n",
        "# print(test_model(encoder2, decoder2, test_loader))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "iT0c0n_UoGbe",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Beam"
      ]
    },
    {
      "metadata": {
        "id": "gTZ_51LeoGbf",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "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()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "SNdz2O6GoGbk",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "random_num = random.choice(range(len(sentence_comparison[1][0])))\n",
        "[i[random_num] for i in sentence_comparison[1]]"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}