seq2seq.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Batched Seq2Seq Example\n",
    "Based on the [`seq2seq-translation-batched.ipynb`](https://github.com/spro/practical-pytorch/blob/master/seq2seq-translation/seq2seq-translation-batched.ipynb) from *practical-pytorch*, but more extra features.\n",
    "\n",
    "This example runs grammatical error correction task where the source sequence is a grammatically erroneuous English sentence and the target sequence is an grammatically correct English sentence. The corpus and evaluation script can be download at: https://github.com/keisks/jfleg.\n",
    "\n",
    "### Extra features\n",
    "- Cleaner codebase\n",
    "- Very detailed comments for learners\n",
    "- Implement Pytorch native dataset and dataloader for batching\n",
    "- Correctly handle the hidden state from bidirectional encoder and past to the decoder as initial hidden state.\n",
    "- Fully batched attention mechanism computation (only implement `general attention` but it's sufficient). Note: The original code still uses for-loop to compute, which is very slow.\n",
    "- Support LSTM instead of only GRU\n",
    "- Shared embeddings (encoder's input embedding and decoder's input embedding)\n",
    "- Pretrained Glove embedding\n",
    "- Fixed embedding\n",
    "- Tie embeddings (decoder's input embedding and decoder's output embedding)\n",
    "- Tensorboard visualization\n",
    "- Load and save checkpoint\n",
    "- Replace unknown words by selecting the source token with the highest attention score. (Translation)\n",
    "\n",
    "### Cons\n",
    "Comparing to the state-of-the-art seq2seq library, OpenNMT-py, there are some stuffs that aren't optimized in this codebase:\n",
    "- Use CuDNN when possible (always on encoder, on decoder when input_feed 0)\n",
    "- Always avoid indexing / loops and use torch primitives.\n",
    "- When possible, batch softmax operations across time. ( this is the second complicated part of the code)\n",
    "- Batch inference and beam search for translation (this is the most complicated part of the code)\n",
    "\n",
    "Thanks to the author of OpenNMT-py @srush for answering the questions for me! See https://github.com/OpenNMT/OpenNMT-py/issues/552"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "import codecs\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Variable\n",
    "from torch import optim\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "\"\"\" Please download from here: \n",
    "1. Install spacy: https://spacy.io/usage/\n",
    "2. Install model: https://spacy.io/usage/models\n",
    "Recommend to install spacy since it is a very powerful NLP tool\n",
    "\"\"\"\n",
    "\n",
    "import spacy\n",
    "nlp = spacy.load('en_core_web_lg') # For the glove embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Use_CUDA=True\n",
      "current_device=0\n"
     ]
    }
   ],
   "source": [
    "\"\"\" Enable GPU training \"\"\"\n",
    "USE_CUDA = torch.cuda.is_available()\n",
    "print('Use_CUDA={}'.format(USE_CUDA))\n",
    "if USE_CUDA:\n",
    "    # You can change device by `torch.cuda.set_device(device_id)`\n",
    "    print('current_device={}'.format(torch.cuda.current_device()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build vocabulary, dataset and data loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import codecs\n",
    "from tqdm import tqdm\n",
    "from collections import Counter, namedtuple\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "PAD = 0\n",
    "BOS = 1\n",
    "EOS = 2\n",
    "UNK = 3\n",
    "\n",
    "class AttrDict(dict):\n",
    "    \"\"\" Access dictionary keys like attribute \n",
    "        https://stackoverflow.com/questions/4984647/accessing-dict-keys-like-an-attribute\n",
    "    \"\"\"\n",
    "    def __init__(self, *av, **kav):\n",
    "        dict.__init__(self, *av, **kav)\n",
    "        self.__dict__ = self\n",
    "\n",
    "class NMTDataset(Dataset):\n",
    "    def __init__(self, src_path, tgt_path, src_vocab=None, tgt_vocab=None, max_vocab_size=50000, share_vocab=True):\n",
    "        \"\"\" Note: If src_vocab, tgt_vocab is not given, it will build both vocabs.\n",
    "            Args: \n",
    "            - src_path, tgt_path: text file with tokenized sentences.\n",
    "            - src_vocab, tgt_vocab: data structure is same as self.build_vocab().\n",
    "        \"\"\"\n",
    "        print('='*100)\n",
    "        print('Dataset preprocessing log:')\n",
    "        \n",
    "        print('- Loading and tokenizing source sentences...')\n",
    "        self.src_sents = self.load_sents(src_path)\n",
    "        print('- Loading and tokenizing target sentences...')\n",
    "        self.tgt_sents = self.load_sents(tgt_path)\n",
    "        \n",
    "        if src_vocab is None or tgt_vocab is None:\n",
    "            print('- Building source counter...')\n",
    "            self.src_counter = self.build_counter(self.src_sents)\n",
    "            print('- Building target counter...')\n",
    "            self.tgt_counter = self.build_counter(self.tgt_sents)\n",
    "\n",
    "            if share_vocab:\n",
    "                print('- Building source vocabulary...')\n",
    "                self.src_vocab = self.build_vocab(self.src_counter + self.tgt_counter, max_vocab_size)\n",
    "                print('- Building target vocabulary...')\n",
    "                self.tgt_vocab = self.src_vocab\n",
    "            else:\n",
    "                print('- Building source vocabulary...')\n",
    "                self.src_vocab = self.build_vocab(self.src_counter, max_vocab_size)\n",
    "                print('- Building target vocabulary...')\n",
    "                self.tgt_vocab = self.build_vocab(self.tgt_counter, max_vocab_size)\n",
    "        else:\n",
    "            self.src_vocab = src_vocab\n",
    "            self.tgt_vocab = tgt_vocab\n",
    "            share_vocab = src_vocab == tgt_vocab\n",
    "                        \n",
    "        print('='*100)\n",
    "        print('Dataset Info:')\n",
    "        print('- Number of source sentences: {}'.format(len(self.src_sents)))\n",
    "        print('- Number of target sentences: {}'.format(len(self.tgt_sents)))\n",
    "        print('- Source vocabulary size: {}'.format(len(self.src_vocab.token2id)))\n",
    "        print('- Target vocabulary size: {}'.format(len(self.tgt_vocab.token2id)))\n",
    "        print('- Shared vocabulary: {}'.format(share_vocab))\n",
    "        print('='*100 + '\\n')\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.src_sents)\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        src_sent = self.src_sents[index]\n",
    "        tgt_sent = self.tgt_sents[index]\n",
    "        src_seq = self.tokens2ids(src_sent, self.src_vocab.token2id, append_BOS=False, append_EOS=True)\n",
    "        tgt_seq = self.tokens2ids(tgt_sent, self.tgt_vocab.token2id, append_BOS=False, append_EOS=True)\n",
    "\n",
    "        return src_sent, tgt_sent, src_seq, tgt_seq\n",
    "    \n",
    "    def load_sents(self, file_path):\n",
    "        sents = []\n",
    "        with codecs.open(file_path) as file:\n",
    "            for sent in tqdm(file.readlines()):\n",
    "                tokens = [token for token in sent.split()]\n",
    "                sents.append(tokens)\n",
    "        return sents\n",
    "    \n",
    "    def build_counter(self, sents):\n",
    "        counter = Counter()\n",
    "        for sent in tqdm(sents):\n",
    "            counter.update(sent)\n",
    "        return counter\n",
    "    \n",
    "    def build_vocab(self, counter, max_vocab_size):\n",
    "        vocab = AttrDict()\n",
    "        vocab.token2id = {'<PAD>': PAD, '<BOS>': BOS, '<EOS>': EOS, '<UNK>': UNK}\n",
    "        vocab.token2id.update({token: _id+4 for _id, (token, count) in tqdm(enumerate(counter.most_common(max_vocab_size)))})\n",
    "        vocab.id2token = {v:k for k,v in tqdm(vocab.token2id.items())}    \n",
    "        return vocab\n",
    "    \n",
    "    def tokens2ids(self, tokens, token2id, append_BOS=True, append_EOS=True):\n",
    "        seq = []\n",
    "        if append_BOS: seq.append(BOS)\n",
    "        seq.extend([token2id.get(token, UNK) for token in tokens])\n",
    "        if append_EOS: seq.append(EOS)\n",
    "        return seq\n",
    "    \n",
    "def collate_fn(data):\n",
    "    \"\"\"\n",
    "    Creates mini-batch tensors from (src_sent, tgt_sent, src_seq, tgt_seq).\n",
    "    We should build a custom collate_fn rather than using default collate_fn,\n",
    "    because merging sequences (including padding) is not supported in default.\n",
    "    Seqeuences are padded to the maximum length of mini-batch sequences (dynamic padding).\n",
    "    \n",
    "    Args:\n",
    "        data: list of tuple (src_sents, tgt_sents, src_seqs, tgt_seqs)\n",
    "        - src_sents, tgt_sents: batch of original tokenized sentences\n",
    "        - src_seqs, tgt_seqs: batch of original tokenized sentence ids\n",
    "    Returns:\n",
    "        - src_sents, tgt_sents (tuple): batch of original tokenized sentences\n",
    "        - src_seqs, tgt_seqs (variable): (max_src_len, batch_size)\n",
    "        - src_lens, tgt_lens (tensor): (batch_size)\n",
    "       \n",
    "    \"\"\"\n",
    "    def _pad_sequences(seqs):\n",
    "        lens = [len(seq) for seq in seqs]\n",
    "        padded_seqs = torch.zeros(len(seqs), max(lens)).long()\n",
    "        for i, seq in enumerate(seqs):\n",
    "            end = lens[i]\n",
    "            padded_seqs[i, :end] = torch.LongTensor(seq[:end])\n",
    "        return padded_seqs, lens\n",
    "\n",
    "    # Sort a list by *source* sequence length (descending order) to use `pack_padded_sequence`.\n",
    "    # The *target* sequence is not sorted <-- It's ok, cause `pack_padded_sequence` only takes\n",
    "    # *source* sequence, which is in the EncoderRNN\n",
    "    data.sort(key=lambda x: len(x[0]), reverse=True)\n",
    "\n",
    "    # Seperate source and target sequences.\n",
    "    src_sents, tgt_sents, src_seqs, tgt_seqs = zip(*data)\n",
    "    \n",
    "    # Merge sequences (from tuple of 1D tensor to 2D tensor)\n",
    "    src_seqs, src_lens = _pad_sequences(src_seqs)\n",
    "    tgt_seqs, tgt_lens = _pad_sequences(tgt_seqs)\n",
    "    \n",
    "    # (batch, seq_len) => (seq_len, batch)\n",
    "    src_seqs = src_seqs.transpose(0,1)\n",
    "    tgt_seqs = tgt_seqs.transpose(0,1)\n",
    "\n",
    "    return src_sents, tgt_sents, src_seqs, tgt_seqs, src_lens, tgt_lens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build models\n",
    "### Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class EncoderRNN(nn.Module):\n",
    "    def __init__(self, embedding=None, rnn_type='LSTM', hidden_size=128, num_layers=1, dropout=0.3, bidirectional=True):\n",
    "        super(EncoderRNN, self).__init__()\n",
    "        \n",
    "        self.num_layers = num_layers\n",
    "        self.dropout = dropout\n",
    "        self.bidirectional = bidirectional\n",
    "        self.num_directions = 2 if bidirectional else 1\n",
    "        self.hidden_size = hidden_size // self.num_directions\n",
    "        \n",
    "        self.embedding = embedding\n",
    "        self.word_vec_size = self.embedding.embedding_dim\n",
    "        \n",
    "        self.rnn_type = rnn_type\n",
    "        self.rnn = getattr(nn, self.rnn_type)(\n",
    "                           input_size=self.word_vec_size,\n",
    "                           hidden_size=self.hidden_size,\n",
    "                           num_layers=self.num_layers,\n",
    "                           dropout=self.dropout, \n",
    "                           bidirectional=self.bidirectional)\n",
    "        \n",
    "    def forward(self, src_seqs, src_lens, hidden=None):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            - src_seqs: (max_src_len, batch_size)\n",
    "            - src_lens: (batch_size)\n",
    "        Returns:\n",
    "            - outputs: (max_src_len, batch_size, hidden_size * num_directions)\n",
    "            - hidden : (num_layers, batch_size, hidden_size * num_directions)\n",
    "        \"\"\"\n",
    "        \n",
    "        # (max_src_len, batch_size) => (max_src_len, batch_size, word_vec_size)\n",
    "        emb = self.embedding(src_seqs)\n",
    "\n",
    "        # packed_emb:\n",
    "        # - data: (sum(batch_sizes), word_vec_size)\n",
    "        # - batch_sizes: list of batch sizes\n",
    "        packed_emb = nn.utils.rnn.pack_padded_sequence(emb, src_lens)\n",
    "\n",
    "        # rnn(gru) returns:\n",
    "        # - packed_outputs: shape same as packed_emb\n",
    "        # - hidden: (num_layers * num_directions, batch_size, hidden_size) \n",
    "        packed_outputs, hidden = self.rnn(packed_emb, hidden)\n",
    "\n",
    "        # outputs: (max_src_len, batch_size, hidden_size * num_directions)\n",
    "        # output_lens == src_lensˇ\n",
    "        outputs, output_lens =  nn.utils.rnn.pad_packed_sequence(packed_outputs)\n",
    "        \n",
    "        if self.bidirectional:\n",
    "            # (num_layers * num_directions, batch_size, hidden_size) \n",
    "            # => (num_layers, batch_size, hidden_size * num_directions)\n",
    "            hidden = self._cat_directions(hidden)\n",
    "        \n",
    "        return outputs, hidden\n",
    "    \n",
    "    def _cat_directions(self, hidden):\n",
    "        \"\"\" If the encoder is bidirectional, do the following transformation.\n",
    "            Ref: https://github.com/IBM/pytorch-seq2seq/blob/master/seq2seq/models/DecoderRNN.py#L176\n",
    "            -----------------------------------------------------------\n",
    "            In: (num_layers * num_directions, batch_size, hidden_size)\n",
    "            (ex: num_layers=2, num_directions=2)\n",
    "\n",
    "            layer 1: forward__hidden(1)\n",
    "            layer 1: backward_hidden(1)\n",
    "            layer 2: forward__hidden(2)\n",
    "            layer 2: backward_hidden(2)\n",
    "\n",
    "            -----------------------------------------------------------\n",
    "            Out: (num_layers, batch_size, hidden_size * num_directions)\n",
    "\n",
    "            layer 1: forward__hidden(1) backward_hidden(1)\n",
    "            layer 2: forward__hidden(2) backward_hidden(2)\n",
    "        \"\"\"\n",
    "        def _cat(h):\n",
    "            return torch.cat([h[0:h.size(0):2], h[1:h.size(0):2]], 2)\n",
    "            \n",
    "        if isinstance(hidden, tuple):\n",
    "            # LSTM hidden contains a tuple (hidden state, cell state)\n",
    "            hidden = tuple([_cat(h) for h in hidden])\n",
    "        else:\n",
    "            # GRU hidden\n",
    "            hidden = _cat(hidden)\n",
    "            \n",
    "        return hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decoder with \"general attention\" mechanism"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class LuongAttnDecoderRNN(nn.Module):\n",
    "    def __init__(self, encoder, embedding=None, attention=True, bias=True, tie_embeddings=False, dropout=0.3):\n",
    "        \"\"\" General attention in `Effective Approaches to Attention-based Neural Machine Translation`\n",
    "            Ref: https://arxiv.org/abs/1508.04025\n",
    "            \n",
    "            Share input and output embeddings:\n",
    "            Ref:\n",
    "                - \"Using the Output Embedding to Improve Language Models\" (Press & Wolf 2016)\n",
    "                   https://arxiv.org/abs/1608.05859\n",
    "                - \"Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling\" (Inan et al. 2016)\n",
    "                   https://arxiv.org/abs/1611.01462\n",
    "        \"\"\"\n",
    "        super(LuongAttnDecoderRNN, self).__init__()\n",
    "        \n",
    "        self.hidden_size = encoder.hidden_size * encoder.num_directions\n",
    "        self.num_layers = encoder.num_layers\n",
    "        self.dropout = dropout\n",
    "        self.embedding = embedding\n",
    "        self.attention = attention\n",
    "        self.tie_embeddings = tie_embeddings\n",
    "        \n",
    "        self.vocab_size = self.embedding.num_embeddings\n",
    "        self.word_vec_size = self.embedding.embedding_dim\n",
    "        \n",
    "        self.rnn_type = encoder.rnn_type\n",
    "        self.rnn = getattr(nn, self.rnn_type)(\n",
    "                            input_size=self.word_vec_size,\n",
    "                            hidden_size=self.hidden_size,\n",
    "                            num_layers=self.num_layers,\n",
    "                            dropout=self.dropout)\n",
    "        \n",
    "        if self.attention:\n",
    "            self.W_a = nn.Linear(encoder.hidden_size * encoder.num_directions,\n",
    "                                 self.hidden_size, bias=bias)\n",
    "            self.W_c = nn.Linear(encoder.hidden_size * encoder.num_directions + self.hidden_size, \n",
    "                                 self.hidden_size, bias=bias)\n",
    "        \n",
    "        if self.tie_embeddings:\n",
    "            self.W_proj = nn.Linear(self.hidden_size, self.word_vec_size, bias=bias)\n",
    "            self.W_s = nn.Linear(self.word_vec_size, self.vocab_size, bias=bias)\n",
    "            self.W_s.weight = self.embedding.weight\n",
    "        else:\n",
    "            self.W_s = nn.Linear(self.hidden_size, self.vocab_size, bias=bias)\n",
    "        \n",
    "    def forward(self, input_seq, decoder_hidden, encoder_outputs, src_lens):\n",
    "        \"\"\" Args:\n",
    "            - input_seq      : (batch_size)\n",
    "            - decoder_hidden : (t=0) last encoder hidden state (num_layers * num_directions, batch_size, hidden_size) \n",
    "                               (t>0) previous decoder hidden state (num_layers, batch_size, hidden_size)\n",
    "            - encoder_outputs: (max_src_len, batch_size, hidden_size * num_directions)\n",
    "        \n",
    "            Returns:\n",
    "            - output           : (batch_size, vocab_size)\n",
    "            - decoder_hidden   : (num_layers, batch_size, hidden_size)\n",
    "            - attention_weights: (batch_size, max_src_len)\n",
    "        \"\"\"        \n",
    "        # (batch_size) => (seq_len=1, batch_size)\n",
    "        input_seq = input_seq.unsqueeze(0)\n",
    "        \n",
    "        # (seq_len=1, batch_size) => (seq_len=1, batch_size, word_vec_size) \n",
    "        emb = self.embedding(input_seq)\n",
    "        \n",
    "        # rnn returns:\n",
    "        # - decoder_output: (seq_len=1, batch_size, hidden_size)\n",
    "        # - decoder_hidden: (num_layers, batch_size, hidden_size)\n",
    "        decoder_output, decoder_hidden = self.rnn(emb, decoder_hidden)\n",
    "\n",
    "        # (seq_len=1, batch_size, hidden_size) => (batch_size, seq_len=1, hidden_size)\n",
    "        decoder_output = decoder_output.transpose(0,1)\n",
    "        \n",
    "        \"\"\" \n",
    "        ------------------------------------------------------------------------------------------\n",
    "        Notes of computing attention scores\n",
    "        ------------------------------------------------------------------------------------------\n",
    "        # For-loop version:\n",
    "\n",
    "        max_src_len = encoder_outputs.size(0)\n",
    "        batch_size = encoder_outputs.size(1)\n",
    "        attention_scores = Variable(torch.zeros(batch_size, max_src_len))\n",
    "\n",
    "        # For every batch, every time step of encoder's hidden state, calculate attention score.\n",
    "        for b in range(batch_size):\n",
    "            for t in range(max_src_len):\n",
    "                # Loung. eq(8) -- general form content-based attention:\n",
    "                attention_scores[b,t] = decoder_output[b].dot(attention.W_a(encoder_outputs[t,b]))\n",
    "\n",
    "        ------------------------------------------------------------------------------------------\n",
    "        # Vectorized version:\n",
    "\n",
    "        1. decoder_output: (batch_size, seq_len=1, hidden_size)\n",
    "        2. encoder_outputs: (max_src_len, batch_size, hidden_size * num_directions)\n",
    "        3. W_a(encoder_outputs): (max_src_len, batch_size, hidden_size)\n",
    "                        .transpose(0,1)  : (batch_size, max_src_len, hidden_size) \n",
    "                        .transpose(1,2)  : (batch_size, hidden_size, max_src_len)\n",
    "        4. attention_scores: \n",
    "                        (batch_size, seq_len=1, hidden_size) * (batch_size, hidden_size, max_src_len) \n",
    "                        => (batch_size, seq_len=1, max_src_len)\n",
    "        \"\"\"\n",
    "        \n",
    "        if self.attention:\n",
    "            # attention_scores: (batch_size, seq_len=1, max_src_len)\n",
    "            attention_scores = torch.bmm(decoder_output, self.W_a(encoder_outputs).transpose(0,1).transpose(1,2))\n",
    "\n",
    "            # attention_mask: (batch_size, seq_len=1, max_src_len)\n",
    "            attention_mask = sequence_mask(src_lens).unsqueeze(1)\n",
    "\n",
    "            # Fills elements of tensor with `-float('inf')` where `mask` is 1.\n",
    "            attention_scores.data.masked_fill_(1 - attention_mask.data, -float('inf'))\n",
    "\n",
    "            # attention_weights: (batch_size, seq_len=1, max_src_len) => (batch_size, max_src_len) for `F.softmax` \n",
    "            # => (batch_size, seq_len=1, max_src_len)\n",
    "            try: # torch 0.3.x\n",
    "                attention_weights = F.softmax(attention_scores.squeeze(1), dim=1).unsqueeze(1)\n",
    "            except:\n",
    "                attention_weights = F.softmax(attention_scores.squeeze(1)).unsqueeze(1)\n",
    "\n",
    "            # context_vector:\n",
    "            # (batch_size, seq_len=1, max_src_len) * (batch_size, max_src_len, encoder_hidden_size * num_directions)\n",
    "            # => (batch_size, seq_len=1, encoder_hidden_size * num_directions)\n",
    "            context_vector = torch.bmm(attention_weights, encoder_outputs.transpose(0,1))\n",
    "\n",
    "            # concat_input: (batch_size, seq_len=1, encoder_hidden_size * num_directions + decoder_hidden_size)\n",
    "            concat_input = torch.cat([context_vector, decoder_output], -1)\n",
    "\n",
    "            # (batch_size, seq_len=1, encoder_hidden_size * num_directions + decoder_hidden_size) => (batch_size, seq_len=1, decoder_hidden_size)\n",
    "            concat_output = F.tanh(self.W_c(concat_input))\n",
    "            \n",
    "            # Prepare returns:\n",
    "            # (batch_size, seq_len=1, max_src_len) => (batch_size, max_src_len)\n",
    "            attention_weights = attention_weights.squeeze(1)\n",
    "        else:\n",
    "            attention_weights = None\n",
    "            concat_output = decoder_output\n",
    "        \n",
    "        # If input and output embeddings are tied,\n",
    "        # project `decoder_hidden_size` to `word_vec_size`.\n",
    "        if self.tie_embeddings:\n",
    "            output = self.W_s(self.W_proj(concat_output))\n",
    "        else:\n",
    "            # (batch_size, seq_len=1, decoder_hidden_size) => (batch_size, seq_len=1, vocab_size)\n",
    "            output = self.W_s(concat_output)    \n",
    "        \n",
    "        # Prepare returns:\n",
    "        # (batch_size, seq_len=1, vocab_size) => (batch_size, vocab_size)\n",
    "        output = output.squeeze(1)\n",
    "        \n",
    "        del src_lens\n",
    "        \n",
    "        return output, decoder_hidden, attention_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def load_spacy_glove_embedding(spacy_nlp, vocab):\n",
    "    \n",
    "    vocab_size = len(vocab.token2id)\n",
    "    word_vec_size = spacy_nlp.vocab.vectors_length\n",
    "    embedding = np.zeros((vocab_size, word_vec_size))\n",
    "    unk_count = 0\n",
    "    \n",
    "    print('='*100)\n",
    "    print('Loading spacy glove embedding:')\n",
    "    print('- Vocabulary size: {}'.format(vocab_size))\n",
    "    print('- Word vector size: {}'.format(word_vec_size))\n",
    "    \n",
    "    for token, index in tqdm(vocab.token2id.items()):\n",
    "        if token == vocab.id2token[PAD]: \n",
    "            continue\n",
    "        elif token in [vocab.id2token[BOS], vocab.id2token[EOS], vocab.id2token[UNK]]: \n",
    "            vector = np.random.rand(word_vec_size,)\n",
    "        elif spacy_nlp.vocab[token].has_vector: \n",
    "            vector = spacy_nlp.vocab[token].vector\n",
    "        else:\n",
    "            vector = embedding[UNK] \n",
    "            unk_count += 1\n",
    "            \n",
    "        embedding[index] = vector\n",
    "        \n",
    "    print('- Unknown word count: {}'.format(unk_count))\n",
    "    print('='*100 + '\\n')\n",
    "        \n",
    "    return torch.from_numpy(embedding).float()\n",
    "\n",
    "def sequence_mask(sequence_length, max_len=None):\n",
    "    \"\"\"\n",
    "    Caution: Input and Return are VARIABLE.\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)\n",
    "    seq_range_expand = Variable(seq_range_expand)\n",
    "    if sequence_length.is_cuda:\n",
    "        seq_range_expand = seq_range_expand.cuda()\n",
    "    seq_length_expand = (sequence_length.unsqueeze(1)\n",
    "                         .expand_as(seq_range_expand))\n",
    "    mask = seq_range_expand < seq_length_expand\n",
    "    \n",
    "    return mask\n",
    "\n",
    "def masked_cross_entropy(logits, target, length):\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",
    "    Returns:\n",
    "        loss: An average loss value masked by the length.\n",
    "        \n",
    "    The code is same as:\n",
    "    \n",
    "    weight = torch.ones(tgt_vocab_size)\n",
    "    weight[padding_idx] = 0\n",
    "    criterion = nn.CrossEntropyLoss(weight.cuda(), size_average)\n",
    "    loss = criterion(logits_flat, losses_flat)\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 = F.log_softmax(logits_flat)\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",
    "    # Note: mask need to bed casted to float!\n",
    "    losses = losses * mask.float()\n",
    "    loss = losses.sum() / mask.float().sum()\n",
    "    \n",
    "    # (batch_size * max_tgt_len,)\n",
    "    pred_flat = log_probs_flat.max(1)[1]\n",
    "    # (batch_size * max_tgt_len,) => (batch_size, max_tgt_len) => (max_tgt_len, batch_size)\n",
    "    pred_seqs = pred_flat.view(*target.size()).transpose(0,1).contiguous()\n",
    "    # (batch_size, max_len) => (batch_size * max_tgt_len,)\n",
    "    mask_flat = mask.view(-1)\n",
    "    \n",
    "    # `.float()` IS VERY IMPORTANT !!!\n",
    "    # https://discuss.pytorch.org/t/batch-size-and-validation-accuracy/4066/3\n",
    "    num_corrects = int(pred_flat.eq(target_flat.squeeze(1)).masked_select(mask_flat).float().data.sum())\n",
    "    num_words = length.data.sum()\n",
    "\n",
    "    return loss, pred_seqs, num_corrects, num_words\n",
    "\n",
    "def load_checkpoint(checkpoint_path):\n",
    "    # It's weird that if `map_location` is not given, it will be extremely slow.\n",
    "    return torch.load(checkpoint_path, map_location=lambda storage, loc: storage)\n",
    "\n",
    "def save_checkpoint(opts, experiment_name, encoder, decoder, encoder_optim, decoder_optim,\n",
    "                    total_accuracy, total_loss, global_step):\n",
    "    checkpoint = {\n",
    "        'opts': opts,\n",
    "        'global_step': global_step,\n",
    "        'encoder_state_dict': encoder.state_dict(),\n",
    "        'decoder_state_dict': decoder.state_dict(),\n",
    "        'encoder_optim_state_dict': encoder_optim.state_dict(),\n",
    "        'decoder_optim_state_dict': decoder_optim.state_dict()\n",
    "    }\n",
    "    \n",
    "    checkpoint_path = 'checkpoints/%s_acc_%.2f_loss_%.2f_step_%d.pt' % (experiment_name, total_accuracy, total_loss, global_step)\n",
    "    \n",
    "    directory, filename = os.path.split(os.path.abspath(checkpoint_path))\n",
    "\n",
    "    if not os.path.exists(directory):\n",
    "        os.makedirs(directory)\n",
    "    \n",
    "    torch.save(checkpoint, checkpoint_path)\n",
    "    \n",
    "    return checkpoint_path\n",
    "\n",
    "def variable2numpy(var):\n",
    "    \"\"\" For tensorboard visualization \"\"\"\n",
    "    return var.data.cpu().numpy()\n",
    "\n",
    "def write_to_tensorboard(writer, global_step, total_loss, total_corrects, total_words, total_accuracy,\n",
    "                         encoder_grad_norm, decoder_grad_norm, clipped_encoder_grad_norm, clipped_decoder_grad_norm,\n",
    "                         encoder, decoder, gpu_memory_usage=None):\n",
    "    # scalars\n",
    "    if gpu_memory_usage is not None:\n",
    "        writer.add_scalar('curr_gpu_memory_usage', gpu_memory_usage['curr'], global_step)\n",
    "        writer.add_scalar('diff_gpu_memory_usage', gpu_memory_usage['diff'], global_step)\n",
    "        \n",
    "    writer.add_scalar('total_loss', total_loss, global_step)\n",
    "    writer.add_scalar('total_accuracy', total_accuracy, global_step)\n",
    "    writer.add_scalar('total_corrects', total_corrects, global_step)\n",
    "    writer.add_scalar('total_words', total_words, global_step)\n",
    "    writer.add_scalar('encoder_grad_norm', encoder_grad_norm, global_step)\n",
    "    writer.add_scalar('decoder_grad_norm', decoder_grad_norm, global_step)\n",
    "    writer.add_scalar('clipped_encoder_grad_norm', clipped_encoder_grad_norm, global_step)\n",
    "    writer.add_scalar('clipped_decoder_grad_norm', clipped_decoder_grad_norm, global_step)\n",
    "    \n",
    "    # histogram\n",
    "    for name, param in encoder.named_parameters():\n",
    "        name = name.replace('.', '/')\n",
    "        writer.add_histogram('encoder/{}'.format(name), variable2numpy(param), global_step, bins='doane')\n",
    "        if param.grad is not None:\n",
    "            writer.add_histogram('encoder/{}/grad'.format(name), variable2numpy(param.grad), global_step, bins='doane')\n",
    "\n",
    "    for name, param in decoder.named_parameters():\n",
    "        name = name.replace('.', '/')\n",
    "        writer.add_histogram('decoder/{}'.format(name), variable2numpy(param), global_step, bins='doane')\n",
    "        if param.grad is not None:\n",
    "            writer.add_histogram('decoder/{}/grad'.format(name), variable2numpy(param.grad), global_step, bins='doane')\n",
    "            \n",
    "def detach_hidden(hidden):\n",
    "    \"\"\" Wraps hidden states in new Variables, to detach them from their history. Prevent OOM.\n",
    "        After detach, the hidden's requires_grad=Fasle and grad_fn=None.\n",
    "    Issues:\n",
    "    - Memory leak problem in LSTM and RNN: https://github.com/pytorch/pytorch/issues/2198\n",
    "    - https://github.com/pytorch/examples/blob/master/word_language_model/main.py\n",
    "    - https://discuss.pytorch.org/t/help-clarifying-repackage-hidden-in-word-language-model/226\n",
    "    - https://discuss.pytorch.org/t/solved-why-we-need-to-detach-variable-which-contains-hidden-representation/1426\n",
    "    - \n",
    "    \"\"\"\n",
    "    if type(hidden) == Variable:\n",
    "        hidden.detach_() # same as creating a new variable.\n",
    "    else:\n",
    "        for h in hidden: h.detach_()\n",
    "\n",
    "def get_gpu_memory_usage(device_id):\n",
    "    \"\"\"Get the current gpu usage. \"\"\"\n",
    "    result = subprocess.check_output(\n",
    "        [\n",
    "            'nvidia-smi', '--query-gpu=memory.used',\n",
    "            '--format=csv,nounits,noheader'\n",
    "        ], encoding='utf-8')\n",
    "    # Convert lines into a dictionary\n",
    "    gpu_memory = [int(x) for x in result.strip().split('\\n')]\n",
    "    gpu_memory_map = dict(zip(range(len(gpu_memory)), gpu_memory))\n",
    "    return gpu_memory_map[device_id]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def compute_grad_norm(parameters, norm_type=2):\n",
    "    \"\"\" Ref: http://pytorch.org/docs/0.3.0/_modules/torch/nn/utils/clip_grad.html#clip_grad_norm\n",
    "    \"\"\"\n",
    "    parameters = list(filter(lambda p: p.grad is not None, parameters))\n",
    "    norm_type = float(norm_type)\n",
    "    if norm_type == float('inf'):\n",
    "        total_norm = max(p.grad.data.abs().max() for p in parameters)\n",
    "    else:\n",
    "        total_norm = 0\n",
    "        for p in parameters:\n",
    "            param_norm = p.grad.data.norm(norm_type)\n",
    "            total_norm += param_norm ** norm_type\n",
    "        total_norm = total_norm ** (1. / norm_type)\n",
    "    return total_norm\n",
    "\n",
    "def train(src_sents, tgt_sents, src_seqs, tgt_seqs, src_lens, tgt_lens,\n",
    "          encoder, decoder, encoder_optim, decoder_optim, opts):    \n",
    "    # -------------------------------------\n",
    "    # Prepare input and output placeholders\n",
    "    # -------------------------------------\n",
    "    # Last batch might not have the same size as we set to the `batch_size`\n",
    "    batch_size = src_seqs.size(1)\n",
    "    assert(batch_size == tgt_seqs.size(1))\n",
    "    \n",
    "    # Pack tensors to variables for neural network inputs (in order to autograd)\n",
    "    src_seqs = Variable(src_seqs)\n",
    "    tgt_seqs = Variable(tgt_seqs)\n",
    "    src_lens = Variable(torch.LongTensor(src_lens))\n",
    "    tgt_lens = Variable(torch.LongTensor(tgt_lens))\n",
    "\n",
    "    # Decoder's input\n",
    "    input_seq = Variable(torch.LongTensor([BOS] * batch_size))\n",
    "    \n",
    "    # Decoder's output sequence length = max target sequence length of current batch.\n",
    "    max_tgt_len = tgt_lens.data.max()\n",
    "    \n",
    "    # Store all decoder's outputs.\n",
    "    # **CRUTIAL** \n",
    "    # Don't set:\n",
    "    # >> decoder_outputs = Variable(torch.zeros(max_tgt_len, batch_size, decoder.vocab_size))\n",
    "    # Varying tensor size could cause GPU allocate a new memory causing OOM, \n",
    "    # so we intialize tensor with fixed size instead:\n",
    "    # `opts.max_seq_len` is a fixed number, unlike `max_tgt_len` always varys.\n",
    "    decoder_outputs = Variable(torch.zeros(opts.max_seq_len, batch_size, decoder.vocab_size))\n",
    "\n",
    "    # Move variables from CPU to GPU.\n",
    "    if USE_CUDA:\n",
    "        src_seqs = src_seqs.cuda()\n",
    "        tgt_seqs = tgt_seqs.cuda()\n",
    "        src_lens = src_lens.cuda()\n",
    "        tgt_lens = tgt_lens.cuda()\n",
    "        input_seq = input_seq.cuda()\n",
    "        decoder_outputs = decoder_outputs.cuda()\n",
    "        \n",
    "    # -------------------------------------\n",
    "    # Training mode (enable dropout)\n",
    "    # -------------------------------------\n",
    "    encoder.train()\n",
    "    decoder.train()\n",
    "    \n",
    "    # -------------------------------------\n",
    "    # Zero gradients, since optimizers will accumulate gradients for every backward.\n",
    "    # -------------------------------------\n",
    "    encoder_optim.zero_grad()\n",
    "    decoder_optim.zero_grad()\n",
    "        \n",
    "    # -------------------------------------\n",
    "    # Forward encoder\n",
    "    # -------------------------------------\n",
    "    encoder_outputs, encoder_hidden = encoder(src_seqs, src_lens.data.tolist())\n",
    "\n",
    "    # -------------------------------------\n",
    "    # Forward decoder\n",
    "    # -------------------------------------\n",
    "    # Initialize decoder's hidden state as encoder's last hidden state.\n",
    "    decoder_hidden = encoder_hidden\n",
    "    \n",
    "    # Run through decoder one time step at a time.\n",
    "    for t in range(max_tgt_len):\n",
    "        \n",
    "        # decoder returns:\n",
    "        # - decoder_output   : (batch_size, vocab_size)\n",
    "        # - decoder_hidden   : (num_layers, batch_size, hidden_size)\n",
    "        # - attention_weights: (batch_size, max_src_len)\n",
    "        decoder_output, decoder_hidden, attention_weights = decoder(input_seq, decoder_hidden,\n",
    "                                                                    encoder_outputs, src_lens)\n",
    "\n",
    "        # Store decoder outputs.\n",
    "        decoder_outputs[t] = decoder_output\n",
    "        \n",
    "        # Next input is current target\n",
    "        input_seq = tgt_seqs[t]\n",
    "        \n",
    "        # Detach hidden state:\n",
    "        detach_hidden(decoder_hidden)\n",
    "        \n",
    "    # -------------------------------------\n",
    "    # Compute loss\n",
    "    # -------------------------------------\n",
    "    loss, pred_seqs, num_corrects, num_words = masked_cross_entropy(\n",
    "        decoder_outputs[:max_tgt_len].transpose(0,1).contiguous(), \n",
    "        tgt_seqs.transpose(0,1).contiguous(),\n",
    "        tgt_lens\n",
    "    )\n",
    "    \n",
    "    pred_seqs = pred_seqs[:max_tgt_len]\n",
    "    \n",
    "    # -------------------------------------\n",
    "    # Backward and optimize\n",
    "    # -------------------------------------\n",
    "    # Backward to get gradients w.r.t parameters in model.\n",
    "    loss.backward()\n",
    "    \n",
    "    # Clip gradients\n",
    "    encoder_grad_norm = nn.utils.clip_grad_norm(encoder.parameters(), opts.max_grad_norm)\n",
    "    decoder_grad_norm = nn.utils.clip_grad_norm(decoder.parameters(), opts.max_grad_norm)\n",
    "    clipped_encoder_grad_norm = compute_grad_norm(encoder.parameters())\n",
    "    clipped_decoder_grad_norm = compute_grad_norm(decoder.parameters())\n",
    "    \n",
    "    # Update parameters with optimizers\n",
    "    encoder_optim.step()\n",
    "    decoder_optim.step()\n",
    "        \n",
    "    return loss.data[0], pred_seqs, attention_weights, num_corrects, num_words,\\\n",
    "           encoder_grad_norm, decoder_grad_norm, clipped_encoder_grad_norm, clipped_decoder_grad_norm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main\n",
    "\n",
    "### Load dataset\n",
    "You can download the small grammatical error correction dataset from [here](https://github.com/keisks/jfleg)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====================================================================================================\n",
      "Dataset preprocessing log:\n",
      "- Loading and tokenizing source sentences...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2443191/2443191 [00:07<00:00, 343302.86it/s]\n",
      "  0%|          | 6142/2443191 [00:00<00:46, 52749.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "- Loading and tokenizing target sentences...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2443191/2443191 [00:06<00:00, 365671.01it/s]\n",
      "  1%|          | 25731/2443191 [00:00<00:09, 257275.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "- Building source counter...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2443191/2443191 [00:08<00:00, 272308.93it/s]\n",
      "  1%|          | 29309/2443191 [00:00<00:08, 293065.82it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "- Building target counter...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2443191/2443191 [00:10<00:00, 240597.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "- Building source vocabulary...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "50000it [00:00, 1461165.22it/s]\n",
      "100%|██████████| 50004/50004 [00:00<00:00, 2374494.52it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "- Building target vocabulary...\n",
      "====================================================================================================\n",
      "Dataset Info:\n",
      "- Number of source sentences: 2443191\n",
      "- Number of target sentences: 2443191\n",
      "- Source vocabulary size: 50004\n",
      "- Target vocabulary size: 50004\n",
      "- Shared vocabulary: True\n",
      "====================================================================================================\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# train_dataset = NMTDataset(src_path='../dataset/jfleg/dev/dev.src',\n",
    "#                            tgt_path='../dataset/jfleg/dev/dev.ref1')\n",
    "\n",
    "train_dataset = NMTDataset(src_path='../dataset/efcamdat/efcamdat2.changed.src.txt',\n",
    "                           tgt_path='../dataset/efcamdat/efcamdat2.changed.tgt.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 754/754 [00:00<00:00, 359334.76it/s]\n",
      "100%|██████████| 754/754 [00:00<00:00, 333597.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====================================================================================================\n",
      "Dataset preprocessing log:\n",
      "- Loading and tokenizing source sentences...\n",
      "- Loading and tokenizing target sentences...\n",
      "====================================================================================================\n",
      "Dataset Info:\n",
      "- Number of source sentences: 754\n",
      "- Number of target sentences: 754\n",
      "- Source vocabulary size: 50004\n",
      "- Target vocabulary size: 50004\n",
      "- Shared vocabulary: True\n",
      "====================================================================================================\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "valid_dataset = NMTDataset(src_path='../dataset/jfleg/dev/dev.src',\n",
    "                           tgt_path='../dataset/jfleg/dev/dev.ref0',\n",
    "                           src_vocab=train_dataset.src_vocab,\n",
    "                           tgt_vocab=train_dataset.tgt_vocab)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Batchify dataset using dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "batch_size = 48\n",
    "\n",
    "train_iter = DataLoader(dataset=train_dataset,\n",
    "                        batch_size=batch_size,\n",
    "                        shuffle=True,\n",
    "                        num_workers=4,\n",
    "                        collate_fn=collate_fn)\n",
    "\n",
    "valid_iter = DataLoader(dataset=valid_dataset,\n",
    "                        batch_size=batch_size, \n",
    "                        shuffle=False,\n",
    "                        num_workers=4,\n",
    "                        collate_fn=collate_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# If enabled, load checkpoint.\n",
    "LOAD_CHECKPOINT = True\n",
    "\n",
    "if LOAD_CHECKPOINT:\n",
    "    # Modify this path.\n",
    "    checkpoint_path = './checkpoints/seq2seq_2018-02-07 20:30:47_acc_88.15_loss_12.85_step_135000.pt'\n",
    "    checkpoint = load_checkpoint(checkpoint_path)\n",
    "    opts = checkpoint['opts']    \n",
    "else:\n",
    "    opts = AttrDict()\n",
    "\n",
    "    # Configure models\n",
    "    opts.word_vec_size = 300\n",
    "    opts.rnn_type = 'LSTM'\n",
    "    opts.hidden_size = 512\n",
    "    opts.num_layers = 2\n",
    "    opts.dropout = 0.3\n",
    "    opts.bidirectional = True\n",
    "    opts.attention = True\n",
    "    opts.share_embeddings = True\n",
    "    opts.pretrained_embeddings = True\n",
    "    opts.fixed_embeddings = True\n",
    "    opts.tie_embeddings = True # Tie decoder's input and output embeddings\n",
    "\n",
    "    # Configure optimization\n",
    "    opts.max_grad_norm = 2\n",
    "    opts.learning_rate = 0.001\n",
    "    opts.weight_decay = 1e-5 # L2 weight regularization\n",
    "    \n",
    "    # Configure training\n",
    "    opts.max_seq_len = 100 # max sequence length to prevent OOM.\n",
    "    opts.num_epochs = 5\n",
    "    opts.print_every_step = 20\n",
    "    opts.save_every_step = 5000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====================================================================================================\n",
      "Options log:\n",
      "- Load from checkpoint: True\n",
      "- Global step: 135000\n",
      "- word_vec_size: 300\n",
      "- rnn_type: LSTM\n",
      "- hidden_size: 512\n",
      "- num_layers: 2\n",
      "- dropout: 0.3\n",
      "- bidirectional: True\n",
      "- attention: True\n",
      "- share_embeddings: True\n",
      "- pretrained_embeddings: True\n",
      "- fixed_embeddings: True\n",
      "- tie_embeddings: True\n",
      "- max_grad_norm: 2\n",
      "- learning_rate: 0.001\n",
      "- weight_decay: 1e-05\n",
      "- max_seq_len: 100\n",
      "- num_epochs: 5\n",
      "- print_every_step: 20\n",
      "- save_every_step: 5000\n",
      "====================================================================================================\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print('='*100)\n",
    "print('Options log:')\n",
    "print('- Load from checkpoint: {}'.format(LOAD_CHECKPOINT))\n",
    "if LOAD_CHECKPOINT: print('- Global step: {}'.format(checkpoint['global_step']))\n",
    "for k,v in opts.items(): print('- {}: {}'.format(k, v))\n",
    "print('='*100 + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize embeddings and models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 38%|███▊      | 19055/50004 [00:00<00:00, 190512.48it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====================================================================================================\n",
      "Loading spacy glove embedding:\n",
      "- Vocabulary size: 50004\n",
      "- Word vector size: 300\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50004/50004 [00:00<00:00, 144659.40it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "- Unknown word count: 9362\n",
      "====================================================================================================\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Initialize vocabulary size.\n",
    "src_vocab_size = len(train_dataset.src_vocab.token2id)\n",
    "tgt_vocab_size = len(train_dataset.tgt_vocab.token2id)\n",
    "\n",
    "# Initialize embeddings.\n",
    "# We can actually put all modules in one module like `NMTModel`)\n",
    "# See: https://github.com/spro/practical-pytorch/issues/34\n",
    "word_vec_size = opts.word_vec_size if not opts.pretrained_embeddings else nlp.vocab.vectors_length\n",
    "src_embedding = nn.Embedding(src_vocab_size, word_vec_size, padding_idx=PAD)\n",
    "tgt_embedding = nn.Embedding(tgt_vocab_size, word_vec_size, padding_idx=PAD)\n",
    "\n",
    "if opts.share_embeddings:\n",
    "    assert(src_vocab_size == tgt_vocab_size)\n",
    "    tgt_embedding.weight = src_embedding.weight\n",
    "\n",
    "# Initialize models.\n",
    "encoder = EncoderRNN(embedding=src_embedding,\n",
    "                     rnn_type=opts.rnn_type,\n",
    "                     hidden_size=opts.hidden_size,\n",
    "                     num_layers=opts.num_layers,\n",
    "                     dropout=opts.dropout,\n",
    "                     bidirectional=opts.bidirectional)\n",
    "\n",
    "decoder = LuongAttnDecoderRNN(encoder, embedding=tgt_embedding,\n",
    "                              attention=opts.attention,\n",
    "                              tie_embeddings=opts.tie_embeddings,\n",
    "                              dropout=opts.dropout)\n",
    "\n",
    "if opts.pretrained_embeddings:\n",
    "    glove_embeddings = load_spacy_glove_embedding(nlp, train_dataset.src_vocab)\n",
    "    encoder.embedding.weight.data.copy_(glove_embeddings)\n",
    "    decoder.embedding.weight.data.copy_(glove_embeddings)\n",
    "    if opts.fixed_embeddings:\n",
    "        encoder.embedding.weight.requires_grad = False\n",
    "        decoder.embedding.weight.requires_grad = False\n",
    "        \n",
    "if LOAD_CHECKPOINT:\n",
    "    encoder.load_state_dict(checkpoint['encoder_state_dict'])\n",
    "    decoder.load_state_dict(checkpoint['decoder_state_dict'])\n",
    "    \n",
    "# Move models to GPU (need time for initial run)\n",
    "if USE_CUDA:\n",
    "    encoder.cuda()\n",
    "    decoder.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fine-tuning embeddings\n",
    "Recommend to use fine-tune after training for a while until the training loss don't decrease.\n",
    "\n",
    "TODO: Should be controlled in training loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "FINE_TUNE = True\n",
    "if FINE_TUNE:\n",
    "    encoder.embedding.weight.requires_grad = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====================================================================================================\n",
      "Model log:\n",
      "\n",
      "EncoderRNN (\n",
      "  (embedding): Embedding(50004, 300, padding_idx=0)\n",
      "  (rnn): LSTM(300, 256, num_layers=2, dropout=0.3, bidirectional=True)\n",
      ")\n",
      "LuongAttnDecoderRNN (\n",
      "  (embedding): Embedding(50004, 300, padding_idx=0)\n",
      "  (rnn): LSTM(300, 512, num_layers=2, dropout=0.3)\n",
      "  (W_a): Linear (512 -> 512)\n",
      "  (W_c): Linear (1024 -> 512)\n",
      "  (W_proj): Linear (512 -> 300)\n",
      "  (W_s): Linear (300 -> 50004)\n",
      ")\n",
      "- Encoder input embedding requires_grad=True\n",
      "- Decoder input embedding requires_grad=True\n",
      "- Decoder output embedding requires_grad=True\n",
      "====================================================================================================\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print('='*100)\n",
    "print('Model log:\\n')\n",
    "print(encoder)\n",
    "print(decoder)\n",
    "print('- Encoder input embedding requires_grad={}'.format(encoder.embedding.weight.requires_grad))\n",
    "print('- Decoder input embedding requires_grad={}'.format(decoder.embedding.weight.requires_grad))\n",
    "print('- Decoder output embedding requires_grad={}'.format(decoder.W_s.weight.requires_grad))\n",
    "print('='*100 + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize optimizers\n",
    "TODO: Different learning rate for fine tuning embeddings: https://discuss.pytorch.org/t/how-to-perform-finetuning-in-pytorch/419/7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Initialize optimizers (we can experiment different learning rates)\n",
    "encoder_optim = optim.Adam([p for p in encoder.parameters() if p.requires_grad], lr=opts.learning_rate, weight_decay=opts.weight_decay)\n",
    "decoder_optim = optim.Adam([p for p in decoder.parameters() if p.requires_grad], lr=opts.learning_rate, weight_decay=opts.weight_decay)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "\"\"\" Open port 6006 and see tensorboard.\n",
    "    Ref:  https://medium.com/@dexterhuang/%E7%B5%A6-pytorch-%E7%94%A8%E7%9A%84-tensorboard-bb341ce3f837\n",
    "\"\"\"\n",
    "from datetime import datetime\n",
    "from tensorboardX import SummaryWriter\n",
    "# --------------------------\n",
    "# Configure tensorboard\n",
    "# --------------------------\n",
    "model_name = 'seq2seq'\n",
    "datetime = ('%s' % datetime.now()).split('.')[0]\n",
    "experiment_name = '{}_{}'.format(model_name, datetime)\n",
    "tensorboard_log_dir = './tensorboard-logs/{}/'.format(experiment_name)\n",
    "writer = SummaryWriter(tensorboard_log_dir)\n",
    "\n",
    "# --------------------------\n",
    "# Configure training\n",
    "# --------------------------\n",
    "num_epochs = opts.num_epochs\n",
    "print_every_step = opts.print_every_step\n",
    "save_every_step = opts.save_every_step\n",
    "# For saving checkpoint and tensorboard\n",
    "global_step = 0 if not LOAD_CHECKPOINT else checkpoint['global_step']\n",
    "\n",
    "# --------------------------\n",
    "# Start training\n",
    "# --------------------------\n",
    "total_loss = 0\n",
    "total_corrects = 0\n",
    "total_words = 0\n",
    "prev_gpu_memory_usage = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for epoch in range(num_epochs):\n",
    "    for batch_id, batch_data in tqdm(enumerate(train_iter)):\n",
    "\n",
    "        # Unpack batch data\n",
    "        src_sents, tgt_sents, src_seqs, tgt_seqs, src_lens, tgt_lens = batch_data\n",
    "        \n",
    "        # Ignore batch if there is a long sequence.\n",
    "        max_seq_len = max(src_lens + tgt_lens)\n",
    "        if max_seq_len > opts.max_seq_len:\n",
    "            print('[!] Ignore batch: sequence length={} > max sequence length={}'.format(max_seq_len, opts.max_seq_len))\n",
    "            continue\n",
    "        \n",
    "        # Train.\n",
    "        loss, pred_seqs, attention_weights, num_corrects, num_words, \\\n",
    "        encoder_grad_norm, decoder_grad_norm, clipped_encoder_grad_norm, clipped_decoder_grad_norm \\\n",
    "        = train(src_sents, tgt_sents, src_seqs, tgt_seqs, src_lens, tgt_lens, encoder, decoder, encoder_optim, decoder_optim, opts)\n",
    "\n",
    "        # Statistics.\n",
    "        global_step += 1\n",
    "        total_loss += loss\n",
    "        total_corrects += num_corrects\n",
    "        total_words += num_words\n",
    "        total_accuracy = 100 * (total_corrects / total_words)\n",
    "        \n",
    "        # Save checkpoint.\n",
    "        if global_step % save_every_step == 0:\n",
    "            \n",
    "            checkpoint_path = save_checkpoint(opts, experiment_name, encoder, decoder, encoder_optim, decoder_optim, \n",
    "                                              total_accuracy, total_loss, global_step)\n",
    "            \n",
    "            print('='*100)\n",
    "            print('Save checkpoint to \"{}\".'.format(checkpoint_path))\n",
    "            print('='*100 + '\\n')\n",
    "\n",
    "        # Print statistics and write to Tensorboard.\n",
    "        if global_step % print_every_step == 0:\n",
    "            \n",
    "            curr_gpu_memory_usage = get_gpu_memory_usage(device_id=torch.cuda.current_device())\n",
    "            diff_gpu_memory_usage = curr_gpu_memory_usage - prev_gpu_memory_usage\n",
    "            prev_gpu_memory_usage = curr_gpu_memory_usage\n",
    "            \n",
    "            print('='*100)\n",
    "            print('Training log:')\n",
    "            print('- Epoch: {}/{}'.format(epoch, num_epochs))\n",
    "            print('- Global step: {}'.format(global_step))\n",
    "            print('- Total loss: {}'.format(total_loss))\n",
    "            print('- Total corrects: {}'.format(total_corrects))\n",
    "            print('- Total words: {}'.format(total_words))\n",
    "            print('- Total accuracy: {}'.format(total_accuracy))\n",
    "            print('- Current GPU memory usage: {}'.format(curr_gpu_memory_usage))\n",
    "            print('- Diff GPU memory usage: {}'.format(diff_gpu_memory_usage))\n",
    "            print('='*100 + '\\n')\n",
    "            \n",
    "            write_to_tensorboard(writer, global_step, total_loss, total_corrects, total_words, total_accuracy,\n",
    "                                 encoder_grad_norm, decoder_grad_norm, clipped_encoder_grad_norm, clipped_decoder_grad_norm,\n",
    "                                 encoder, decoder,\n",
    "                                 gpu_memory_usage={\n",
    "                                     'curr': curr_gpu_memory_usage,\n",
    "                                     'diff': diff_gpu_memory_usage\n",
    "                                 })\n",
    "            \n",
    "            total_loss = 0\n",
    "            total_corrects = 0\n",
    "            total_words = 0\n",
    "\n",
    "        # Free memory\n",
    "        del src_sents, tgt_sents, src_seqs, tgt_seqs, src_lens, tgt_lens, \\\n",
    "            loss, pred_seqs, attention_weights, num_corrects, num_words, \\\n",
    "            encoder_grad_norm, decoder_grad_norm, clipped_encoder_grad_norm, clipped_decoder_grad_norm\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "checkpoint_path = save_checkpoint(opts, experiment_name, encoder, decoder, encoder_optim, decoder_optim, \n",
    "                                              total_accuracy, total_loss, global_step)\n",
    "            \n",
    "print('='*100)\n",
    "print('Save checkpoint to \"{}\".'.format(checkpoint_path))\n",
    "print('='*100 + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def evaluate(src_sents, tgt_sents, src_seqs, tgt_seqs, src_lens, tgt_lens, encoder, decoder):\n",
    "    # -------------------------------------\n",
    "    # Prepare input and output placeholders\n",
    "    # -------------------------------------\n",
    "    # Last batch might not have the same size as we set to the `batch_size`\n",
    "    batch_size = src_seqs.size(1)\n",
    "    assert(batch_size == tgt_seqs.size(1))\n",
    "    \n",
    "    # Pack tensors to variables for neural network inputs (in order to autograd)\n",
    "    src_seqs = Variable(src_seqs, volatile=True)\n",
    "    tgt_seqs = Variable(tgt_seqs, volatile=True)\n",
    "    src_lens = Variable(torch.LongTensor(src_lens), volatile=True)\n",
    "    tgt_lens = Variable(torch.LongTensor(tgt_lens), volatile=True)\n",
    "\n",
    "    # Decoder's input\n",
    "    input_seq = Variable(torch.LongTensor([BOS] * batch_size), volatile=True)\n",
    "    \n",
    "    # Decoder's output sequence length = max target sequence length of current batch.\n",
    "    max_tgt_len = tgt_lens.data.max()\n",
    "    \n",
    "    # Store all decoder's outputs.\n",
    "    # **CRUTIAL** \n",
    "    # Don't set:\n",
    "    # >> decoder_outputs = Variable(torch.zeros(max_tgt_len, batch_size, decoder.vocab_size))\n",
    "    # Varying tensor size could cause GPU allocate a new memory causing OOM, \n",
    "    # so we intialize tensor with fixed size instead:\n",
    "    # `opts.max_seq_len` is a fixed number, unlike `max_tgt_len` always varys.\n",
    "    decoder_outputs = Variable(torch.zeros(opts.max_seq_len, batch_size, decoder.vocab_size), volatile=True)\n",
    "\n",
    "    # Move variables from CPU to GPU.\n",
    "    if USE_CUDA:\n",
    "        src_seqs = src_seqs.cuda()\n",
    "        tgt_seqs = tgt_seqs.cuda()\n",
    "        src_lens = src_lens.cuda()\n",
    "        tgt_lens = tgt_lens.cuda()\n",
    "        input_seq = input_seq.cuda()\n",
    "        decoder_outputs = decoder_outputs.cuda()\n",
    "        \n",
    "    # -------------------------------------\n",
    "    # Evaluation mode (disable dropout)\n",
    "    # -------------------------------------\n",
    "    encoder.eval()\n",
    "    decoder.eval()\n",
    "            \n",
    "    # -------------------------------------\n",
    "    # Forward encoder\n",
    "    # -------------------------------------\n",
    "    encoder_outputs, encoder_hidden = encoder(src_seqs, src_lens.data.tolist())\n",
    "    \n",
    "    # -------------------------------------\n",
    "    # Forward decoder\n",
    "    # -------------------------------------\n",
    "    # Initialize decoder's hidden state as encoder's last hidden state.\n",
    "    decoder_hidden = encoder_hidden\n",
    "    \n",
    "    # Run through decoder one time step at a time.\n",
    "    for t in range(max_tgt_len):\n",
    "        \n",
    "        # decoder returns:\n",
    "        # - decoder_output   : (batch_size, vocab_size)\n",
    "        # - decoder_hidden   : (num_layers, batch_size, hidden_size)\n",
    "        # - attention_weights: (batch_size, max_src_len)\n",
    "        decoder_output, decoder_hidden, attention_weights = decoder(input_seq, decoder_hidden,\n",
    "                                                                    encoder_outputs, src_lens)\n",
    "\n",
    "        # Store decoder outputs.\n",
    "        decoder_outputs[t] = decoder_output\n",
    "        \n",
    "        # Next input is current target\n",
    "        input_seq = tgt_seqs[t]\n",
    "        \n",
    "        # Detach hidden state (may not need this, since no BPTT)\n",
    "        detach_hidden(decoder_hidden)\n",
    "        \n",
    "    # -------------------------------------\n",
    "    # Compute loss\n",
    "    # -------------------------------------\n",
    "    loss, pred_seqs, num_corrects, num_words = masked_cross_entropy(\n",
    "        decoder_outputs[:max_tgt_len].transpose(0,1).contiguous(), \n",
    "        tgt_seqs.transpose(0,1).contiguous(),\n",
    "        tgt_lens\n",
    "    )\n",
    "    \n",
    "    pred_seqs = pred_seqs[:max_tgt_len]\n",
    "    \n",
    "    return loss.data[0], pred_seqs, attention_weights, num_corrects, num_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "16it [00:04,  3.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====================================================================================================\n",
      "Validation log:\n",
      "- Total loss: 23.030829787254333\n",
      "- Total corrects: 11675\n",
      "- Total words: 14994\n",
      "- Total accuracy: 77.86447912498332\n",
      "====================================================================================================\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "total_loss = 0\n",
    "total_corrects = 0\n",
    "total_words = 0\n",
    "\n",
    "for batch_id, batch_data in tqdm(enumerate(valid_iter)):\n",
    "    src_sents, tgt_sents, src_seqs, tgt_seqs, src_lens, tgt_lens = batch_data\n",
    "    \n",
    "    loss, pred_seqs, attention_weights, num_corrects, num_words \\\n",
    "        = evaluate(src_sents, tgt_sents, src_seqs, tgt_seqs, src_lens, tgt_lens, encoder, decoder)\n",
    "        \n",
    "    total_loss += loss\n",
    "    total_corrects += num_corrects\n",
    "    total_words += num_words\n",
    "    total_accuracy = 100 * (total_corrects / total_words)\n",
    "\n",
    "print('='*100)\n",
    "print('Validation log:')\n",
    "print('- Total loss: {}'.format(total_loss))\n",
    "print('- Total corrects: {}'.format(total_corrects))\n",
    "print('- Total words: {}'.format(total_words))\n",
    "print('- Total accuracy: {}'.format(total_accuracy))\n",
    "print('='*100 + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Translate (Inference)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def translate(src_text, train_dataset, encoder, decoder, max_seq_len, replace_unk=True):\n",
    "    # -------------------------------------\n",
    "    # Prepare input and output placeholders\n",
    "    # -------------------------------------\n",
    "    # Like dataset's `__getitem__()` and dataloader's `collate_fn()`.\n",
    "    src_sent = src_text.split()\n",
    "    src_seqs = torch.LongTensor([train_dataset.tokens2ids(tokens=src_text.split(),\n",
    "                                                          token2id=train_dataset.src_vocab.token2id,\n",
    "                                                          append_BOS=False, append_EOS=True)]).transpose(0,1)\n",
    "    src_lens = [len(src_seqs)]\n",
    "    \n",
    "    # Last batch might not have the same size as we set to the `batch_size`\n",
    "    batch_size = src_seqs.size(1)\n",
    "    \n",
    "    # Pack tensors to variables for neural network inputs (in order to autograd)\n",
    "    src_seqs = Variable(src_seqs, volatile=True)\n",
    "    src_lens = Variable(torch.LongTensor(src_lens), volatile=True)\n",
    "\n",
    "    # Decoder's input\n",
    "    input_seq = Variable(torch.LongTensor([BOS] * batch_size), volatile=True)\n",
    "    # Store output words and attention states\n",
    "    out_sent = []\n",
    "    all_attention_weights = torch.zeros(max_seq_len, len(src_seqs))\n",
    "    \n",
    "    # Move variables from CPU to GPU.\n",
    "    if USE_CUDA:\n",
    "        src_seqs = src_seqs.cuda()\n",
    "        src_lens = src_lens.cuda()\n",
    "        input_seq = input_seq.cuda()\n",
    "        \n",
    "    # -------------------------------------\n",
    "    # Evaluation mode (disable dropout)\n",
    "    # -------------------------------------\n",
    "    encoder.eval()\n",
    "    decoder.eval()\n",
    "        \n",
    "    # -------------------------------------\n",
    "    # Forward encoder\n",
    "    # -------------------------------------\n",
    "    encoder_outputs, encoder_hidden = encoder(src_seqs, src_lens.data.tolist())\n",
    "\n",
    "    # -------------------------------------\n",
    "    # Forward decoder\n",
    "    # -------------------------------------\n",
    "    # Initialize decoder's hidden state as encoder's last hidden state.\n",
    "    decoder_hidden = encoder_hidden\n",
    "    \n",
    "    # Run through decoder one time step at a time.\n",
    "    for t in range(max_seq_len):\n",
    "        \n",
    "        # decoder returns:\n",
    "        # - decoder_output   : (batch_size, vocab_size)\n",
    "        # - decoder_hidden   : (num_layers, batch_size, hidden_size)\n",
    "        # - attention_weights: (batch_size, max_src_len)\n",
    "        decoder_output, decoder_hidden, attention_weights = decoder(input_seq, decoder_hidden,\n",
    "                                                                    encoder_outputs, src_lens)\n",
    "\n",
    "        # Store attention weights.\n",
    "        # .squeeze(0): remove `batch_size` dimension since batch_size=1\n",
    "        all_attention_weights[t] = attention_weights.squeeze(0).cpu().data \n",
    "        \n",
    "        # Choose top word from decoder's output\n",
    "        prob, token_id = decoder_output.data.topk(1)\n",
    "        token_id = token_id[0][0] # get value\n",
    "        if token_id == EOS:\n",
    "            break\n",
    "        else:\n",
    "            if token_id == UNK and replace_unk:\n",
    "                # Replace unk by selecting the source token with the highest attention score.\n",
    "                score, idx = all_attention_weights[t].max(0)\n",
    "                token = src_sent[idx[0]]\n",
    "            else:\n",
    "                # <UNK>\n",
    "                token = train_dataset.tgt_vocab.id2token[token_id]\n",
    "            \n",
    "            out_sent.append(token)\n",
    "        \n",
    "        # Next input is chosen word\n",
    "        input_seq = Variable(torch.LongTensor([token_id]), volatile=True)\n",
    "        if USE_CUDA: input_seq = input_seq.cuda()\n",
    "            \n",
    "        # Repackage hidden state (may not need this, since no BPTT)\n",
    "        detach_hidden(decoder_hidden)\n",
    "    \n",
    "    src_text = ' '.join([train_dataset.src_vocab.id2token[token_id] for token_id in src_seqs.data.squeeze(1).tolist()])\n",
    "    out_text = ' '.join(out_sent)\n",
    "        \n",
    "    # all_attention_weights: (out_len, src_len)\n",
    "    return src_text, out_text, all_attention_weights[:len(out_sent)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Small test for translation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('He have a car <EOS>', 'He has a car', \n",
       "  0.8339  0.0667  0.0158  0.0305  0.0530\n",
       "  0.0471  0.8918  0.0325  0.0141  0.0144\n",
       "  0.0654  0.2109  0.5132  0.1534  0.0572\n",
       "  0.0083  0.0270  0.0291  0.8793  0.0564\n",
       " [torch.FloatTensor of size 4x5])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "src_text, out_text, all_attention_weights = translate('He have a car', train_dataset, encoder, decoder, max_seq_len=opts.max_seq_len)\n",
    "src_text, out_text, all_attention_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# check attention weight sum == 1\n",
    "[all_attention_weights[t].sum() for t in range(all_attention_weights.size(0))]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Translate a given text file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "test_src_texts = []\n",
    "with codecs.open('../dataset/jfleg/test/test.src', 'r', 'utf-8') as f:\n",
    "    test_src_texts = f.readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['New and new technology has been introduced to the society .\\n',\n",
       " 'One possible outcome is that an environmentally-induced reduction in motorization levels in the richer countries will outweigh any rise in motorization levels in the poorer countries .\\n',\n",
       " 'Every person needs to know a bit about math , sciences , arts , literature and history in order to stand out in society .\\n',\n",
       " 'While the travel company will most likely show them some interesting sites in order for their customers to advertise for their company to their family and friends , it is highly unlikely , that the company will tell about the sites that were not included in the tour -- for example due to entrance fees that would make the total package price overly expensive .\\n',\n",
       " 'Disadvantage is parking their car is very difficult .\\n']"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_src_texts[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "out_texts = []\n",
    "for src_text in test_src_texts:\n",
    "    _, out_text, _ = translate(src_text.strip(), train_dataset, encoder, decoder, max_seq_len=opts.max_seq_len)\n",
    "    out_texts.append(out_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['The new and new technology has been introduced to the society .',\n",
       " 'One possible outcome is that an environmentally-induced reduction in motorization levels in the higher countries will outweigh any rise in motorization levels in the high countries .',\n",
       " 'Every person needs to know a bit about math , sciences , arts , literature and history in order to stand out in society .',\n",
       " 'While the travel company will most likely show them some interesting sites in order for their customers to advertise for their company to their family and friends , it is highly unlikely , that the company will tell about the sites that were not included in the tour -- for example due to entrance fees that would make the total of price overly expensive .',\n",
       " 'The price is parking their cars are very difficult .']"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_texts[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save the predictions to text file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with codecs.open('./pred.txt', 'w', 'utf-8') as f:\n",
    "    for text in out_texts:\n",
    "        f.write(text + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate with GLEU metric\n",
    "If you're playing with grammatical error correction (GEC) corpus (jfleg),\n",
    "it has an evaluation script specifically for GEC task:\n",
    "\n",
    "Run:\n",
    "```\n",
    "python jfleg/eval/gleu.py \\\n",
    "-s jfleg/test/test.src \\\n",
    "-r jfleg/test/test.ref[0-3] \\\n",
    "--hyp ./pred.txt\n",
    "```\n",
    "\n",
    "Output (GLEU score, std, confidence interval):\n",
    "Note: The OpenNMT-py can further achieves ~0.49 GLEU score with the same model settings.\n",
    "TODO: Try to optimize the code.\n",
    "```\n",
    "Running GLEU...\n",
    "./pred.txt\n",
    "[['0.451747', '0.007620', '(0.437,0.467)']]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Notes:\n",
    "- Set `MAX_LENGTH` to training sequence is important to prevent OOM.\n",
    "    - Will effect：`decoder_outputs = Variable(torch.zeros(max_tgt_len, batch_size, decoder.vocab_size))`\n",
    "- Do not `next(iter(data_loader))` in training for-loop，could be very slow.\n",
    "- When computing `num_corrects`, need to cast `ByteTensor` using `.float()` in order to do `.sum()`, otherwise the result will overflow. Ref: https://discuss.pytorch.org/t/batch-size-and-validation-accuracy/4066/3\n",
    "- Very crutial to GPU memory usage: Don'T set `MAX_LENGTH` to `max(tgt_lens)`. Varying tensor size could cause GPU allocate a new memory, so we fixed tensor size instead: `decoder_outputs = Variable(torch.zeros(**MAX_LENGTH**, batch_size, decoder.vocab_size))`\n",
    "- Be careful if you only want to get `Variable`'s data and do some operations, for example, `sum()`, you should use `Variable(...).data.sum()` instead of `Variable(...).sum().data[0]`. This will create a new computational graph and if you do this in for-loop, it might increase memory.\n",
    "- Be careful to misuse `Variable`.\n",
    "- Do `detach` for RNN's hidden states, or it might increase memory when doing backprop.\n",
    "- If restart but GPU memory is not returned, kill all python processes: `>> ps x |grep python|awk '{print $1}'|xargs kill`\n",
    "- Forward decoder is time-consuming (for-loop).\n",
    "- Calling `backward()` free memory: https://discuss.pytorch.org/t/calling-loss-backward-reduce-memory-usage/2735\n",
    "\n",
    "### Try to:\n",
    "- Implement schedule sampling for training.\n",
    "- Implement beam search for evaluation and translation.\n",
    "- Understand and interpret param visualization on tensorboard.\n",
    "- Implement more RNN optimizing and regularization tricks:\n",
    "    - Set `max_seq_len` for preventing RNN OOM \n",
    "    - Xavier initializer\n",
    "    - Weight normalization and layer normalization: https://github.com/pytorch/pytorch/issues/1601\n",
    "    - Embedding dropout\n",
    "    - Weight dropping\n",
    "    - Variational dropout: [part1](https://becominghuman.ai/learning-note-dropout-in-recurrent-networks-part-1-57a9c19a2307), [part2](https://towardsdatascience.com/learning-note-dropout-in-recurrent-networks-part-2-f209222481f8), [part3](https://towardsdatascience.com/learning-note-dropout-in-recurrent-networks-part-3-1b161d030cd4)\n",
    "    - Zoneout\n",
    "    - Fraternal dropout\n",
    "    - Activation regularization (AR), and temporal activation regularization (TAR)\n",
    "    - Read more: [Regularizing and Optimizing LSTM Language Models](https://arxiv.org/pdf/1708.02182.pdf)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}