Seq2SeqAttention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.init as init\n",
    "from torch import optim\n",
    "import numpy as np\n",
    "\n",
    "import random\n",
    "import os\n",
    "\n",
    "seed = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TEST_FILE = \"data/test.txt\"\n",
    "TRAIN_FILE = \"data/train.txt\"\n",
    "F_VOCAB_FILE = \"data/vocab.f.txt\"\n",
    "Q_VOCAB_FILE = \"data/vocab.q.txt\"\n",
    "CHECKPOINT_FILE = \"out/Attention/WithGradientClippingAndLearningRateDecay/checkpoints/model_seq2seq_attention.pt\"\n",
    "LOSS_DIR = \"out/Attention/WithGradientClippingAndLearningRateDecay/losses\"\n",
    "METRICS_DIR = \"out/Attention/WithGradientClippingAndLearningRateDecay/metrics\"\n",
    "ATTENTION_ALIGNMENTS_DIR = \"out/Attention/WithGradientClippingAndLearningRateDecay/attention_alignments\"\n",
    "\n",
    "if not os.path.exists(os.path.dirname(CHECKPOINT_FILE)):\n",
    "    os.makedirs(os.path.dirname(CHECKPOINT_FILE))\n",
    "if not os.path.exists(LOSS_DIR):\n",
    "    os.makedirs(LOSS_DIR)\n",
    "if not os.path.exists(METRICS_DIR):\n",
    "    os.makedirs(METRICS_DIR)\n",
    "if not os.path.exists(ATTENTION_ALIGNMENTS_DIR):\n",
    "    os.makedirs(ATTENTION_ALIGNMENTS_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Options:\n",
    "    rnn_size = 50\n",
    "    init_weight = 0.08\n",
    "    decay_rate = 0.985\n",
    "    learning_rate = 0.01\n",
    "    plot_every = 10\n",
    "    print_every = 200\n",
    "    grad_clip = 5\n",
    "    dropout = 0\n",
    "    dropoutrec = 0\n",
    "    learning_rate_decay =  0.985\n",
    "    learning_rate_decay_after = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Language:\n",
    "    def __init__(self, question_vocab, form_vocab):\n",
    "        self.w2i = Language.__read_vocab(question_vocab)\n",
    "        self.lf2i = Language.__read_vocab(form_vocab)\n",
    "        self.i2lf = {self.lf2i[i] : i for i in self.lf2i}\n",
    "        \n",
    "    @staticmethod\n",
    "    def __read_vocab(filename):\n",
    "        t2i = {\"<s>\": 0, \"</s>\":1, \"UNK\": 2}\n",
    "        with open(filename) as target:\n",
    "            for line in target:\n",
    "                token = line.strip().split()[0]\n",
    "                if token not in t2i:\n",
    "                    t2i[token] = len(t2i)\n",
    "        return t2i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Entry:\n",
    "    def __init__(self, sentence, form, language):\n",
    "        self.sentence = sentence\n",
    "        self.form = form\n",
    "        self.sentence_tensor = Entry.__create_index_tensor(sentence, language.w2i)\n",
    "        self.form_tensor = Entry.__create_index_tensor(form, language.lf2i)\n",
    "        self.predicted_form = None\n",
    "    \n",
    "    @staticmethod\n",
    "    def __create_index_tensor(sequence, dictionary):\n",
    "        tensor = torch.zeros((1, len(sequence) + 2), dtype=torch.long)\n",
    "        tensor[0][0] = dictionary[\"<s>\"]\n",
    "        for idx, token in enumerate(sequence):\n",
    "            token_index = dictionary[token] if token in dictionary else dictionary[\"UNK\"]\n",
    "            tensor[0][idx+1] = token_index\n",
    "        tensor[0][-1] = dictionary[\"</s>\"]\n",
    "        return tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Data:\n",
    "    def __init__(self, train_file, test_file, question_vocab, form_vocab, shuffle_train=False):\n",
    "        self.language = Language(question_vocab, form_vocab)\n",
    "        self.train_entries = self.__read_data(train_file, self.language)\n",
    "        self.test_entries = self.__read_data(test_file, self.language)\n",
    "        if shuffle_train:\n",
    "            random.shuffle(self.train_entries)\n",
    "            \n",
    "    @staticmethod\n",
    "    def __read_data(file, language):\n",
    "        entries = []\n",
    "        with open(file) as target:\n",
    "            for line in target:\n",
    "                sentence, lf = line.strip().split(\"\\t\")\n",
    "                sentence = sentence.split()\n",
    "                lf = lf.split()\n",
    "                entries.append(Entry(sentence, lf, language))\n",
    "        return entries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTM(nn.Module):\n",
    "    def __init__(self, opt):\n",
    "        super(LSTM, self).__init__()\n",
    "        self.opt = opt\n",
    "        self.i2h = nn.Linear(opt.rnn_size, 4 * opt.rnn_size)\n",
    "        self.h2h = nn.Linear(opt.rnn_size, 4 * opt.rnn_size)\n",
    "        if opt.dropoutrec > 0:\n",
    "            self.dropout = nn.Dropout(opt.dropoutrec)\n",
    "            \n",
    "    def forward(self, x, prev_c, prev_h):\n",
    "        gates = self.i2h(x) + self.h2h(prev_h)\n",
    "        ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)\n",
    "        ingate = torch.sigmoid(ingate)\n",
    "        forgetgate = torch.sigmoid(forgetgate)\n",
    "        cellgate = torch.tanh(cellgate)\n",
    "        outgate = torch.sigmoid(outgate)\n",
    "        if self.opt.dropoutrec > 0:\n",
    "            cellgate = self.dropout(cellgate)\n",
    "        cy = (forgetgate * prev_c) + (ingate * cellgate)\n",
    "        hy = outgate * torch.tanh(cy)  # n_b x hidden_dim\n",
    "        return cy, hy\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNN(nn.Module):\n",
    "    def __init__(self, opt, input_size):\n",
    "        super(RNN, self).__init__()\n",
    "        self.opt = opt\n",
    "        self.hidden_size = opt.rnn_size\n",
    "        self.embedding = nn.Embedding(input_size, self.hidden_size)\n",
    "        self.lstm = LSTM(self.opt)\n",
    "        if opt.dropout > 0:\n",
    "            self.dropout = nn.Dropout(opt.dropout)\n",
    "        self.__initParameters()\n",
    "\n",
    "    def __initParameters(self):\n",
    "        for name, param in self.named_parameters():\n",
    "            if param.requires_grad:\n",
    "                init.uniform_(param, -self.opt.init_weight, self.opt.init_weight)\n",
    "                \n",
    "    def forward(self, input_src, prev_c, prev_h):\n",
    "        src_emb = self.embedding(input_src) # batch_size x src_length x emb_size\n",
    "        if self.opt.dropout > 0:\n",
    "            src_emb = self.dropout(src_emb)\n",
    "        prev_cy, prev_hy = self.lstm(src_emb, prev_c, prev_h)\n",
    "        return prev_cy, prev_hy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Attention(nn.Module):\n",
    "    def __init__(self, opt, output_size):\n",
    "        super(Attention, self).__init__()\n",
    "        self.opt = opt\n",
    "        self.hidden_size = opt.rnn_size\n",
    "\n",
    "        self.linear_att = nn.Linear(2*self.hidden_size, self.hidden_size)\n",
    "        self.linear_out = nn.Linear(self.hidden_size, output_size)\n",
    "        if opt.dropout > 0:\n",
    "            self.dropout = nn.Dropout(opt.dropout)\n",
    "\n",
    "        self.softmax = nn.Softmax(dim=1)\n",
    "        self.logsoftmax = nn.LogSoftmax(dim=1)\n",
    "        self.__initParameters()\n",
    "    \n",
    "    def __initParameters(self):\n",
    "        for name, param in self.named_parameters():\n",
    "            if param.requires_grad:\n",
    "                init.uniform_(param, -self.opt.init_weight, self.opt.init_weight)\n",
    "\n",
    "    def forward(self, enc_s_top, dec_s_top):\n",
    "        dot = torch.bmm(enc_s_top, dec_s_top.unsqueeze(2))\n",
    "        attention = self.softmax(dot.squeeze(2)).unsqueeze(2)\n",
    "        enc_attention = torch.bmm(enc_s_top.permute(0,2,1), attention)\n",
    "        hid = torch.tanh(self.linear_att(torch.cat((enc_attention.squeeze(2),dec_s_top), 1)))\n",
    "        h2y_in = hid\n",
    "        if self.opt.dropout > 0:\n",
    "            h2y_in = self.dropout(h2y_in)\n",
    "        h2y = self.linear_out(h2y_in)\n",
    "        pred = self.logsoftmax(h2y)\n",
    "        return pred, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model:\n",
    "    def __init__(self):\n",
    "        self.opt = None\n",
    "        self.encoder = None\n",
    "        self.decoder = None\n",
    "        self.attention = None\n",
    "        self.optimizers = {}\n",
    "        self.criterion = None\n",
    "        \n",
    "    def create(self, opt, data):\n",
    "        self.opt = opt\n",
    "        self.encoder = RNN(self.opt, len(data.language.w2i))\n",
    "        self.decoder = RNN(self.opt, len(data.language.lf2i))\n",
    "        self.attention = Attention(self.opt, len(data.language.lf2i))\n",
    "        self.optimizers[\"encoder_optimizer\"] = optim.RMSprop(self.encoder.parameters(), lr=self.opt.learning_rate, alpha=self.opt.decay_rate)\n",
    "        self.optimizers[\"decoder_optimizer\"] = optim.RMSprop(self.decoder.parameters(), lr=self.opt.learning_rate, alpha=self.opt.decay_rate)\n",
    "        self.optimizers[\"attention_optimizer\"] = optim.RMSprop(self.attention.parameters(), lr=self.opt.learning_rate, alpha=self.opt.decay_rate)\n",
    "        self.criterion = nn.NLLLoss(ignore_index=0)\n",
    " \n",
    "    def train(self):\n",
    "        self.encoder.train()\n",
    "        self.decoder.train()\n",
    "        self.attention.train()\n",
    "        \n",
    "    def eval(self):\n",
    "        self.encoder.eval()\n",
    "        self.decoder.eval()\n",
    "        self.attention.eval()\n",
    "    \n",
    "    def step(self):\n",
    "        for optimizer in self.optimizers:\n",
    "            self.optimizers[optimizer].step()\n",
    "            \n",
    "    def zero_grad(self):\n",
    "        for optimizer in self.optimizers:\n",
    "            self.optimizers[optimizer].zero_grad()\n",
    "            \n",
    "    def rate_decay(self):\n",
    "        for optimizer in self.optimizers:\n",
    "            for param_group in self.optimizers[optimizer].param_groups:\n",
    "                param_group['lr'] = param_group['lr'] * self.opt.learning_rate_decay\n",
    "        \n",
    "    def grad_clip(self):\n",
    "        torch.nn.utils.clip_grad_value_(self.encoder.parameters(), self.opt.grad_clip)\n",
    "        torch.nn.utils.clip_grad_value_(self.decoder.parameters(), self.opt.grad_clip)\n",
    "        torch.nn.utils.clip_grad_value_(self.attention.parameters(), self.opt.grad_clip)\n",
    "            \n",
    "    def save(self, filename):\n",
    "        checkpoint = {}\n",
    "        checkpoint[\"opt\"] = self.opt\n",
    "        checkpoint[\"encoder\"] = self.encoder.state_dict()\n",
    "        checkpoint[\"decoder\"] = self.decoder.state_dict()\n",
    "        checkpoint[\"attention\"] = self.attention.state_dict()\n",
    "        for o in self.optimizers:\n",
    "             checkpoint[o] = self.optimizers[o].state_dict()\n",
    "        torch.save(checkpoint, filename)\n",
    "    \n",
    "    def load(self, filename, data):\n",
    "        checkpoint = torch.load(filename)\n",
    "    \n",
    "        opt = checkpoint[\"opt\"]\n",
    "        self.create(opt, data)\n",
    "        \n",
    "        self.encoder.load_state_dict(checkpoint[\"encoder\"])\n",
    "        self.decoder.load_state_dict(checkpoint[\"decoder\"])\n",
    "        self.attention.load_state_dict(checkpoint[\"attention\"])\n",
    "        for o in self.optimizers:\n",
    "            self.optimizers[o].load_state_dict(checkpoint[o])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, sentence, form):\n",
    "    model.zero_grad()\n",
    "\n",
    "    outputs = torch.zeros((1, sentence.size(1), model.opt.rnn_size), dtype=torch.float, requires_grad=False)\n",
    "    c = torch.zeros((1, model.opt.rnn_size), dtype=torch.float, requires_grad=True)\n",
    "    h = torch.zeros((1, model.opt.rnn_size), dtype=torch.float, requires_grad=True)\n",
    "    for i in range(sentence.size(1)):\n",
    "        c, h = model.encoder(sentence[:, i], c, h)\n",
    "        outputs[:,i, :] = h\n",
    "        \n",
    "    loss = 0\n",
    "    for i in range(form.size(1)-1):\n",
    "        c, h = model.decoder(form[:, i], c, h)\n",
    "        pred, attention_weights = model.attention(outputs, h)\n",
    "        loss += model.criterion(pred, form[:, i+1])\n",
    "        \n",
    "    loss.backward()\n",
    "    if model.opt.grad_clip != -1:\n",
    "        model.grad_clip()\n",
    "    model.step()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(model, sentence, lf2i):\n",
    "    c = torch.zeros((1, model.opt.rnn_size), dtype=torch.float, requires_grad=True)\n",
    "    h = torch.zeros((1, model.opt.rnn_size), dtype=torch.float, requires_grad=True)\n",
    "    outputs = torch.zeros((1, sentence.size(1), model.opt.rnn_size), dtype=torch.float, requires_grad=True)\n",
    "\n",
    "    for i in range(sentence.size(1)):\n",
    "        c, h = model.encoder(sentence[:, i], c, h)\n",
    "        outputs[:,i, :] = h\n",
    "            \n",
    "    prev = torch.tensor([lf2i['<s>']], dtype=torch.long)\n",
    "    predicted_form = []\n",
    "    counter = 0\n",
    "    max_length = 100\n",
    "    decoder_attentions = torch.zeros(max_length, sentence.size(1))\n",
    "    while True:\n",
    "        c, h = model.decoder(prev, c, h)\n",
    "        pred, attention_weights = model.attention(outputs, h)\n",
    "        decoder_attentions[counter] = attention_weights.view(-1)\n",
    "        form_id = pred.argmax().item()\n",
    "        prev = torch.tensor([form_id], dtype=torch.long)\n",
    "        counter += 1\n",
    "        if form_id == lf2i[\"</s>\"] or counter >= 100:\n",
    "            break\n",
    "        predicted_form.append(form_id)\n",
    "    return predicted_form, decoder_attentions[:counter]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as ticker\n",
    "\n",
    "\n",
    "def showPlot(points, fig_name, extra_info):\n",
    "    plt.figure()\n",
    "    fig, ax = plt.subplots()\n",
    "    # this locator puts ticks at regular intervals\n",
    "    loc = ticker.MultipleLocator(base=0.2)\n",
    "    ax.yaxis.set_major_locator(loc)\n",
    "    plt.title(extra_info) \n",
    "    plt.plot(points)\n",
    "    plt.savefig(\"{}.png\".format(fig_name))\n",
    "    plt.close()\n",
    "\n",
    "def showAttention(fig_name, input_sentence, output_words, attentions):\n",
    "    # Set up figure with colorbar\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    cax = ax.matshow(attentions, cmap='bone')\n",
    "    fig.colorbar(cax)\n",
    "\n",
    "    # Set up axes\n",
    "    input_sentence = ['']  + ['<s>'] + input_sentence+ ['</s>']\n",
    "    output_words = ['']  +output_words + [\"</s>\"]\n",
    "    ax.set_xticklabels(input_sentence, rotation=90)\n",
    "    ax.set_yticklabels( output_words)\n",
    "\n",
    "    # Show label at every tick\n",
    "    ax.xaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "    plt.savefig(\"{}.png\".format(fig_name))\n",
    "    plt.close()\n",
    "\n",
    "def evaluateAndShowAttention(fig_name, input_sentence, input_tensor, preds, attentions, i2lf):\n",
    "    output_words = [i2lf[i] for i in preds]\n",
    "    showAttention(fig_name, input_sentence, output_words, attentions.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_and_test(epoch_num):\n",
    "    random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    np.random.seed(seed)\n",
    "\n",
    "    data = Data(TRAIN_FILE, TRAIN_FILE, Q_VOCAB_FILE, F_VOCAB_FILE, shuffle_train=True)\n",
    "    model = Model()\n",
    "    model.create(Options(), data)\n",
    "\n",
    "    losses = []\n",
    "    max_acc = 0\n",
    "    maxAccEpochId = 0\n",
    "    accuracies = []\n",
    "    chunkend_losses = []\n",
    "    for epoch in range(epoch_num):\n",
    "        print(\"---Epoch {}---\\n\".format(epoch+1))\n",
    "        \n",
    "        print(\"Training...\")\n",
    "        model.train()\n",
    "        plot_data = []\n",
    "        for index, entry in enumerate(data.train_entries):\n",
    "            loss = train(model,\n",
    "                         entry.sentence_tensor,\n",
    "                         entry.form_tensor)\n",
    "            if index != 0:\n",
    "                if index % model.opt.plot_every == 0:     \n",
    "                    plot_data.append(np.mean(losses[epoch*len(data.train_entries)+index-model.opt.plot_every:]))\n",
    "                    chunkend_losses.append(np.mean(plot_data))\n",
    "                if index % model.opt.print_every == 0:\n",
    "                    print(\"Index {} Loss {}\".format(index,\n",
    "                                                    np.mean(losses[epoch*len(data.train_entries)+index-model.opt.print_every:])))\n",
    "            losses.append(loss.item())\n",
    "\n",
    "        if model.opt.learning_rate_decay < 1:\n",
    "            if epoch >= model.opt.learning_rate_decay_after:\n",
    "                model.rate_decay()\n",
    "\n",
    "        print(\"Predicting..\")\n",
    "        model.eval()\n",
    "        correct = 0.0\n",
    "        with torch.no_grad():\n",
    "            for index, entry in enumerate(data.test_entries):\n",
    "                prediction, attention_weights = predict(model, entry.sentence_tensor, data.language.lf2i)\n",
    "                entry.prediction = [data.language.i2lf[p] for p in prediction]\n",
    "                if len(entry.form) == len(entry.prediction):\n",
    "                    same = True\n",
    "                    for g, p in zip(entry.form, entry.prediction):\n",
    "                        if g != p:\n",
    "                            same = False\n",
    "                    if same:\n",
    "                        correct += 1\n",
    "  \n",
    "        accuracy = 100*(correct/len(data.test_entries))\n",
    "        accuracies.append(accuracy)\n",
    "        if accuracy > max_acc:\n",
    "            max_acc = accuracy\n",
    "            maxAccEpochId = epoch\n",
    "            model.save(CHECKPOINT_FILE)\n",
    "\n",
    "        print(\"Accuracy: {} Max Accuracy {}\".format(accuracy, max_acc))\n",
    "        file_name = \"{}/epoch.{}\".format(LOSS_DIR, epoch)\n",
    "        extra = \"Mean Loss {0:.2f}\".format(np.mean(plot_data))\n",
    "        showPlot(plot_data, file_name, extra)\n",
    "\n",
    "    file_name = \"{}/{}\".format(METRICS_DIR, \"accuracies\")\n",
    "    extra = \"Maximum Accuracy {0:.2f} at epoch {1}\".format(np.max(accuracies), maxAccEpochId)\n",
    "    showPlot(accuracies, file_name, extra)\n",
    "    file_name = \"{}/{}\".format(LOSS_DIR, \"all_losses\") \n",
    "    extra = \"Mean Loss {0:.2f}\".format(np.mean(chunkend_losses))\n",
    "    showPlot(chunkend_losses, file_name, extra)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_predictions(model, data):\n",
    "    model.eval()\n",
    "    correct = 0.0\n",
    "    with torch.no_grad():\n",
    "        for index, entry in enumerate(data.test_entries):\n",
    "            prediction, attention_weights = predict(model, entry.sentence_tensor, data.language.lf2i)\n",
    "            entry.prediction = [data.language.i2lf[p] for p in prediction]\n",
    "            showAttention(\"{}/entry.{}\".format(ATTENTION_ALIGNMENTS_DIR, index+1), entry.sentence, entry.prediction, attention_weights.numpy())\n",
    "            if len(entry.form) == len(entry.prediction):\n",
    "                same = True\n",
    "                for g, p in zip(entry.form, entry.prediction):\n",
    "                    if g != p:\n",
    "                        same = False\n",
    "                if same:\n",
    "                    correct += 1\n",
    "\n",
    "    accuracy = 100*(correct/len(data.test_entries))\n",
    "    print(\"Accuracy : {}\".format(accuracy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_and_test(epoch_num=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = Data(TRAIN_FILE, TEST_FILE, Q_VOCAB_FILE, F_VOCAB_FILE)\n",
    "model = Model()\n",
    "model.load(CHECKPOINT_FILE, data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_predictions(model, data)"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}