sequence_classification_lstm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "colab_type": "code",
    "id": "CbU_OgLCVN_0",
    "outputId": "47a3c006-7c1e-4b0e-8f1d-1bf4cd528a6d"
   },
   "outputs": [],
   "source": [
    "# download the Large IMDB Movie Review Dataset\n",
    "# the task is binary classification: positive or negative review\n",
    "\n",
    "! wget http://ai.stanford.edu/%7Eamaas/data/sentiment/aclImdb_v1.tar.gz\n",
    "! tar -xzf aclImdb_v1.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mif0OHyAVcqM"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import math\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torch import optim\n",
    "import os\n",
    "from collections import namedtuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(42)\n",
    "torch.manual_seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yvd2UgrY3lSU"
   },
   "outputs": [],
   "source": [
    "# let's set some parameters\n",
    "\n",
    "train_path = \"aclImdb/train/\" \n",
    "test_path = \"aclImdb/test/\"\n",
    "\n",
    "batch_size = 100\n",
    "max_len = 300\n",
    "embedding_size = 300\n",
    "min_count = 2\n",
    "device = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ckgNJjRgqW2k"
   },
   "source": [
    "## Load the dataset\n",
    "- 25000 train and test sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_jphx0DE3iOO"
   },
   "outputs": [],
   "source": [
    "Sentence = namedtuple('Sentence', ['index', 'tokens', 'label'])\n",
    "\n",
    "def read_imdb_movie_dataset(dataset_path):\n",
    "\n",
    "    indices = []\n",
    "    text = []\n",
    "    rating = []\n",
    "\n",
    "    i = 0\n",
    "\n",
    "    for filename in os.listdir(os.path.join(dataset_path, \"pos\")):\n",
    "        file_path = os.path.join(dataset_path, \"pos\", filename)\n",
    "        data = open(file_path, 'r', encoding=\"ISO-8859-1\").read()\n",
    "        indices.append(i)\n",
    "        text.append(data)\n",
    "        rating.append(1)\n",
    "        i = i + 1\n",
    "\n",
    "    for filename in os.listdir(os.path.join(dataset_path, \"neg\")):\n",
    "        file_path = os.path.join(dataset_path, \"neg\", filename)\n",
    "        data = open(file_path, 'r', encoding=\"ISO-8859-1\").read()\n",
    "        indices.append(i)\n",
    "        text.append(data)\n",
    "        rating.append(0)\n",
    "        i = i + 1\n",
    "\n",
    "    sentences = [ Sentence(index, text.split(), rating)\n",
    "                  for index, text, rating in zip(indices, text, rating)]\n",
    "\n",
    "    return sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "colab_type": "code",
    "id": "CNpRcBNj3pxA",
    "outputId": "fbf0cb7a-8a48-49d4-fe18-7fdb22603110"
   },
   "outputs": [],
   "source": [
    "train_examples = read_imdb_movie_dataset(train_path)\n",
    "test_examples = read_imdb_movie_dataset(test_path)\n",
    "\n",
    "print(len(train_examples))\n",
    "print(len(train_examples))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VWv1Lme8qhQW"
   },
   "source": [
    "## Mapping our words to unique identifiers: the Vocabulary object\n",
    "- We will create an object to manage a mapping between words (or more generally tokens) and unique indices. \n",
    "- There are a few special symbols that we will be adding to handle special cases.\n",
    "  - The first key special case is the `UNK` token, wich will represent all tokens that we do not have in our vocabulary. This is needed as we will build our vocabulary only using the training examples, and during validation or testing (or if we deploy our model in production) we may encounter new words that also need to be represented somehow.\n",
    "  - The `PAD` token, which we will use to create even-sized batches of sentences of different length (more on this below). \n",
    "  - The beginning-of-sentence or `BOS` token, which we may use to denote the beginning of a sentence in some special cases\n",
    "  - The end-of-sentence or `EOS` token, which as in the previous case is useful for certain tasks.\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-iLmpHZEVJHY"
   },
   "outputs": [],
   "source": [
    "# Define the string of special tokens we will need \n",
    "UNK = '<UNK>'\n",
    "PAD = '<PAD>'\n",
    "BOS = '<BOS>'\n",
    "EOS = '<EOS>'\n",
    "\n",
    "\n",
    "class VocabItem:\n",
    "\n",
    "    def __init__(self, string, hash=None):\n",
    "        \"\"\"\n",
    "        Our token object, representing a term in our vocabulary.\n",
    "        \"\"\"\n",
    "        self.string = string\n",
    "        self.count = 0\n",
    "        self.hash = hash\n",
    "\n",
    "    def __str__(self):\n",
    "        \"\"\"\n",
    "        For pretty-printing of our object\n",
    "        \"\"\"\n",
    "        return 'VocabItem({})'.format(self.string)\n",
    "\n",
    "    def __repr__(self):\n",
    "        \"\"\"\n",
    "        For pretty-printing of our object\n",
    "        \"\"\"\n",
    "        return self.__str__()\n",
    "\n",
    "\n",
    "class Vocab:\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        min_count=0,\n",
    "        no_unk=False,\n",
    "        add_padding=False,\n",
    "        add_bos=False,\n",
    "        add_eos=False,\n",
    "        unk=None):\n",
    "\n",
    "        \"\"\"\n",
    "        :param min_count: The minimum frequency count threshold for a token\n",
    "                          to be added to our mapping. Only useful if\n",
    "                          the unk parameter is None.\n",
    "\n",
    "        :param add_padding: If we should add the special `PAD` token.\n",
    "\n",
    "        :param add_bos: If we should add the special `BOS` token.\n",
    "\n",
    "        :param add_eos: If we should add the special `EOS` token.\n",
    "\n",
    "        :param no_unk: If we should not add the `UNK` token to our Vocab.\n",
    "\n",
    "        :param unk: A string with the unknown token, in case our\n",
    "                    sentences have already been processed for this,\n",
    "                    or `None` to use our default `UNK` token.\n",
    "        \"\"\"\n",
    "\n",
    "        self.no_unk = no_unk\n",
    "        self.vocab_items = []\n",
    "        self.vocab_hash = {}\n",
    "        self.word_count = 0\n",
    "        self.special_tokens = []\n",
    "        self.min_count = min_count\n",
    "        self.add_padding = add_padding\n",
    "        self.add_bos = add_bos\n",
    "        self.add_eos = add_eos\n",
    "        self.unk = unk\n",
    "\n",
    "        self.UNK = None\n",
    "        self.PAD = None\n",
    "        self.BOS = None\n",
    "        self.EOS = None\n",
    "\n",
    "        self.index2token = []\n",
    "        self.token2index = {}\n",
    "\n",
    "        self.finished = False\n",
    "\n",
    "    def add_tokens(self, tokens):\n",
    "        if self.finished:\n",
    "            raise RuntimeError('Vocabulary is finished')\n",
    "\n",
    "        for token in tokens:\n",
    "            if token not in self.vocab_hash:\n",
    "                self.vocab_hash[token] = len(self.vocab_items)\n",
    "                self.vocab_items.append(VocabItem(token))\n",
    "\n",
    "            self.vocab_items[self.vocab_hash[token]].count += 1\n",
    "            self.word_count += 1\n",
    "\n",
    "    def finish(self):\n",
    "\n",
    "        token2index = self.token2index\n",
    "        index2token = self.index2token\n",
    "\n",
    "        tmp = []\n",
    "\n",
    "        if not self.no_unk:\n",
    "\n",
    "            # we add/handle the special `UNK` token\n",
    "            # and set it to have index 0 in our mapping\n",
    "            if self.unk:\n",
    "                self.UNK = VocabItem(self.unk, hash=0)\n",
    "                self.UNK.count = self.vocab_items[self.vocab_hash[self.unk]].count\n",
    "                index2token.append(self.UNK)\n",
    "                self.special_tokens.append(self.UNK)\n",
    "\n",
    "                for token in self.vocab_items:\n",
    "                    if token.string != self.unk:\n",
    "                        tmp.append(token)\n",
    "\n",
    "            else:\n",
    "                self.UNK = VocabItem(UNK, hash=0)\n",
    "                index2token.append(self.UNK)\n",
    "                self.special_tokens.append(self.UNK)\n",
    "\n",
    "                for token in self.vocab_items:\n",
    "                    if token.count <= self.min_count:\n",
    "                        self.UNK.count += token.count\n",
    "                    else:\n",
    "                        tmp.append(token)\n",
    "        else:\n",
    "            for token in self.vocab_items:\n",
    "                tmp.append(token)\n",
    "\n",
    "        # we sort our vocab. items by frequency\n",
    "        # so for the same corpus, the indices of our words\n",
    "        # are always the same\n",
    "        tmp.sort(key=lambda token: token.count, reverse=True)\n",
    "\n",
    "        # we always add our additional special tokens\n",
    "        # at the end of our mapping\n",
    "        if self.add_bos:\n",
    "            self.BOS = VocabItem(BOS)\n",
    "            tmp.append(self.BOS)\n",
    "            self.special_tokens.append(self.BOS)\n",
    "\n",
    "        if self.add_eos:\n",
    "            self.EOS = VocabItem(EOS)\n",
    "            tmp.append(self.EOS)\n",
    "            self.special_tokens.append(self.EOS)\n",
    "\n",
    "        if self.add_padding:\n",
    "            self.PAD = VocabItem(PAD)\n",
    "            tmp.append(self.PAD)\n",
    "            self.special_tokens.append(self.PAD)\n",
    "\n",
    "        index2token += tmp\n",
    "\n",
    "        # we update the vocab_hash for each\n",
    "        # VocabItem object in our list\n",
    "        # based on their frequency\n",
    "        for i, token in enumerate(self.index2token):\n",
    "            token2index[token.string] = i\n",
    "            token.hash = i\n",
    "\n",
    "        self.index2token = index2token\n",
    "        self.token2index = token2index\n",
    "\n",
    "        if not self.no_unk:\n",
    "            print('Unknown vocab size:', self.UNK.count)\n",
    "\n",
    "        print('Vocab size: %d' % len(self))\n",
    "\n",
    "        self.finished = True\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        return self.index2token[i]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.index2token)\n",
    "\n",
    "    def __iter__(self):\n",
    "        return iter(self.index2token)\n",
    "\n",
    "    def __contains__(self, key):\n",
    "        return key in self.token2index\n",
    "\n",
    "    def tokens2indices(self, tokens, add_bos=False, add_eos=False):\n",
    "        \"\"\"\n",
    "        Returns a list of mapping indices by processing the given string\n",
    "        with our `tokenizer` and `token_function`, and defaulting to our\n",
    "        special `UNK` token whenever we found an unseen term.\n",
    "\n",
    "        :param string: A sentence string we wish to map into our vocabulary.\n",
    "\n",
    "        :param add_bos: If we should add the `BOS` at the beginning.\n",
    "\n",
    "        :param add_eos: If we should add the `EOS` at the end.\n",
    "\n",
    "        :return: A list of ints, with the indices of each token in the\n",
    "                given string.\n",
    "        \"\"\"\n",
    "        string_seq = []\n",
    "        if add_bos:\n",
    "            string_seq.append(self.BOS.hash)\n",
    "        for token in tokens:\n",
    "            if self.no_unk:\n",
    "                string_seq.append(self.token2index[token])\n",
    "            else:\n",
    "                string_seq.append(self.token2index.get(token, self.UNK.hash))\n",
    "        if add_eos:\n",
    "            string_seq.append(self.EOS.hash)\n",
    "        return string_seq\n",
    "\n",
    "    def indices2tokens(self, indices, ignore_ids=()):\n",
    "        \"\"\"\n",
    "        Returns a list of strings by mapping back every index to our\n",
    "        vocabulary.\n",
    "\n",
    "        :param indices: A list of ints.\n",
    "\n",
    "        :param ignore_ids: An itereable with indices to ignore, meaning\n",
    "                           that we will not look for them in our mapping.\n",
    "\n",
    "        :return: A list of strings.\n",
    "\n",
    "        Will raise a KeyException whenever we pass an index that we\n",
    "        do not have in our mapping, except when provided with `ignore_ids`.\n",
    "\n",
    "        \"\"\"\n",
    "        tokens = []\n",
    "        for idx in indices:\n",
    "            if idx in ignore_ids:\n",
    "                continue\n",
    "            tokens.append(self.index2token[idx].string)\n",
    "\n",
    "        return tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VygQN-dXAeoA"
   },
   "source": [
    "- Now we can instance our vocabulary objects and add the data.\n",
    "- We will use one vocabulary for the input data (the sentences), and another vocabulary object for the output data, the class labels. In this way our code is generic and should work out-of-the-box for any number of output labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8wDnRfy1fp-o"
   },
   "outputs": [],
   "source": [
    "# for the input vocabulary\n",
    "# we set a minimum frequency, therefore adding the `UNK` special token\n",
    "# and we also add the `PAD` special token, as we will need it later  \n",
    "src_vocab = Vocab(min_count=min_count, add_padding=True)\n",
    "\n",
    "# for the output vocabulary\n",
    "# we do not need the `UNK` token (we know all the classes), or the `PAD`\n",
    "tgt_vocab = Vocab(no_unk=True, add_padding=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 71
    },
    "colab_type": "code",
    "id": "Dr_C7CAr5Enw",
    "outputId": "52654ce8-c420-49cb-cb59-a13f0c9c3513"
   },
   "outputs": [],
   "source": [
    "for sentence in train_examples:\n",
    "    src_vocab.add_tokens(sentence.tokens[:max_len])\n",
    "    tgt_vocab.add_tokens([sentence.label])\n",
    "\n",
    "src_vocab.finish()\n",
    "tgt_vocab.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "CDcHKXM3DT4K",
    "outputId": "19850996-bfcb-4927-f5cf-b5b3d54bcce1"
   },
   "outputs": [],
   "source": [
    "src_vocab.tokens2indices('the movie was bad'.split())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KfsFW2h27k0O"
   },
   "outputs": [],
   "source": [
    "Vocabs = namedtuple('Vocabs', ['src', 'tgt'])\n",
    "vocabs = Vocabs(src_vocab, tgt_vocab)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mSGDb0MH9A_G"
   },
   "source": [
    "## Representing words using sparse vectors: Word Embeddings\n",
    "- One of the major breakthroughs in NLP with deep models came after the conception of word embeddings, which changed the way in which we represent each word in our machine learning models.\n",
    "- We start by simply assigning an initially random vector to each word in our vocabulary.These vectors are stacked together into a big matrix, usually referred to as the *embedding* matrix. After we have built our vocabulary, all we have to do is to create a big tensor of shape (`vocab_size`, `embedding_size`).\n",
    "- In theory, whenever we need to obtain the vector for a given word, we could build a one-hot vector of our word and multply this vector by our *embedding* matrix. All but one value in this one-hot vector are zeroes, the result of this product will correspond exactly to the vector that represents our word.\n",
    "- Our *embeddings* will be treated as parameters of our models and are trained with it. This is possible because the *embedding* mechanism as has a well-defined derivative, so we are  allowed to use backpropagation to train these vectors.\n",
    "- Note that in practice, however, the one-hot-based behavior can be achieved by simply selecting row vectors from our *embedding* matrix, given our indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wF9Ef_2n82ny"
   },
   "outputs": [],
   "source": [
    "embeddings = nn.Embedding(\n",
    "    len(src_vocab),\n",
    "    embedding_size,\n",
    "    padding_idx=src_vocab.PAD.hash\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "4MRXrQ66J230",
    "outputId": "18ba3578-e90f-4f61-ea51-e8a30fbbbad3"
   },
   "outputs": [],
   "source": [
    "print(embeddings.weight.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Batch objects\n",
    " - To easily access all the data in a batch, let's create a special Batch object that will give us access to all the information we may require during training.\n",
    " - Let's begin creating a more friendly object that contains a numeric representation of our inputs and outputs.\n",
    "- By default we will use numpy objects, but we will also add a function to translate the contents of the object to PyTorch.\n",
    "- We will create this object to be generic enough so we can use it with tasks other than classification, too. \n",
    " - This object will work like a dictionary, but it will also allow us to access each component using an attribute with the same name.\n",
    "  The main principle is that this dictionary-like batch will hold `numpy` objects as values, and that after calling the `to_torch_()` function, they will be turned into `pytorch` objects and moved to the corresponding provided device. In this way, we know that all our elements inside the batch object are in the right place.\n",
    " - We will combine our `Batch` object with a `BatchTuple` object that will hold data relevant to a specific input of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Batch(dict):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super(Batch, self).__init__(*args, **kwargs)\n",
    "        self.__dict__ = self\n",
    "        self._is_torch = False\n",
    "\n",
    "    def to_torch_(self, device):\n",
    "        self._is_torch = False\n",
    "        for key in self.keys():\n",
    "            value = self[key]\n",
    "            # we move `numpy` objects to `pytorch`\n",
    "            if isinstance(value, BatchTuple):\n",
    "                value.to_torch_(device)\n",
    "            # we also move our BatchTuple objects to `pytorch`\n",
    "            if isinstance(value, np.ndarray):\n",
    "                self[key] = torch.from_numpy(value).to(device)\n",
    "\n",
    "\n",
    "class BatchTuple(object):\n",
    "    def __init__(self, sequences, lengths, sublengths, masks):\n",
    "        self.sequences = sequences\n",
    "        self.lengths = lengths\n",
    "        self.sublengths = sublengths\n",
    "        self.masks = masks\n",
    "        self._is_torch = False\n",
    "\n",
    "    def to_torch_(self, device):\n",
    "        if not self._is_torch:\n",
    "            self.sequences = torch.tensor(\n",
    "                self.sequences, device=device, dtype=torch.long\n",
    "            )\n",
    "\n",
    "            if self.lengths is not None:\n",
    "                self.lengths = torch.tensor(\n",
    "                    self.lengths, device=device, dtype=torch.long\n",
    "                )\n",
    "\n",
    "            if self.sublengths is not None:\n",
    "                self.sublengths = torch.tensor(\n",
    "                    self.sublengths, device=device, dtype=torch.long\n",
    "                )\n",
    "            if self.masks is not None:\n",
    "                self.masks = torch.tensor(\n",
    "                    self.masks, device=device, dtype=torch.float\n",
    "                )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gBhJPDdsDNdJ"
   },
   "source": [
    "### The padding function\n",
    "- Let's suppose we have these two sentences to build a batch:\n",
    "  - the dog barks $\\rightarrow [1, 2 ,3]$\n",
    "  - the cat likes to sleep $\\rightarrow [1, 4, 5, 6, 7]$\n",
    "  \n",
    "  In order to put these two examples in a batch Tensor, we will need to *pad* the shortest sentence to have the same length of the longest one. \n",
    "  - the dog barks $\\rightarrow [1, 2 ,3, 0 , 0]$\n",
    "  - the cat likes to sleep $\\rightarrow [1, 4, 5, 6, 7]$\n",
    "  \n",
    "  Finally, our batch Tensor will look like this: \n",
    "  - $\\begin{bmatrix}1 & 2 & 3 & 0 & 0 \\\\ 1 & 4 & 5 & 6 & 7\\end{bmatrix}$\n",
    "  \n",
    "  where its first dimension represents the size of the batch, and its second dimension has the length of the longest sentence in our batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "J8eXpGbS8BJW"
   },
   "outputs": [],
   "source": [
    "def pad_list(\n",
    "    sequences,\n",
    "    dim0_pad=None,\n",
    "    dim1_pad=None,\n",
    "    align_right=False,\n",
    "    pad_value=0\n",
    "):\n",
    "    \"\"\"\n",
    "    Receives a list of lists and returns a padded 2d ndarray,\n",
    "    and a list of lengths. \n",
    "    \n",
    "    sequences: a list of lists. len(sequences) = M, and N is the max\n",
    "               length of any of the lists contained in sequences.\n",
    "               e.g.: [[2,45,3,23,54], [12,4,2,2], [4], [45, 12]]\n",
    "   \n",
    "    Returns a numpy ndarray of dimension (M, N) corresponding to the padded\n",
    "    sequences and a list of the original lengths.\n",
    "    \n",
    "    Returns:\n",
    "       - out: a torch tensor of dimension (M, N) \n",
    "       - lengths: a list of ints containing the lengths of each element\n",
    "                  in sequences\n",
    "       \n",
    "    \"\"\"\n",
    "    \n",
    "    sequences = [np.asarray(sublist) for sublist in sequences]\n",
    "\n",
    "    if not dim0_pad:\n",
    "        dim0_pad = len(sequences)\n",
    "\n",
    "    if not dim1_pad:\n",
    "        dim1_pad = max(len(seq) for seq in sequences)\n",
    "\n",
    "    out = np.full(shape=(dim0_pad, dim1_pad), fill_value=pad_value)\n",
    "\n",
    "    lengths = []\n",
    "    for i in range(len(sequences)):\n",
    "        data_length = len(sequences[i])\n",
    "        lengths.append(data_length)\n",
    "        offset = dim1_pad - data_length if align_right else 0\n",
    "        np.put(out[i], range(offset, offset + data_length), sequences[i])\n",
    "\n",
    "    lengths = np.array(lengths)\n",
    "\n",
    "    return out, lengths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The BatchBuilder object\n",
    "- On top of the Batch object we create our own SequenceClassificationBatchBuilder which will be the in charge of transforming input raw examples into a collection of our Batch objects that our model can handle.\n",
    "- This object will do all the heavy-lifting, turning our string examples into our batch objects, which PyTorch can later handle.\n",
    "- We will combine this object with the `DataLoader` util from `pytorch`, using as a function for the [`collate_fn` parameter](https://pytorch.org/docs/stable/data.html#working-with-collate-fn), which allows us to provide a custom function to create this funcion. In our case, this is achieved by implementing the `__call__` function in the `BatchBuilder` object, which will esentally turn the [instance into a function](https://docs.python.org/3/reference/datamodel.html#emulating-callable-objects)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SequenceClassificationBatchBuilder(object):\n",
    "    # Because the `__call__` function needs to only recieve \n",
    "    # one parameter (due to restrictions of the `DataLoader`\n",
    "    # we can use the constructor we can pass any additional\n",
    "    # inputs we may require when building our batches\n",
    "    def __init__(self, vocabs, max_len=None):\n",
    "        self.vocabs = vocabs\n",
    "        self.max_len = max_len\n",
    "    \n",
    "    # This will the function called by the `DataLoader` object\n",
    "    # that only accepts the `examples` parameter\n",
    "    def __call__(self, examples):\n",
    "\n",
    "        ids_batch = [int(sentence.index) for sentence in examples]\n",
    "\n",
    "        src_examples = [\n",
    "            self.vocabs.src.tokens2indices(sentence.tokens[: self.max_len])\n",
    "            for sentence in examples\n",
    "        ]\n",
    "\n",
    "        tgt_examples = [\n",
    "            self.vocabs.tgt.token2index[sentence.label] for sentence in examples\n",
    "        ]\n",
    "\n",
    "        src_padded, src_lengths = pad_list(\n",
    "            src_examples, pad_value=self.vocabs.src.PAD.hash\n",
    "        )\n",
    "\n",
    "        src_batch_tuple = BatchTuple(src_padded, src_lengths, None, None)\n",
    "\n",
    "        tgt_batch_tuple = BatchTuple(tgt_examples, None, None, None)\n",
    "\n",
    "        return Batch(\n",
    "            indices=ids_batch, src=src_batch_tuple, tgt=tgt_batch_tuple\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "l1V-6QNUDnFA"
   },
   "source": [
    "Let's instance our `batch_builder`, feed it into the `DataLoader` object alongside the  training and test examples, and let's inspect a single batch of examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OQTU7ZG691DG"
   },
   "outputs": [],
   "source": [
    "batch_builder = SequenceClassificationBatchBuilder(\n",
    "    vocabs, max_len=max_len\n",
    ")\n",
    "\n",
    "train_batches = DataLoader(\n",
    "    train_examples,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    num_workers=0,\n",
    "    collate_fn=batch_builder,\n",
    ")\n",
    "\n",
    "test_batches = DataLoader(\n",
    "    test_examples,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False,\n",
    "    num_workers=0,\n",
    "    collate_fn=batch_builder,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wp4GHiIQB45s"
   },
   "outputs": [],
   "source": [
    "train_batches_iter = iter(train_batches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mXk_zXIZC451"
   },
   "outputs": [],
   "source": [
    "train_batch = next(train_batches_iter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 143
    },
    "colab_type": "code",
    "id": "qsZBQdp19m6b",
    "outputId": "e7305591-43b6-4d7c-9436-3c0acc1f002b"
   },
   "outputs": [],
   "source": [
    "train_batch.src.sequences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KxR6muZJD1sG"
   },
   "source": [
    "## The Pytorch Model\n",
    "### The LSTM\n",
    "![An unrolled RNN.](http://colah.github.io/posts/2015-08-Understanding-LSTMs/img/RNN-unrolled.png)\n",
    "- The LSTM is a special kind of Recurrent Neural Network that will process sequence data and return a vector for each input in our sequence. In the example, given a sequence of inputs $X=x_1, \\ldots , x_t$, the LSTM will give us a sequence of $t$ vectors also called hidden states $H= h_1, \\ldots, h_t$.\n",
    "- The LSTM is a complex beast, in this tutorial we will be skipping details on how exactly it works. For more details, visit http://pytorch.org/docs/master/nn.html#lstm\n",
    "- If we think of our input sequence as our word vectors for a given sentence, we can think of the output as a kind of enriched or contextualized version of the input, which will contain not only information about the word each vector represents, but also about its previous words.\n",
    "- In PyTorch, LSTMs will return both the set of output vectors $H$ but also some additional output that we will not pay attention to.\n",
    "- Because we need a fixed-size vector to classify our sentences, we will have to use some kind of pooling function over our hidden states to achieve this. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hqnxzQFEJ_xW"
   },
   "outputs": [],
   "source": [
    "def mean_pooling(batch_hidden_states, batch_lengths):\n",
    "    '''\n",
    "    :param batch_hidden_states: torch.Tensor(batch_size, seq_len, hidden_size)\n",
    "    :param batch_lengths: list(batch_size)\n",
    "    :return:\n",
    "    '''\n",
    "    batch_lengths = batch_lengths.float()\n",
    "    batch_lengths = batch_lengths.unsqueeze(1)\n",
    "    if batch_hidden_states.is_cuda:\n",
    "        batch_lengths = batch_lengths.cuda()\n",
    "\n",
    "    pooled_batch = torch.sum(batch_hidden_states, 1)\n",
    "    pooled_batch = pooled_batch / batch_lengths.expand_as(pooled_batch)\n",
    "\n",
    "    return pooled_batch\n",
    "\n",
    "\n",
    "def max_pooling(batch_hidden_states):\n",
    "    '''\n",
    "    :param batch_hidden_states: torch.Tensor(batch_size, seq_len, hidden_size)\n",
    "    :return:\n",
    "    '''\n",
    "    pooled_batch, _ = torch.max(batch_hidden_states, 1)\n",
    "    return pooled_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aLc-X3xDG9NW"
   },
   "source": [
    "- The next key util functions are related to the fact that we are using batches of sentences to train.\n",
    "- To make the training efficient, Pytorch asks us to sort the examples in our batch by sequence length and build a special object.\n",
    "- We will use the function `pack_padded_sequence()` to build this special `PackedSequence` object given our sorted padded batch and the lengths of each sentence on it\n",
    "- Conversely, we will use the `pad_packed_sequence()` function to turn the output of the `nn.LSTM`, a `PackedSequence` object, into a regular Pytorch tensor. This tensor will have zeroes in all padding positions, so we can later directy use our pooling functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cCyAAJaCG9ey"
   },
   "outputs": [],
   "source": [
    "def pack_rnn_input(embedded_sequence_batch, sequence_lengths):\n",
    "    \"\"\"\n",
    "    Prepares the special `PackedSequence` object that can be\n",
    "    efficiently processed by the `nn.LSTM`.\n",
    "\n",
    "    :param embedded_sequence_batch: torch.Tensor(seq_len, batch_size)\n",
    "\n",
    "    :param sequence_lengths: list(batch_size)\n",
    "\n",
    "    :return:\n",
    "      - `PackedSequence` object containing our padded batch\n",
    "      - indices to sort back our sentences to their original order\n",
    "    \"\"\"\n",
    "\n",
    "    sequence_lengths = sequence_lengths.cpu().numpy()\n",
    "\n",
    "    sorted_sequence_lengths = np.sort(sequence_lengths)[::-1]\n",
    "    sorted_sequence_lengths = torch.from_numpy(\n",
    "        sorted_sequence_lengths.copy()\n",
    "    )\n",
    "\n",
    "    idx_sort = np.argsort(-sequence_lengths)\n",
    "    idx_unsort = np.argsort(idx_sort)\n",
    "\n",
    "    idx_sort = torch.from_numpy(idx_sort)\n",
    "    idx_unsort = torch.from_numpy(idx_unsort)\n",
    "\n",
    "    if embedded_sequence_batch.is_cuda:\n",
    "        idx_sort = idx_sort.cuda()\n",
    "        idx_unsort = idx_unsort.cuda()\n",
    "\n",
    "    embedded_sequence_batch = embedded_sequence_batch.index_select(\n",
    "        0, idx_sort\n",
    "    )\n",
    "\n",
    "    # Handling padding in Recurrent Networks\n",
    "    packed_rnn_input = nn.utils.rnn.pack_padded_sequence(\n",
    "        embedded_sequence_batch, \n",
    "        sorted_sequence_lengths,\n",
    "        batch_first=True\n",
    "    )\n",
    "\n",
    "    return packed_rnn_input, idx_unsort\n",
    "\n",
    "  \n",
    "def unpack_rnn_output(packed_rnn_output, indices):\n",
    "    \"\"\"\n",
    "     Recover a regular tensor given a `PackedSequence` as returned\n",
    "     by  `nn.LSTM`\n",
    "\n",
    "    :param packed_rnn_output: torch object\n",
    "\n",
    "    :param indices: Variable(LongTensor) of indices to sort output\n",
    "\n",
    "    :return:\n",
    "      - Padded tensor\n",
    "\n",
    "    \"\"\"\n",
    "    encoded_sequence_batch, _ = nn.utils.rnn.pad_packed_sequence(\n",
    "        packed_rnn_output, batch_first=True\n",
    "    )\n",
    "\n",
    "    encoded_sequence_batch = encoded_sequence_batch.index_select(0, indices)\n",
    "\n",
    "    return encoded_sequence_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "e70TNEbkANUI"
   },
   "source": [
    "- To build the model, we extend the `nn.Module`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6Bf9pMNtdHWo"
   },
   "outputs": [],
   "source": [
    "class BiLSTM(nn.Module):\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        embeddings,\n",
    "        hidden_size,\n",
    "        num_labels,\n",
    "        input_dropout=0,\n",
    "        output_dropout=0,\n",
    "        bidirectional=True,\n",
    "        num_layers=2,\n",
    "        pooling='mean'\n",
    "    ):\n",
    "\n",
    "        super(BiLSTM, self).__init__()\n",
    "\n",
    "        self.embeddings = embeddings\n",
    "        self.pooling = pooling\n",
    "\n",
    "        self.input_dropout = nn.Dropout(input_dropout)\n",
    "        self.output_dropout = nn.Dropout(output_dropout)\n",
    "\n",
    "        self.bidirectional = bidirectional\n",
    "        self.num_layers = num_layers\n",
    "        self.num_labels = num_labels\n",
    "\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        self.input_size = self.embeddings.embedding_dim\n",
    "\n",
    "        self.lstm = nn.LSTM(\n",
    "            self.input_size,\n",
    "            hidden_size,\n",
    "            bidirectional=bidirectional,\n",
    "            num_layers=num_layers,\n",
    "            batch_first=True\n",
    "        )\n",
    "\n",
    "        self.total_hidden_size = self.hidden_size \n",
    "        if self.bidirectional:\n",
    "            self.total_hidden_size += self.hidden_size\n",
    "\n",
    "        self.output_layer = nn.Linear(\n",
    "            self.total_hidden_size,\n",
    "            self.num_labels)\n",
    "\n",
    "        self.loss_function = nn.CrossEntropyLoss()\n",
    "\n",
    "        \n",
    "    def forward(self, src_batch, tgt_batch=None):\n",
    "\n",
    "        src_sequences = src_batch.sequences\n",
    "        src_lengths = src_batch.lengths\n",
    "\n",
    "        embedded_sequence_batch = self.embeddings(src_sequences)\n",
    "        embedded_sequence_batch = self.input_dropout(\n",
    "            embedded_sequence_batch\n",
    "        )\n",
    "\n",
    "        packed_rnn_input, indices = pack_rnn_input(\n",
    "            embedded_sequence_batch, src_lengths\n",
    "        )\n",
    "\n",
    "        rnn_packed_output, _ = self.lstm(packed_rnn_input)\n",
    "        encoded_sequence_batch = unpack_rnn_output(\n",
    "            rnn_packed_output, indices\n",
    "        )\n",
    "\n",
    "        if self.pooling == \"mean\":\n",
    "            # batch_size, hidden_x_dirs\n",
    "            pooled_batch = mean_pooling(encoded_sequence_batch,\n",
    "                                        src_lengths)\n",
    "\n",
    "        elif self.pooling == \"max\":\n",
    "            # batch_size, hidden_x_dirs\n",
    "            pooled_batch = max_pooling(encoded_sequence_batch)\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "\n",
    "        logits = self.output_layer(pooled_batch)\n",
    "        _, predictions = logits.max(1)\n",
    "\n",
    "        if tgt_batch is not None:\n",
    "            targets = tgt_batch.sequences\n",
    "            loss = self.loss_function(logits, targets)\n",
    "        else:\n",
    "            loss = None\n",
    "\n",
    "        return loss, predictions, logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ipl1kOiKEkac"
   },
   "source": [
    "### Instancing our model\n",
    "- Let's define the hyperparameters of our model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0qoZmYg689qG"
   },
   "outputs": [],
   "source": [
    "epochs = 10\n",
    "hidden_size = 300\n",
    "log_interval = 10\n",
    "num_labels = 2\n",
    "input_dropout = 0.5\n",
    "output_dropout = 0.5\n",
    "bidirectional = True\n",
    "num_layers = 2\n",
    "pooling = 'mean'\n",
    "lr = 0.001\n",
    "gradient_clipping = 0.25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 160
    },
    "colab_type": "code",
    "id": "4x850E1O9kM2",
    "outputId": "d5be47a0-ea20-4771-8da4-6f018ae27fec"
   },
   "outputs": [],
   "source": [
    "model = BiLSTM(\n",
    "    embeddings=embeddings,\n",
    "    hidden_size=hidden_size,\n",
    "    num_labels=num_labels,\n",
    "    input_dropout=input_dropout,\n",
    "    output_dropout=output_dropout,\n",
    "    bidirectional=bidirectional,\n",
    "    num_layers=num_layers,\n",
    "    pooling=pooling\n",
    ")\n",
    "\n",
    "model.to(device)\n",
    "    \n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rBnvsLNC91Ak"
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1034
    },
    "colab_type": "code",
    "id": "yaJjF4H2908k",
    "outputId": "b1b9894b-f07a-46ca-a004-46ad690b3e38"
   },
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    \n",
    "    epoch_correct = 0\n",
    "    epoch_total = 0\n",
    "    epoch_loss = 0\n",
    "    i = 0\n",
    "\n",
    "    model.train()\n",
    "\n",
    "    for batch in train_batches:\n",
    "\n",
    "        batch.to_torch_(device)\n",
    "\n",
    "        ids_batch = batch.indices\n",
    "        src_batch = batch.src\n",
    "        tgt_batch = batch.tgt\n",
    "\n",
    "        loss, predictions, logits = model.forward(\n",
    "            src_batch,\n",
    "            tgt_batch=tgt_batch\n",
    "        )\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        torch.nn.utils.clip_grad_norm_(\n",
    "            model.parameters(),\n",
    "            gradient_clipping)\n",
    "\n",
    "        optimizer.step()\n",
    "        correct = (predictions == tgt_batch.sequences).long().sum()\n",
    "        total = tgt_batch.sequences.size(0)\n",
    "        epoch_correct += correct.item()\n",
    "        epoch_total += total\n",
    "        epoch_loss += loss.item()\n",
    "        i += 1\n",
    "\n",
    "    accuracy  = 100 * epoch_correct / epoch_total\n",
    "\n",
    "    print('Epoch {}'.format(epoch))\n",
    "    print('Train Loss: {}'.format(epoch_loss / len(train_batches)))\n",
    "    print('Train Accuracy: {}'.format(accuracy))\n",
    "\n",
    "    test_epoch_correct = 0\n",
    "    test_epoch_total = 0\n",
    "    test_epoch_loss = 0\n",
    "\n",
    "    model.eval()\n",
    "\n",
    "    for batch in test_batches:\n",
    "\n",
    "        ids_batch = batch.indices\n",
    "        src_batch = batch.src\n",
    "        tgt_batch = batch.tgt\n",
    "        \n",
    "        batch.to_torch_(device)\n",
    "\n",
    "        loss, predictions, logits = model.forward(\n",
    "            src_batch,\n",
    "            tgt_batch=tgt_batch)\n",
    "\n",
    "        correct = (predictions == tgt_batch.sequences).long().sum()\n",
    "        total = tgt_batch.sequences.size(0)\n",
    "        test_epoch_correct += correct.item()\n",
    "        test_epoch_total += total\n",
    "        test_epoch_loss += loss.item()\n",
    "\n",
    "    test_accuracy = 100 * test_epoch_correct / test_epoch_total\n",
    "\n",
    "    print('\\n---------------------')\n",
    "    print('Test Loss: {}'.format(test_epoch_loss / len(test_batches)))\n",
    "    print('Test Accuracy: {}'.format(test_accuracy))\n",
    "    print('---------------------\\n')\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "sequence_classification_lstm.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}