encoder-decoder-with-n-wo-attn-zh-en.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Import Required Lib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "import unicodedata\n",
    "import string\n",
    "import re\n",
    "import random\n",
    "from collections import Counter\n",
    "import pickle\n",
    "import random\n",
    "import pdb\n",
    "import logging\n",
    "import itertools\n",
    "import argparse'\n",
    "import time\n",
    "import os\n",
    "import math\n",
    "import copy\n",
    "import seaborn\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from torch import optim\n",
    "import torchtext\n",
    "from torch.utils.data import BatchSampler\n",
    "from torch.utils.data import SequentialSampler\n",
    "from torch.utils.data import Sampler\n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
    "\n",
    "from bleu_score import BLEU_SCORE\n",
    "from load_dataset_zh_wcharoption import *\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Build data loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Chinese(Dataset):\n",
    "    def __init__(self, df, val = False):\n",
    "        self.df = df\n",
    "        self.val = val\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "    def __getitem__(self, idx):\n",
    "        english = self.df.iloc[idx,:]['en_idized']\n",
    "        chin = self.df.iloc[idx,:]['zh_idized']\n",
    "        en_len = self.df.iloc[idx,:]['en_len']\n",
    "        zh_len = self.df.iloc[idx,:]['zh_len']\n",
    "        if self.val:\n",
    "            en_data = self.df.iloc[idx,:]['en_data'].lower()\n",
    "            return [chin,english,zh_len,en_len,en_data]\n",
    "        else:\n",
    "            return [chin,english,zh_len,en_len]\n",
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "    \n",
    "def vocab_collate_func(batch):\n",
    "    MAX_LEN_EN = 48\n",
    "    MAX_LEN_zh = 48\n",
    "    en_data = []\n",
    "    zh_data = []\n",
    "    en_len = []\n",
    "    zh_len = []\n",
    "    for datum in batch:\n",
    "        en_len.append(datum[3])\n",
    "        zh_len.append(datum[2])\n",
    "    max_batch_length_en = max(en_len)\n",
    "    max_batch_length_zh = max(zh_len)\n",
    "    if max_batch_length_en < MAX_LEN_EN:\n",
    "        MAX_LEN_EN = max_batch_length_en\n",
    "    if max_batch_length_zh < MAX_LEN_zh:\n",
    "        MAX_LEN_zh = max_batch_length_zh\n",
    "    # padding\n",
    "    for datum in batch:\n",
    "        if datum[2]>MAX_LEN_zh:\n",
    "            padded_vec_s1 = np.array(datum[0])[:MAX_LEN_zh]\n",
    "        else:\n",
    "            padded_vec_s1 = np.pad(np.array(datum[0]),\n",
    "                                pad_width=((0,MAX_LEN_zh - datum[2])),\n",
    "                                mode=\"constant\", constant_values=PAD_IDX)\n",
    "        if datum[3]>MAX_LEN_EN:\n",
    "            padded_vec_s2 = np.array(datum[1])[:MAX_LEN_EN]\n",
    "        else:\n",
    "            padded_vec_s2 = np.pad(np.array(datum[1]),\n",
    "                                pad_width=((0,MAX_LEN_EN - datum[3])),\n",
    "                                mode=\"constant\", constant_values=PAD_IDX)\n",
    "        en_data.append(padded_vec_s2)\n",
    "        zh_data.append(padded_vec_s1)\n",
    "    zh_data = np.array(zh_data)\n",
    "    en_data = np.array(en_data)\n",
    "    zh_len = np.array(zh_len)\n",
    "    en_len = np.array(en_len)\n",
    "\n",
    "    zh_len[zh_len>MAX_LEN_zh] = MAX_LEN_zh\n",
    "    en_len[en_len>MAX_LEN_EN] = MAX_LEN_EN\n",
    "        \n",
    "    return [torch.from_numpy(zh_data), torch.from_numpy(en_data),\n",
    "            torch.from_numpy(zh_len), torch.from_numpy(en_len)]\n",
    "\n",
    "def convert_idx_2_sent(tensor, lang_obj):\n",
    "    word_list = []\n",
    "    for i in tensor:\n",
    "        if i.item() not in set([PAD_IDX,EOS_token,SOS_token]):\n",
    "            word_list.append(lang_obj.index2word[i.item()])\n",
    "    return (' ').join(word_list)\n",
    "\n",
    "def convert_id_list_2_sent(list_idx, lang_obj):\n",
    "    word_list = []\n",
    "    if type(list_idx) == list:\n",
    "        for i in list_idx:\n",
    "            if i not in set([EOS_token]):\n",
    "                word_list.append(lang_obj.index2word[i])\n",
    "    else:\n",
    "        for i in list_idx:\n",
    "            if i.item() not in set([EOS_token,SOS_token,PAD_IDX]):\n",
    "                word_list.append(lang_obj.index2word[i.item()])\n",
    "    return (' ').join(word_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vocab_collate_func_val(batch):\n",
    "    return [torch.from_numpy(np.array(batch[0][0])).unsqueeze(0), torch.from_numpy(np.array(batch[0][1])).unsqueeze(0),\n",
    "            torch.from_numpy(np.array(batch[0][2])).unsqueeze(0), torch.from_numpy(np.array(batch[0][3])).unsqueeze(0),batch[0][4]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 Use character level Chinese or not"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MAX_LEN = 57\n",
    "train,val,test,en_lang,zh_lang = train_val_load(57, \"\", '/scratch/ark576/machine_translation_data/', char=False)\n",
    "# use character level Chinese\n",
    "#train,val,test,en_lang,zh_lang = train_val_load(57, \"\", '/scratch/ark576/machine_translation_data/', char=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs_dict = {'train':128,'validate':1, 'train_val':1,'val_train':128, 'test':1}\n",
    "shuffle_dict = {'train':True,'validate':False, 'train_val':False,'val_train':True, 'test':False}\n",
    "\n",
    "train_used = train\n",
    "val_used = val\n",
    "\n",
    "collate_fn_dict = {'train':vocab_collate_func, 'validate':vocab_collate_func_val,\\\n",
    "                   'train_val':vocab_collate_func_val,'val_train':vocab_collate_func,'test': vocab_collate_func_val}\n",
    "transformed_dataset = {'train': Chinese(train_used),\n",
    "                       'validate': Chinese(val_used, val = True),\n",
    "                       'train_val':Chinese(train.iloc[:50], val = True),\n",
    "                       'val_train':Chinese(val_used),\n",
    "                       'test':Chinese(test, val= True)\n",
    "                                               }\n",
    "\n",
    "dataloader = {x: DataLoader(transformed_dataset[x], batch_size=bs_dict[x], collate_fn=collate_fn_dict[x],\n",
    "                    shuffle=shuffle_dict[x], num_workers=0) for x in ['train', 'validate', 'train_val','val_train', 'test']}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Build Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderRNN(nn.Module):\n",
    "    def __init__(self, input_size, embed_dim, hidden_size,n_layers, rnn_type = 'lstm', device = 'cuda'):\n",
    "        super(EncoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.embedding = Embedding(input_size, embed_dim, PAD_IDX)\n",
    "        self.rnn_type =  rnn_type\n",
    "        self.dropout_in = nn.Dropout(p = 0.1)\n",
    "        self.n_layers = n_layers\n",
    "        self.device = device\n",
    "        if rnn_type == 'gru':\n",
    "            self.rnn = nn.GRU(embed_dim, hidden_size,batch_first=True,bidirectional=True, num_layers = self.n_layers, dropout = 0.2)\n",
    "        elif rnn_type == 'lstm':\n",
    "            self.rnn = LSTM(embed_dim, hidden_size, batch_first=True,bidirectional=True, num_layers = n_layers,dropout = 0.2)\n",
    "\n",
    "    def forward(self, enc_inp, src_len):\n",
    "        sorted_idx = torch.sort(src_len, descending=True)[1]\n",
    "        orig_idx = torch.sort(sorted_idx)[1]\n",
    "        embedded = self.embedding(enc_inp)\n",
    "        bs = embedded.size(0)\n",
    "        output = self.dropout_in(embedded)\n",
    "        if self.rnn_type == 'gru':\n",
    "            hidden =  self.initHidden(bs)\n",
    "            sorted_output = output[sorted_idx]\n",
    "            sorted_len = src_len[sorted_idx]\n",
    "            packed_output = nn.utils.rnn.pack_padded_sequence(sorted_output, sorted_len.data.tolist(), batch_first = True)\n",
    "            packed_outs, hiddden = self.rnn(packed_output,(hidden, c))\n",
    "            hidden = hidden[:,orig_idx,:]\n",
    "            output, _ = nn.utils.rnn.pad_packed_sequence(packed_outs, padding_value=PAD_IDX, batch_first = True)\n",
    "            output = output[orig_idx]\n",
    "            hidden = hidden.view(self.n_layers, 2, bs, -1).transpose(1, 2).contiguous().view(self.n_layers, bs, -1)\n",
    "            return output, hidden, hidden\n",
    "        elif self.rnn_type == 'lstm':\n",
    "            hidden, c = self.initHidden(bs)\n",
    "            sorted_output = output[sorted_idx]\n",
    "            sorted_len = src_len[sorted_idx]\n",
    "            packed_output = nn.utils.rnn.pack_padded_sequence(sorted_output, sorted_len.data.tolist(), batch_first = True)\n",
    "            packed_outs, (hiddden, c) = self.rnn(packed_output,(hidden, c))\n",
    "            hidden = hidden[:,orig_idx,:]\n",
    "            c = c[:,orig_idx,:]\n",
    "            output, _ = nn.utils.rnn.pad_packed_sequence(packed_outs, padding_value=PAD_IDX, batch_first = True)\n",
    "            output = output[orig_idx]\n",
    "            c = c.view(self.n_layers, 2, bs, -1).transpose(1, 2).contiguous().view(self.n_layers, bs, -1)\n",
    "            hidden = hidden.view(self.n_layers, 2, bs, -1).transpose(1, 2).contiguous().view(self.n_layers, bs, -1)\n",
    "            return output, hidden, c\n",
    "        \n",
    "    def initHidden(self,bs):\n",
    "        if self.rnn_type == 'gru' :\n",
    "            return torch.zeros(self.n_layers*2, bs, self.hidden_size).to(self.device)\n",
    "        elif self.rnn_type == 'lstm':\n",
    "            return torch.zeros(self.n_layers*2,bs,self.hidden_size).to(self.device),torch.zeros(self.n_layers*2,bs,self.hidden_size).to(self.device)\n",
    "\n",
    "class Attention_Module(nn.Module):\n",
    "    def __init__(self, hidden_dim, output_dim, device = 'cuda'):\n",
    "        super(Attention_Module, self).__init__()\n",
    "        self.l1 = Linear(hidden_dim, output_dim, bias = False)\n",
    "        self.l2 = Linear(hidden_dim+output_dim, output_dim, bias =  False)\n",
    "        self.device = device\n",
    "        \n",
    "    def forward(self, hidden, encoder_outs, src_lens):\n",
    "        ''' hiddden: bsz x hidden_dim\n",
    "        encoder_outs: bsz x sq_len x encoder dim (output_dim)\n",
    "        src_lens: bsz\n",
    "        \n",
    "        x: bsz x output_dim\n",
    "        attn_score: bsz x sq_len'''\n",
    "        x = self.l1(hidden)\n",
    "        att_score = (encoder_outs.transpose(0,1) * x.unsqueeze(0)).sum(dim = 2)\n",
    "        seq_mask = sequence_mask(src_lens, max_len = max(src_lens).item(), device = self.device).transpose(0,1)\n",
    "        masked_att = seq_mask*att_score\n",
    "        masked_att[masked_att==0] = -1e10\n",
    "        attn_scores = F.softmax(masked_att, dim=0)\n",
    "        x = (attn_scores.unsqueeze(2) * encoder_outs.transpose(0,1)).sum(dim=0)\n",
    "        x = torch.tanh(self.l2(torch.cat((x, hidden), dim=1)))\n",
    "        return x, attn_scores\n",
    "        \n",
    "class AttentionDecoderRNN(nn.Module):\n",
    "    def __init__(self, output_size, embed_dim, hidden_size, n_layers = 1, attention = True, device = 'cuda'):\n",
    "        super(AttentionDecoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        encoder_output_size = hidden_size\n",
    "        self.embedding = Embedding(output_size, embed_dim, PAD_IDX)\n",
    "        self.dropout = nn.Dropout(p=0.1)\n",
    "        self.n_layers = n_layers\n",
    "        self.device = device\n",
    "        self.att_layer = Attention_Module(self.hidden_size, encoder_output_size,self.device) if attention else None\n",
    "        self.layers = nn.ModuleList([\n",
    "            LSTMCell(\n",
    "                input_size=self.hidden_size + embed_dim if ((layer == 0) and attention) else embed_dim if layer == 0 else hidden_size,\n",
    "                hidden_size=hidden_size,\n",
    "            )\n",
    "            for layer in range(self.n_layers)\n",
    "        ])\n",
    "        self.fc_out = nn.Linear(self.hidden_size, output_size)\n",
    "        self.softmax = nn.LogSoftmax(dim=1)\n",
    "        \n",
    "    def forward(self, dec_input,context_vector, prev_hiddens,prev_cs,encoder_outputs,src_len):\n",
    "        bsz = dec_input.size(0)\n",
    "        output = self.embedding(dec_input)\n",
    "        output = self.dropout(output)\n",
    "        if self.att_layer is not None:\n",
    "            cated_input = torch.cat([output.squeeze(1),context_vector], dim = 1)\n",
    "        else:\n",
    "            cated_input = output.squeeze(1)\n",
    "        new_hiddens = []\n",
    "        new_cs = []\n",
    "        for i, rnn in enumerate(self.layers):\n",
    "            hidden, c = rnn(cated_input, (prev_hiddens[i], prev_cs[i]))\n",
    "            cated_input = self.dropout(hidden)\n",
    "            new_hiddens.append(hidden.unsqueeze(0))\n",
    "            new_cs.append(c.unsqueeze(0))\n",
    "        new_hiddens = torch.cat(new_hiddens, dim = 0)\n",
    "        new_cs = torch.cat(new_cs, dim = 0)\n",
    "\n",
    "        # apply attention using the last layer's hidden state\n",
    "        if self.att_layer is not None:\n",
    "            out, attn_score = self.att_layer(hidden, encoder_outputs, src_len)\n",
    "        else:\n",
    "            out = hidden\n",
    "            attn_score = None\n",
    "        context_vec = out\n",
    "        out = self.dropout(out)\n",
    "        out_vocab = self.softmax(self.fc_out(out))\n",
    "\n",
    "        return out_vocab, context_vec, new_hiddens, new_cs, attn_score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Utilites functions, training and evaluation functions "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Embedding(num_embeddings, embedding_dim, padding_idx):\n",
    "    m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)\n",
    "    nn.init.uniform_(m.weight, -0.1, 0.1)\n",
    "    nn.init.constant_(m.weight[padding_idx], 0)\n",
    "    return m\n",
    "\n",
    "\n",
    "def LSTM(input_size, hidden_size, **kwargs):\n",
    "    m = nn.LSTM(input_size, hidden_size,**kwargs)\n",
    "    for name, param in m.named_parameters():\n",
    "        if 'weight' in name or 'bias' in name:\n",
    "            param.data.uniform_(-0.1, 0.1)\n",
    "    return m\n",
    "\n",
    "\n",
    "def LSTMCell(input_size, hidden_size, **kwargs):\n",
    "    m = nn.LSTMCell(input_size, hidden_size,**kwargs)\n",
    "    for name, param in m.named_parameters():\n",
    "        if 'weight' in name or 'bias' in name:\n",
    "            param.data.uniform_(-0.1, 0.1)\n",
    "    return m\n",
    "\n",
    "\n",
    "def Linear(in_features, out_features, bias=True, dropout=0):\n",
    "    \"\"\"Linear layer (input: N x T x C)\"\"\"\n",
    "    m = nn.Linear(in_features, out_features, bias=bias)\n",
    "    m.weight.data.uniform_(-0.1, 0.1)\n",
    "    if bias:\n",
    "        m.bias.data.uniform_(-0.1, 0.1)\n",
    "    return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sequence_mask(sequence_length, max_len=None, device = 'cuda'):\n",
    "    if max_len is None:\n",
    "        max_len = sequence_length.max().item()\n",
    "    batch_size = sequence_length.size(0)\n",
    "    seq_range = torch.arange(0, max_len).long()\n",
    "    seq_range_expand = seq_range.unsqueeze(0).repeat([batch_size,1])\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).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode_decode(encoder,decoder,data_en,data_de,src_len,tar_len,rand_num = 0.95, val = False):\n",
    "    if not val:\n",
    "        use_teacher_forcing = True if random.random() < rand_num else False\n",
    "    #     print(\"tar_len\",tar_len)\n",
    "        bss = data_en.size(0)\n",
    "        en_out,en_hid,en_c = encoder(data_en, src_len)\n",
    "        max_src_len_batch = max(src_len).item()\n",
    "        max_tar_len_batch = max(tar_len).item()\n",
    "        prev_hiddens = en_hid\n",
    "        prev_cs = en_c\n",
    "        decoder_input = torch.tensor([[SOS_token]]*bss).to(device)\n",
    "        prev_output = torch.zeros((bss, en_out.size(-1))).to(device)\n",
    "        if use_teacher_forcing:\n",
    "            d_out = []\n",
    "            for i in range(max_tar_len_batch):\n",
    "                out_vocab, prev_output,prev_hiddens, prev_cs, attention_score = decoder(decoder_input,prev_output, \\\n",
    "                                                                                        prev_hiddens,prev_cs, en_out,\\\n",
    "                                                                                        src_len)\n",
    "                d_out.append(out_vocab.unsqueeze(-1))\n",
    "                decoder_input = data_de[:,i].view(-1,1)\n",
    "            d_out = torch.cat(d_out,dim=-1)\n",
    "        else:\n",
    "            d_out = []\n",
    "            for i in range(max_tar_len_batch):\n",
    "                out_vocab, prev_output,prev_hiddens, prev_cs, attention_score = decoder(decoder_input,prev_output, \\\n",
    "                                                                                        prev_hiddens,prev_cs, en_out,\\\n",
    "                                                                                        src_len)\n",
    "                d_out.append(out_vocab.unsqueeze(-1))\n",
    "                topv, topi = out_vocab.topk(1)\n",
    "                decoder_input = topi.squeeze().detach().view(-1,1)\n",
    "            d_out = torch.cat(d_out,dim=-1)\n",
    "        return d_out\n",
    "    else:\n",
    "        encoder.eval()\n",
    "        decoder.eval()\n",
    "        bss = data_en.size(0)\n",
    "        en_out,en_hid,en_c = encoder(data_en, src_len)\n",
    "        max_src_len_batch = max(src_len).item()\n",
    "        max_tar_len_batch = max(tar_len).item()\n",
    "        prev_hiddens = en_hid\n",
    "        prev_cs = en_c\n",
    "        decoder_input = torch.tensor([[SOS_token]]*bss).to(device)\n",
    "        prev_output = torch.zeros((bss, en_out.size(-1))).to(device)\n",
    "        d_out = []\n",
    "        for i in range(max_tar_len_batch):\n",
    "            out_vocab, prev_output,prev_hiddens, prev_cs, attention_score = decoder(decoder_input,prev_output, \\\n",
    "                                                                                    prev_hiddens,prev_cs, en_out,\\\n",
    "                                                                                    src_len)\n",
    "            d_out.append(out_vocab.unsqueeze(-1))\n",
    "            topv, topi = out_vocab.topk(1)\n",
    "            decoder_input = topi.squeeze().detach().view(-1,1)\n",
    "        d_out = torch.cat(d_out,dim=-1)\n",
    "        return d_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def flatten_cel_loss(input,target,nll):\n",
    "    input = input.transpose(1,2)\n",
    "    bs, sl = input.size()[:2]\n",
    "    return nll(input.contiguous().view(bs*sl,-1),target.contiguous().view(-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validation_new(encoder, decoder, val_dataloader, lang_en,lang_zh,m_type, verbose = False, replace_unk = False):\n",
    "    encoder.eval()\n",
    "    decoder.eval()\n",
    "    pred_corpus = []\n",
    "    true_corpus = []\n",
    "    src_corpus = []\n",
    "    running_loss = 0\n",
    "    running_total = 0\n",
    "    bl = BLEU_SCORE()\n",
    "    attention_scores_for_all_val = []\n",
    "    for data in val_dataloader:\n",
    "        encoder_i = data[0].to(device)\n",
    "        src_len = data[2].to(device)\n",
    "        bs,sl = encoder_i.size()[:2]\n",
    "        en_out,en_hid,en_c = encoder(encoder_i,src_len)\n",
    "        max_src_len_batch = max(src_len).item()\n",
    "        prev_hiddens = en_hid\n",
    "        prev_cs = en_c\n",
    "        decoder_input = torch.tensor([[SOS_token]]*bs).to(device)\n",
    "        prev_output = torch.zeros((bs, en_out.size(-1))).to(device)\n",
    "        d_out = []\n",
    "        attention_scores = []\n",
    "        for i in range(sl*2):\n",
    "            out_vocab, prev_output,prev_hiddens, prev_cs, attention_score = decoder(decoder_input,prev_output, \\\n",
    "                                                                                    prev_hiddens,prev_cs, en_out,\\\n",
    "                                                                                    src_len)\n",
    "            topv, topi = out_vocab.topk(1)\n",
    "            d_out.append(topi.item())\n",
    "            decoder_input = topi.squeeze().detach().view(-1,1)\n",
    "            if m_type == 'attention':\n",
    "                attention_scores.append(attention_score.unsqueeze(-1))\n",
    "            if topi.item() == EOS_token:\n",
    "                break\n",
    "        \n",
    "        if replace_unk:\n",
    "            true_sent = convert_id_list_2_sent(data[1][0],lang_en)\n",
    "            true_corpus.append(true_sent)\n",
    "        else:\n",
    "            true_corpus.append(data[-1])\n",
    "        src_sent = convert_id_list_2_sent(data[0][0],lang_zh)\n",
    "        src_corpus.append(src_sent)\n",
    "        pred_sent = convert_id_list_2_sent(d_out,lang_en)\n",
    "        pred_corpus.append(pred_sent)\n",
    "        if m_type == 'attention':\n",
    "            attention_scores = torch.cat(attention_scores, dim = -1)\n",
    "            attention_scores_for_all_val.append(attention_scores)\n",
    "        if verbose:\n",
    "            print(\"True Sentence:\",data[-1])\n",
    "            print(\"Pred Sentence:\", pred_sent)\n",
    "            print('-*'*50)\n",
    "    score = bl.corpus_bleu(pred_corpus,[true_corpus],lowercase=True)[0]\n",
    "    return score, attention_scores_for_all_val, pred_corpus, src_corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validation_beam_search(encoder, decoder, val_dataloader,lang_en,lang_zh,m_type, beam_size, verbose = False,\\\n",
    "                           device = 'cuda', replace_unk = False):\n",
    "    encoder.eval()\n",
    "    decoder.eval()\n",
    "    encoder = encoder.to(device)\n",
    "    decoder = decoder.to(device)\n",
    "    pred_corpus = []\n",
    "    true_corpus = []\n",
    "    src_corpus = []\n",
    "    running_loss = 0\n",
    "    running_total = 0\n",
    "    bl = BLEU_SCORE()\n",
    "    j = 0\n",
    "    attention_scores_for_all_val = []\n",
    "    for data in val_dataloader:\n",
    "#         print(j)\n",
    "        encoder_i = data[0].to(device)\n",
    "        src_len = data[2].to(device)\n",
    "        bs,sl = encoder_i.size()[:2]\n",
    "        en_out,en_hid,en_c = encoder(encoder_i,src_len)\n",
    "        max_src_len_batch = max(src_len).item()\n",
    "        prev_hiddens = en_hid\n",
    "        prev_cs = en_c\n",
    "        decoder_input = torch.tensor([[SOS_token]]*bs).to(device)\n",
    "        prev_output = torch.zeros((bs, en_out.size(-1))).to(device)\n",
    "        list_decoder_input = [None]*beam_size\n",
    "        beam_stop_flags = [False]*beam_size\n",
    "        beam_score = torch.zeros((bs,beam_size)).to(device)\n",
    "        list_d_outs = [[] for _ in range(beam_size)]\n",
    "        select_beam_size = beam_size\n",
    "        attention_scores = [[] for _ in range(beam_size)]\n",
    "        for i in range(sl+20):\n",
    "            if i == 0:\n",
    "                out_vocab, prev_output,prev_hiddens, prev_cs, attention_score = decoder(decoder_input,prev_output, \\\n",
    "                                                                                    prev_hiddens,prev_cs, en_out,\\\n",
    "                                                                                    src_len)\n",
    "                bss, vocab_size = out_vocab.size()\n",
    "                topv, topi = out_vocab.topk(beam_size)\n",
    "                list_prev_output = [prev_output]*beam_size\n",
    "                list_prev_hiddens = [prev_hiddens]*beam_size\n",
    "                list_prev_cs = [prev_cs]*beam_size\n",
    "                for b in range(beam_size):\n",
    "                    beam_score[0][b] = topv[0][b].item()\n",
    "                    list_decoder_input[b] = topi[0][b].squeeze().detach().view(-1,1)\n",
    "                    list_d_outs[b].append(topi[0][b].item())\n",
    "                    if m_type == 'attention':\n",
    "                        attention_scores[b].append(attention_score.unsqueeze(-1))\n",
    "                    if topi[0][b].item() == EOS_token:\n",
    "                        beam_stop_flags[b] = True\n",
    "            else:\n",
    "                beam_out_vocab = [None]*beam_size\n",
    "                temp_out = [None]*beam_size\n",
    "                temp_hid = [None]*beam_size\n",
    "                temp_c = [None]*beam_size\n",
    "                temp_attention_score = [[] for _ in range(beam_size)]\n",
    "                prev_d_outs = copy.deepcopy(list_d_outs)\n",
    "                for b in range(beam_size):\n",
    "                    if not beam_stop_flags[b]:\n",
    "                        beam_out_vocab[b], temp_out[b], temp_hid[b], temp_c[b], temp_attention_score[b] =\\\n",
    "                            decoder(list_decoder_input[b],list_prev_output[b],list_prev_hiddens[b],list_prev_cs[b],\\\n",
    "                                    en_out,src_len)\n",
    "                        beam_out_vocab[b] = beam_out_vocab[b] + beam_score[0][b]\n",
    "                    if beam_stop_flags[b]:\n",
    "                        beam_out_vocab[b] = torch.zeros(bss,vocab_size).fill_(float('-inf')).to(device)\n",
    "                beam_out_vocab = torch.cat(beam_out_vocab,dim = 1)\n",
    "                \n",
    "                topv, topi = beam_out_vocab.topk(beam_size)\n",
    "                id_for_hid = topi//vocab_size\n",
    "                topi_idx = topi%vocab_size\n",
    "                for b in range(beam_size):\n",
    "                    if not beam_stop_flags[b]:\n",
    "                        beam_score[0][b] = topv[0][b].item()\n",
    "                        list_decoder_input[b] = topi_idx[0][b].squeeze().detach().view(-1,1)\n",
    "                        list_d_outs[b] = copy.deepcopy(prev_d_outs[id_for_hid[0][b]])\n",
    "                        list_d_outs[b].append(topi_idx[0][b].item())\n",
    "                        if m_type == 'attention':\n",
    "                            attention_scores[b].append(temp_attention_score[b].unsqueeze(-1))\n",
    "                        if topi_idx[0][b].item() == EOS_token:\n",
    "                            beam_stop_flags[b] = True\n",
    "                        else:\n",
    "                            list_prev_output[b] = temp_out[id_for_hid[0][b]]\n",
    "                            list_prev_hiddens[b] = temp_hid[id_for_hid[0][b]]\n",
    "                            list_prev_cs[b] = temp_c[id_for_hid[0][b]]\n",
    "                if all(beam_stop_flags):\n",
    "                    break\n",
    "        id_max_score = np.argmax(beam_score)\n",
    "        d_out = list_d_outs[id_max_score]\n",
    "        if m_type == 'attention':\n",
    "            att_score = attention_scores[id_max_score]\n",
    "            att_score = torch.cat(att_score, dim = -1)\n",
    "            attention_scores_for_all_val.append(att_score)\n",
    "        if replace_unk:\n",
    "            true_sent = convert_id_list_2_sent(data[1][0],lang_en)\n",
    "            true_corpus.append(true_sent)\n",
    "        else:\n",
    "            true_corpus.append(data[-1])\n",
    "        pred_sent = convert_id_list_2_sent(d_out,lang_en)\n",
    "        pred_corpus.append(pred_sent)\n",
    "        src_sent = convert_id_list_2_sent(data[0][0], lang_zh)\n",
    "        src_corpus.append(src_sent)\n",
    "        if verbose:\n",
    "            print(\"True Sentence:\",data[-1])\n",
    "            print(\"Pred Sentence:\", pred_sent)\n",
    "            print('-*'*50)\n",
    "    score = bl.corpus_bleu(pred_corpus,[true_corpus],lowercase=True)[0]\n",
    "    return score, attention_scores_for_all_val, pred_corpus, src_corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(encoder_optimizer,decoder_optimizer, encoder, decoder, loss_fun,m_type, dataloader, en_lang,zh_lang,\\\n",
    "                num_epochs=60, val_every = 1, train_bleu_every = 10,clip = 0.1, rm = 0.8, enc_scheduler = None,\\\n",
    "               dec_scheduler = None, enc_dec_fn = encode_decode, val_fn = validation_new):\n",
    "    best_score = 0\n",
    "    best_bleu = 0\n",
    "    loss_hist = {'train': [], 'val_train': []}\n",
    "    bleu_hist = {'train': [], 'validate': []}\n",
    "    best_encoder_wts = None\n",
    "    best_decoder_wts = None\n",
    "    phases = ['train','val_train']\n",
    "    for epoch in range(num_epochs):\n",
    "        for ex, phase in enumerate(phases):\n",
    "            start = time.time()\n",
    "            total = 0\n",
    "            top1_correct = 0\n",
    "            running_loss = 0\n",
    "            running_total = 0\n",
    "            if phase == 'train':\n",
    "                encoder.train()\n",
    "                decoder.train()\n",
    "            else:\n",
    "                encoder.eval()\n",
    "                decoder.eval()\n",
    "            for data in dataloader[phase]:\n",
    "                encoder_optimizer.zero_grad()\n",
    "                decoder_optimizer.zero_grad()\n",
    "\n",
    "                encoder_i = data[0].to(device)\n",
    "                decoder_i = data[1].to(device)\n",
    "                src_len = data[2].to(device)\n",
    "                tar_len = data[3].to(device)\n",
    "                if phase == 'val_train':  \n",
    "                    out = enc_dec_fn(encoder,decoder,encoder_i,decoder_i,src_len,tar_len,rand_num=rm,val = True )\n",
    "                else:\n",
    "                    out = enc_dec_fn(encoder,decoder,encoder_i,decoder_i,src_len,tar_len,rand_num=rm,val = False )\n",
    "                N = decoder_i.size(0)\n",
    "                loss = loss_fun(out.float(), decoder_i.long())\n",
    "                running_loss += loss.item() * N\n",
    "                \n",
    "                total += N\n",
    "                if phase == 'train':\n",
    "                    loss.backward()\n",
    "                    torch.nn.utils.clip_grad_norm_(encoder.parameters(), clip)\n",
    "                    torch.nn.utils.clip_grad_norm_(decoder.parameters(), clip)\n",
    "                    encoder_optimizer.step()\n",
    "                    decoder_optimizer.step()\n",
    "                    \n",
    "            epoch_loss = running_loss / total \n",
    "            loss_hist[phase].append(epoch_loss)\n",
    "            print(\"epoch {} {} loss = {}, time = {}\".format(epoch, phase, epoch_loss,\n",
    "                                                                           time.time() - start))\n",
    "        if (enc_scheduler is not None) and (dec_scheduler is not None):\n",
    "            enc_scheduler.step(loss_hist['train'][-1])\n",
    "            dec_scheduler.step(loss_hist['train'][-1])\n",
    "        if epoch%val_every == 0:\n",
    "            val_bleu_score, _,_,_ = val_fn(encoder,decoder,dataloader['validate'],\\\n",
    "                                                                               en_lang,zh_lang,m_type, verbose=False, \\\n",
    "                                                                               replace_unk=True)\n",
    "            bleu_hist['validate'].append(val_bleu_score)\n",
    "            print(\"validation BLEU = \", val_bleu_score)\n",
    "            if val_bleu_score > best_bleu:\n",
    "                best_bleu = val_bleu_score\n",
    "                best_encoder_wts = encoder.state_dict()\n",
    "                best_decoder_wts = decoder.state_dict()\n",
    "        print('='*50)\n",
    "    encoder.load_state_dict(best_encoder_wts)\n",
    "    decoder.load_state_dict(best_decoder_wts)\n",
    "    print(\"Training completed. Best BLEU is {}\".format(best_bleu))\n",
    "    return encoder,decoder,loss_hist,bleu_hist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams['font.sans-serif'] = ['SimHei']\n",
    "mpl.rcParams['font.serif'] = ['SimHei']\n",
    "seaborn.set_style(\"darkgrid\",{\"Droid Sans Fallback\":['simhei', 'Arial']})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw(att_mat, x, y):\n",
    "    x = x.split(' ')\n",
    "    y = y.split(' ')\n",
    "    x += ['EOS']\n",
    "    y += ['EOS']\n",
    "    att_mat = att_mat.data.cpu().squeeze(1).numpy()\n",
    "    seaborn.heatmap(att_mat, \n",
    "                    xticklabels=x, square=True, yticklabels=y, vmin=0.0, vmax=1.0, \n",
    "                    cbar=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Without attention (Training and Evaluation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "encoder_wo_att = EncoderRNN(zh_lang.n_words,512,512, 1).to(device)\n",
    "decoder_wo_att = AttentionDecoderRNN(en_lang.n_words,512,1024,n_layers=1, attention = False).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# encoder_optimizer = optim.Adam(encoder_wo_att.parameters(), lr = 5e-3)\n",
    "# decoder_optimizer = optim.Adam(decoder_wo_att.parameters(), lr = 5e-3)\n",
    "encoder_optimizer = optim.SGD(encoder_wo_att.parameters(), lr=0.25,nesterov=True, momentum = 0.99)\n",
    "enc_scheduler = ReduceLROnPlateau(encoder_optimizer, min_lr=1e-4,factor = 0.5,  patience=0)\n",
    "decoder_optimizer = optim.SGD(decoder_wo_att.parameters(), lr=0.25,nesterov=True, momentum = 0.99)\n",
    "dec_scheduler = ReduceLROnPlateau(decoder_optimizer, min_lr=1e-4,factor = 0.5,  patience=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "encoder_wo_att, decoder_wo_att, loss_hist, acc_hist = train_model(encoder_optimizer, decoder_optimizer, encoder_wo_att, decoder_wo_att, criterion,\\\n",
    "                                            \"non_attention\", dataloader,en_lang,zh_lang, num_epochs = 15, rm = 0.95,\\\n",
    "                                           enc_scheduler = enc_scheduler, dec_scheduler = dec_scheduler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(encoder_wo_att.state_dict(), 'lstm_wo_att_enc_1_layer.pth')\n",
    "torch.save(decoder_wo_att.state_dict(), 'lstm_wo_att_dec_1_layer.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation using BLEU with beam size 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bleu_3_no_unk, att_score_wo, pred_wo, src_wo = validation_beam_search(encoder_wo_att, decoder_wo_att,dataloader['validate'],en_lang,\\\n",
    "                                                                      zh_lang, 'no_attention',3,verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bleu_3_no_unk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bleu_3_unk, att_score_wo, pred_wo, src_wo = validation_beam_search(encoder_wo_att, decoder_wo_att,dataloader['validate'],en_lang,\\\n",
    "                                                                      zh_lang, 'no_attention',3,verbose=False,\\\n",
    "                                                                  replace_unk = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bleu_3_unk"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6. With Attention (Training and Evaluation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "encoder_w_att = EncoderRNN(zh_lang.n_words,512,512, 1).to(device)\n",
    "decoder_w_att = AttentionDecoderRNN(en_lang.n_words,512,1024,n_layers=1, attention = True).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# encoder_optimizer = optim.Adam(encoder_wo_att.parameters(), lr = 5e-3)\n",
    "# decoder_optimizer = optim.Adam(decoder_wo_att.parameters(), lr = 5e-3)\n",
    "encoder_optimizer = optim.SGD(encoder_w_att.parameters(), lr=0.25,nesterov=True, momentum = 0.99)\n",
    "enc_scheduler = ReduceLROnPlateau(encoder_optimizer, min_lr=1e-4,factor = 0.5,  patience=0)\n",
    "decoder_optimizer = optim.SGD(decoder_w_att.parameters(), lr=0.25,nesterov=True, momentum = 0.99)\n",
    "dec_scheduler = ReduceLROnPlateau(decoder_optimizer, min_lr=1e-4,factor = 0.5,  patience=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "encoder_w_att, decoder_w_att, loss_hist, acc_hist = train_model(encoder_optimizer, decoder_optimizer, encoder_w_att, decoder_w_att, criterion,\\\n",
    "                                            \"non_attention\", dataloader,en_lang,zh_lang, num_epochs = 15, rm = 0.95,\\\n",
    "                                           enc_scheduler = enc_scheduler, dec_scheduler = dec_scheduler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(encoder_w_att.state_dict(), 'lstm_w_att_enc_1_layer.pth')\n",
    "torch.save(decoder_w_att.state_dict(), 'lstm_w_att_dec_1_layer.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation using BLEU with beam size 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bleu_10_no_unk, att_score_w, pred_w, src_w = validation_beam_search(encoder_w_att, decoder_w_att,dataloader['validate'],en_lang,\\\n",
    "                                                                      zh_lang, 'attention',10,verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bleu_10_unk, att_score_w, pred_w, src_wo = validation_beam_search(encoder_wo_att, decoder_wo_att,dataloader['validate'],en_lang,\\\n",
    "                                                                      zh_lang, 'attention',10,verbose=False,\\\n",
    "                                                                   replace_unk = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## plotting Attention maps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 5\n",
    "draw(attention_scores[n],pred_corpus[n],src_corpus[n])"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}