9.seq2seq_models.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "\n",
    "import time\n",
    "import math\n",
    "import numpy as np\n",
    "\n",
    "import torch as t\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable as V\n",
    "import torch.utils.data as Data\n",
    "import torch.optim as optim\n",
    "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\n",
    "\n",
    "from pyfile.text_loader import TextDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "MAX_LENGTH = 100\n",
    "SOS_token = chr(0)\n",
    "EOS_token = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Helper function to create Variable based on\n",
    "# the cuda availability\n",
    "\n",
    "def cuda_variable(tensor):\n",
    "    # Do cuda() before wrapping with variable\n",
    "    if torch.cuda.is_available():\n",
    "        return V(tensor.cuda())\n",
    "    else:\n",
    "        return V(tensor)\n",
    "\n",
    "# Sting to char tensor\n",
    "def str2tensor(msg, eos=False):\n",
    "    tensor = [ord(c) for c in msg]\n",
    "    if eos:\n",
    "        tensor.append(EOS_token)\n",
    "    return cuda_available(t.LongTensor(tensor))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# To demonstrate seq2seq, we don't handle batch in code,\n",
    "# and our encoder runs this one step at a time\n",
    "# It's extremely slow, and please do not use in practice.\n",
    "# We need to use (1) batch and (2) Data parallelism\n",
    "\n",
    "class EncoderRNN(nn.Module):\n",
    "    \n",
    "    def __init__(self, input_size, hidden_size, n_layers=1):\n",
    "        self.hidden_size = hidden_size\n",
    "        self.n_layers = n_layers\n",
    "        \n",
    "        super(EncoderRNN, self).__init__()\n",
    "        self.embedding = nn.Embedding(input_size, hidden_size)\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size, n_layers)\n",
    "        \n",
    "    def forward(self, word_inputs, hidden):\n",
    "        # Note: we run this all at once (over the whole input sequence)\n",
    "        seq_len = len(word_inputs)\n",
    "        # input shape: S * B (=1) * I (input size)\n",
    "        embedded = self.embedding(word_inputs.view(seq_len, 1, -1))\n",
    "        output, hidden = self.gru(embedded, hidden)\n",
    "        return output, hidden\n",
    "    \n",
    "    def init_hidden(self):\n",
    "        # (num_layers * num_directions, batch, hidden_size)\n",
    "        return cuda_available(t.zeros(self.n_layers, 1, self.hidden_size))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class DecoderRNN(nn.Module):\n",
    "    \n",
    "    def __init__(self, hidden_size, output_size, n_layers=1):\n",
    "        super(DecoderRNN, self).__init__()\n",
    "        \n",
    "        self.embedding = nn.Embedding(output_size, hidden_size)\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size, n_layers)\n",
    "        self.out = nn.Linear(hidden_size, output_size)\n",
    "    \n",
    "    def forward(self, inputs, hidden):\n",
    "        # input shape: S(=1) * B(=1) * I(input size)\n",
    "        # Note: we run this one step at a time. (Sequence size = 1)\n",
    "        outputs = self.embedding(inputs).view(1, 1, -1)\n",
    "        output, hidden = self.gru(output, hidden)\n",
    "        output = self.out(output[0])\n",
    "        # No need softmax, since we are using CrossEntropyLoss\n",
    "        return output, hidden\n",
    "    \n",
    "    def init_hidden(self):\n",
    "        # (num_layers * num_directions, batch, hidden_size)\n",
    "        return cuda_variable(t.zeros(self.n_layers, 1, self.hidden_size))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class AttnDecoderRNN(nn.Module):\n",
    "    \n",
    "    def __init__(self, hidden_size, output_size, n_layers=1, dropout_p=0.1):\n",
    "        super(AttnDecoderRNN, self).__init__()\n",
    "        \n",
    "        # Linear for attention\n",
    "        self.attn = nn.Linear(hidden_size, hidden_size)\n",
    "        \n",
    "        # Define layers\n",
    "        self.embedding = nn.Embedding(output_size, hidden_size)\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size, n_layers, dropout=dropout_p)\n",
    "        self.out = nn.Linear(hidden_size * 2, output_size)\n",
    "    \n",
    "    def forward(self, word_input, last_hidden, encoder_hiddens):\n",
    "        # Note: we run this one step (S=1) at a time\n",
    "        # Get the embedding of the current input word (last output word)\n",
    "        rnn_input = self.embedding(word_input).view(1, 1, -1)   # S = 1 * B * I\n",
    "        rnn_output, hidden = self.gru(rnn_input, last_hidden)\n",
    "        \n",
    "        # Calculate attention from current RNN state and all encoder outputs;\n",
    "        # apply to encoder outputs\n",
    "        attn_weights = self.get_att_weight(rnn_output.squeeze(0), encoder_hiddens)\n",
    "        context = attn_weights.bmm(encoder_hiddens.transpose(0, 1)) # B * S(=1) * I\n",
    "        \n",
    "        # Final output layer (next word prediction) using the RNN hidden state\n",
    "        # and context vector\n",
    "        rnn_output = rnn_output.squeeze(0)     # S(=1) * B * I ---> B * I\n",
    "        context = context.squeeze(1)           # B * S(=1) * I ---> B * I\n",
    "        output = self.out(t.cat(rnn_output, context), 1)\n",
    "        \n",
    "        # Return final output, hidden_state, and attention weights(for visualization)\n",
    "        return output, hidden, attn_weights\n",
    "    \n",
    "    def get_att_weight(self, hidden, encoder_hiddens):\n",
    "        seq_len = len(encoder_hiddens)\n",
    "        \n",
    "        # Create variables to store attention energies\n",
    "        attn_scores = cuda_variable(t.zeros(seq_len))   # B * 1 * S\n",
    "        \n",
    "        # Calculate energies for each encoder hidden\n",
    "        for i in range(seq_len):\n",
    "            attn_scores[i] = self.get_att_score(hidden, encoder_hiddens[i])\n",
    "        \n",
    "        # Normalize scores to weights in range 0 to 1,\n",
    "        # resize to 1 * 1 * seq_len\n",
    "        return F.softmax(attn_scores).view(1, 1, -1)\n",
    "    \n",
    "    # score = h^T W h^e = h dot (W h^e)\n",
    "    # TODO: We need to implement different score models\n",
    "    def get_att_score(self, hidden, encoder_hidden):\n",
    "        score = self.attn(encoder_hidden)\n",
    "        return t.dot(hidden.view(-1), score.view(-1))  "
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda root]",
   "language": "python",
   "name": "conda-root-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}