lesson4-lstm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sentiment classification with LSTM\n",
    "In this notebook we will use LSTMs to do sentiment classification on the [imdb dataset](http://ai.stanford.edu/~amaas/data/sentiment/). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "import pandas as pd \n",
    "import os\n",
    "import spacy\n",
    "import string\n",
    "import re\n",
    "import numpy as np\n",
    "from spacy.symbols import ORTH\n",
    "from collections import Counter\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get the data: <br>\n",
    "`wget http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def unpack_dataset():\n",
    "    ! mkdir -p data/aclImdb\n",
    "    ! wget http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\n",
    "    ! tar -zxvf aclImdb_v1.tar.gz -C data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#unpack_dataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[PosixPath('data/aclImdb/imdbEr.txt'),\n",
       " PosixPath('data/aclImdb/imdb.vocab'),\n",
       " PosixPath('data/aclImdb/train'),\n",
       " PosixPath('data/aclImdb/test'),\n",
       " PosixPath('data/aclImdb/README')]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "PATH = Path(\"data/aclImdb/\")\n",
    "list(PATH.iterdir())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Bromwell High is a cartoon comedy. It ran at the same time as some other programs about school life, such as \"Teachers\". My 35 years in the teaching profession lead me to believe that Bromwell High\\'s satire is much closer to reality than is \"Teachers\". The scramble to survive financially, the insightful students who can see right through their pathetic teachers\\' pomp, the pettiness of the whole situation, all remind me of the schools I knew and their students. When I saw the episode in which a student repeatedly tried to burn down the school, I immediately recalled ......... at .......... High. A classic line: INSPECTOR: I\\'m here to sack one of your teachers. STUDENT: Welcome to Bromwell High. I expect that many adults of my age think that Bromwell High is far fetched. What a pity that it isn\\'t!'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "path = PATH/\"train/pos/0_9.txt\"\n",
    "path.read_text()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first time run this\n",
    "#!python3 -m spacy download en"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "re_br = re.compile(r'<\\s*br\\s*/?>', re.IGNORECASE)\n",
    "def sub_br(x): return re_br.sub(\"\\n\", x)\n",
    "\n",
    "my_tok = spacy.load('en')\n",
    "def spacy_tok(x): return [tok.text for tok in my_tok.tokenizer(sub_br(x))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Bromwell', 'High', 'is', 'a', 'cartoon', 'comedy', '.', 'It', 'ran', 'at']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "path = PATH/\"train/pos/0_9.txt\"\n",
    "spacy_tok(path.read_text())[:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing vocab2index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[PosixPath('data/aclImdb/train/pos/10544_8.txt'),\n",
       " PosixPath('data/aclImdb/train/pos/9530_9.txt'),\n",
       " PosixPath('data/aclImdb/train/pos/9901_8.txt'),\n",
       " PosixPath('data/aclImdb/train/pos/11951_8.txt'),\n",
       " PosixPath('data/aclImdb/train/pos/7441_7.txt')]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pos_files = list((PATH/\"train\"/\"pos\").iterdir())\n",
    "neg_files = list((PATH/\"train\"/\"neg\").iterdir())\n",
    "all_files = pos_files + neg_files\n",
    "all_files[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "counts = Counter()\n",
    "for path in all_files:\n",
    "    counts.update(spacy_tok(path.read_text()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "103688"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(counts.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "for word in list(counts):\n",
    "    if counts[word] < 5:\n",
    "        del counts[word]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "33914"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(counts.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab2index = {\"\":0, \"UNK\":1}\n",
    "words = [\"\", \"UNK\"]\n",
    "for word in counts:\n",
    "    vocab2index[word] = len(words)\n",
    "    words.append(word)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#vocab2index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# note that spacy_tok takes a while run it just once\n",
    "def encode_sentence(path, vocab2index, N=400, padding_start=True):\n",
    "    x = spacy_tok(path.read_text())\n",
    "    enc = np.zeros(N, dtype=np.int32)\n",
    "    enc1 = np.array([vocab2index.get(w, vocab2index[\"UNK\"]) for w in x])\n",
    "    l = min(N, len(enc1))\n",
    "    if padding_start:\n",
    "        enc[:l] = enc1[:l]\n",
    "    else:\n",
    "        enc[N-l:] = enc1[:l]\n",
    "    return enc, l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([    0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "            0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "            0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "            0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "            0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "            0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "            0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "            0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "            0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "            0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "            1,  1782,     4,  2723,     4,    29,    36,    37,    15,\n",
       "         2388,   320,     6,    23,   351,     8,   316,   254,   680,\n",
       "         1382,   103,   334,  7723,   103, 12825,     3,   369,   951,\n",
       "            3, 16948,    19,  1791,    55,  1573,    36,     1,   153,\n",
       "           16,  9449,  1070,   349,    79,    31,  3864,   118,   179,\n",
       "         3504,    23,    13,    44,  6231,    23,  5043,   214,  3491,\n",
       "          301,   254,  1967,    16,  1749,  3525,   174,    12,    78,\n",
       "         1774,  6181,    19,  1070,    70,    36,  2723,  1070,  2007,\n",
       "        23765,    31,   731,   128,  1031,    55,   372,   174,   316,\n",
       "         2459,    19,  1840,    31,    23,   654,   271,  1635,    19,\n",
       "         6647,    31,   186,    14,    22,    23,  6502,  1184,    13,\n",
       "        14376,    31,   195,   243, 16948,   194, 13986, 14400,    19,\n",
       "          174,    78,   262,   976,    70,    50,    31, 27234,   500,\n",
       "          249,   195,    10,     1,   325, 12859,    22,    78,    68,\n",
       "           84,    55,  7244,    91,    31,  2057,  1205,     4,    50,\n",
       "            4,   316,  1966,   134,    16,  2492,    10, 14699,    44,\n",
       "          350,   117,  3776,    36,  1629,   751,  1291,   153,    96,\n",
       "         3659,  4531,    31,   502,  5259,    31,    40,   127,   128,\n",
       "         1031,    55,   372,    36,    37,    15,  6072, 28537,  5475,\n",
       "           19,   358,    31,    76,    23,  1201,  1727,    47,    22,\n",
       "         1170,    16,  1749,    13,    69,    67,    23,  1555,    68,\n",
       "         1507,  7723, 12825,    31,   186,  1849,    46,  1070,   153,\n",
       "          199,  1070, 10712,    19,   350,    46,  1140,   214,    32,\n",
       "          128, 22123,    95,   214,   276,   143,    55,    16,   992,\n",
       "           22,    23,    13,  3247,   103, 32525,   458,    16,   844,\n",
       "         7470,  1749,   729,    31,     6,   343,   184,    16,  1749,\n",
       "           19,   142,     1,    36,  1289,  2447,    12,  1070,  1669,\n",
       "            8,   171,    72,    36,   908,  7219,   729,  1070,    16,\n",
       "        13986,   751,  1354,     3,    36,  1135,   566,  8803,   745,\n",
       "           19,    16,  7712,  9449,  1549,    31,   794,   142,  1967,\n",
       "         7723, 12825,    71,    10,  1171,    12,    31,    99,    43,\n",
       "          142,   128,   200,   115,    69,  2723,  2832,    59,  1070,\n",
       "          974,    43,    31,  2057, 23617,   196,   174,   564,    36,\n",
       "        18583,  1388,   249,   171,    10,    96,   939,   215,    15,\n",
       "        13481,    59,   174,    31], dtype=int32), 310)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "path = PATH/\"train/neg/211_4.txt\"\n",
    "encode_sentence(path, vocab2index, N=400, padding_start=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImdbDataset(Dataset):\n",
    "    def __init__(self, PATH, train=\"train\", N=400, padding_start=True):\n",
    "        self.path_to_images = PATH/train\n",
    "        self.pos_files = list((self.path_to_images/\"pos\").iterdir())\n",
    "        self.neg_files = list((self.path_to_images/\"neg\").iterdir())\n",
    "        self.files = self.pos_files + self.neg_files\n",
    "        # pos 1, neg 0\n",
    "        self.y = np.concatenate((np.ones(len(self.pos_files), dtype=int),\n",
    "                                np.zeros(len(self.neg_files), dtype=int)), axis=0)\n",
    "        # it is important to run encode_sentence in the init\n",
    "        self.X = [encode_sentence(path, vocab2index, N, padding_start) for path in self.files]\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.y)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        x, s = self.X[idx]\n",
    "        return x, s, self.y[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds_v0 = ImdbDataset(PATH, padding_start=False)\n",
    "valid_ds_v0 = ImdbDataset(PATH, \"test\", padding_start=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 1000\n",
    "train_dl_v0 = DataLoader(train_ds_v0, batch_size=batch_size, shuffle=True)\n",
    "valid_dl_v0 = DataLoader(valid_ds_v0, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
       "          0,   0,   0,   0, 181, 182, 183, 184,  91,  36, 185,  31, 186,\n",
       "         50, 187, 188, 189,  84, 190, 191,  91, 174,  12, 192, 193, 194,\n",
       "         31, 195, 196, 197, 198,  70, 174,  78,  14,  15,  16, 199, 103,\n",
       "        200, 201, 202, 142, 203,  25, 204,  31,  76, 163, 164,   3,  16,\n",
       "        205, 206,  68, 207,  27, 142, 208, 209, 210,  19, 211, 212,  95,\n",
       "         16, 213, 214, 186, 215,  15, 216,  59,  71,  78,  34, 217,   3,\n",
       "        218, 219, 220, 218, 160, 221, 220, 218, 222,  31,  76, 223,  15,\n",
       "         16,  54,  59,  71,   3, 224,   3,  72,   4,  69, 225,   4,  70,\n",
       "        210,  78,  36, 226, 227, 103, 228, 115, 229, 230, 212,  31, 210,\n",
       "         85, 231, 232, 233,  78, 220, 234,  84, 235, 236,  55, 237, 238,\n",
       "         36, 239,  57, 240,   3, 241, 242, 243, 244, 240, 245,  55, 118,\n",
       "         16, 246, 228,  59, 247, 214, 248, 249,  76, 250,  69,  70,  16,\n",
       "        205,  78,  19,  70,  18,  70,  16,  66,  78,   3,  16, 251, 252,\n",
       "         15, 174, 187,  78, 253, 182,   3, 254, 255, 256, 257, 258, 256,\n",
       "         91,  84, 259,  31,  61, 260,  78,  34, 261,  31], dtype=int32),\n",
       " 188,\n",
       " 1)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_ds_v0[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding LSTMs "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input dim is the dimension of the embedding for each word (2 in the example)\n",
    "# Output dim is the dimension of the hidden layer (4 in this example)\n",
    "# batch_first – If True, then the input and output tensors are provided as (batch, seq, feature). \n",
    "lstm = nn.LSTM(2, 4, batch_first=True)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.0611,  0.9548],\n",
       "         [-0.3836,  1.2201],\n",
       "         [-1.4101, -1.9007],\n",
       "         [ 1.9819, -1.6735],\n",
       "         [-1.1701,  1.1169]]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inputs = [torch.randn(1, 2) for _ in range(5)] # make a sequence of length 5\n",
    "inputs = torch.cat(inputs).view(1, len(inputs), -1)\n",
    "inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 5, 2])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# RNNs with batch_first=True assume this input shape\n",
    "# input shape should be bash_size x seq_len x embedding dimension\n",
    "inputs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "out, (hidden, cell) = lstm(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 4])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.0772,  0.1723,  0.0313, -0.1882],\n",
       "         [-0.1118,  0.2238, -0.0149, -0.1851],\n",
       "         [ 0.0273, -0.1614, -0.3221,  0.0415],\n",
       "         [ 0.2766, -0.5186, -0.0278,  0.1104],\n",
       "         [ 0.0777,  0.0040, -0.1215,  0.0532]]], grad_fn=<TransposeBackward0>)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(out.shape)\n",
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.0777,  0.0040, -0.1215,  0.0532]]], grad_fn=<StackBackward>)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LSTM V0 model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTMV0Model(torch.nn.Module) :\n",
    "    def __init__(self, vocab_size, embedding_dim, hidden_dim) :\n",
    "        super(LSTMV0Model,self).__init__()\n",
    "        self.embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)\n",
    "        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)\n",
    "        self.linear = nn.Linear(hidden_dim, 1)\n",
    "        self.dropout = nn.Dropout(0.5)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.embeddings(x)\n",
    "        x = self.dropout(x)\n",
    "        out_pack, (ht, ct) = self.lstm(x)\n",
    "        return self.linear(ht[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epocs_v0(model, epochs=10, lr=0.001):\n",
    "    parameters = filter(lambda p: p.requires_grad, model.parameters())\n",
    "    optimizer = torch.optim.Adam(parameters, lr=lr)\n",
    "    for i in range(epochs):\n",
    "        model.train()\n",
    "        sum_loss = 0.0\n",
    "        total = 0\n",
    "        for x, s, y in train_dl:\n",
    "            # s is not used in this model\n",
    "            x = x.long().cuda()\n",
    "            y = y.float().cuda()\n",
    "            y_pred = model(x)\n",
    "            optimizer.zero_grad()\n",
    "            loss = F.binary_cross_entropy_with_logits(y_pred, y.unsqueeze(1))\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            sum_loss += loss.item()*y.shape[0]\n",
    "            total += y.shape[0]\n",
    "        val_loss, val_acc = val_metrics_v0(model, val_dl)\n",
    "        if i % 5 == 1:\n",
    "            print(\"train loss %.3f val loss %.3f and val accuracy %.3f\" % (sum_loss/total, val_loss, val_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def val_metrics_v0(model, valid_dl):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    sum_loss = 0.0\n",
    "    for x, s, y in valid_dl:\n",
    "        # s is not used here\n",
    "        x = x.long().cuda()\n",
    "        y = y.float().unsqueeze(1).cuda()\n",
    "        y_hat = model(x)\n",
    "        loss = F.binary_cross_entropy_with_logits(y_hat, y)\n",
    "        y_pred = y_hat > 0\n",
    "        correct += (y_pred.float() == y).float().sum()\n",
    "        total += y.shape[0]\n",
    "        sum_loss += loss.item()*y.shape[0]\n",
    "    return sum_loss/total, correct/total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 5000\n",
    "train_dl = DataLoader(train_ds_v0, batch_size=batch_size, shuffle=True)\n",
    "val_dl = DataLoader(valid_ds_v0, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "33916\n"
     ]
    }
   ],
   "source": [
    "vocab_size = len(words)\n",
    "print(vocab_size)\n",
    "model_v0 = LSTMV0Model(vocab_size, 50, 50).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train loss 0.683 val loss 0.675 and val accuracy 0.570\n",
      "train loss 0.594 val loss 0.863 and val accuracy 0.584\n",
      "train loss 0.485 val loss 0.856 and val accuracy 0.664\n",
      "train loss 0.359 val loss 1.108 and val accuracy 0.631\n",
      "train loss 0.270 val loss 1.236 and val accuracy 0.649\n",
      "train loss 0.236 val loss 0.763 and val accuracy 0.761\n"
     ]
    }
   ],
   "source": [
    "train_epocs_v0(model_v0, epochs=30, lr=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train loss 0.201 val loss 0.840 and val accuracy 0.766\n",
      "train loss 0.183 val loss 0.878 and val accuracy 0.757\n",
      "train loss 0.179 val loss 0.924 and val accuracy 0.753\n",
      "train loss 0.171 val loss 0.992 and val accuracy 0.740\n",
      "train loss 0.163 val loss 0.965 and val accuracy 0.753\n",
      "train loss 0.161 val loss 1.058 and val accuracy 0.744\n"
     ]
    }
   ],
   "source": [
    "train_epocs_v0(model_v0, epochs=30, lr=0.001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LSTM model with variable length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset with padding at the end\n",
    "train_ds = ImdbDataset(PATH)\n",
    "valid_ds = ImdbDataset(PATH, \"test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Debugging our model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 7\n",
    "train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "x,s,y = next(iter(train_dl)) # here s is the length of the sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([7, 400]), torch.Size([7]))"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape, s.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([162, 342, 168,  59, 166, 400, 181])"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 1, 1, 1, 1, 0, 0])"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sort by length so we can use pack_padded_sequence\n",
    "s, index = s.sort(0, descending=True)\n",
    "x = x[index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([400, 342, 181, 168, 166, 162,  59])"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([5, 1, 6, 2, 4, 0, 3])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 1, 0, 1, 1, 0, 1])"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y[index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_size = len(words)\n",
    "embedding_dim = 10\n",
    "embed = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([7, 400, 10])"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = embed(x.long())\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_dim = 9\n",
    "lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RNN will not perform calculation on pad elements if pack_padded_sequence is used\n",
    "x_pack = pack_padded_sequence(x, list(s), batch_first=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_pack, (ht, ct) = lstm(x_pack)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 7, 9])"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## final hidden layer\n",
    "ht.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([7, 9])"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ht[-1].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0057],\n",
       "        [ 0.0407],\n",
       "        [ 0.1412],\n",
       "        [ 0.1036],\n",
       "        [ 0.0486],\n",
       "        [ 0.0872],\n",
       "        [ 0.1097]], grad_fn=<AddmmBackward>)"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "linear = nn.Linear(hidden_dim, 1)\n",
    "y_hat = linear(ht[-1])\n",
    "y_hat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([7, 1])"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index.unsqueeze(1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "# takes back to the original ordering\n",
    "h = torch.zeros_like(y_hat).scatter_(0, index.unsqueeze(1), y_hat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0872],\n",
       "        [ 0.0407],\n",
       "        [ 0.1036],\n",
       "        [ 0.1097],\n",
       "        [ 0.0486],\n",
       "        [-0.0057],\n",
       "        [ 0.1412]], grad_fn=<ScatterBackward0>)"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0057],\n",
       "        [ 0.0407],\n",
       "        [ 0.1412],\n",
       "        [ 0.1036],\n",
       "        [ 0.0486],\n",
       "        [ 0.0872],\n",
       "        [ 0.1097]], grad_fn=<AddmmBackward>)"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_hat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([5, 1, 6, 2, 4, 0, 3])"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model v1\n",
    "Running this model on the GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTMModel(torch.nn.Module) :\n",
    "    def __init__(self, vocab_size, embedding_dim, hidden_dim) :\n",
    "        super(LSTMModel,self).__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.dropout = nn.Dropout(0.5)\n",
    "        self.embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)\n",
    "        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)\n",
    "        self.linear = nn.Linear(hidden_dim, 1)\n",
    "        \n",
    "    def forward(self, x, s):\n",
    "        # sorting\n",
    "        s, sort_index = torch.sort(s, 0,descending=True)\n",
    "        s = s.numpy().tolist()\n",
    "        x = x[sort_index]\n",
    "        x = self.embeddings(x)\n",
    "        x = self.dropout(x)\n",
    "        x_pack = pack_padded_sequence(x, s, batch_first=True)\n",
    "        out_pack, (ht, ct) = self.lstm(x_pack)\n",
    "        out = self.linear(ht[-1])\n",
    "        return torch.zeros_like(out).scatter_(0, sort_index.unsqueeze(1).cuda(), out) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epocs(model, epochs=10, lr=0.001):\n",
    "    parameters = filter(lambda p: p.requires_grad, model.parameters())\n",
    "    optimizer = torch.optim.Adam(parameters, lr=lr)\n",
    "    for i in range(epochs):\n",
    "        model.train()\n",
    "        sum_loss = 0.0\n",
    "        total = 0\n",
    "        for x, s, y in train_dl:\n",
    "            x = x.long().cuda()\n",
    "            y = y.float().cuda()\n",
    "            y_pred = model(x, s)\n",
    "            optimizer.zero_grad()\n",
    "            loss = F.binary_cross_entropy_with_logits(y_pred, y.unsqueeze(1))\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            sum_loss += loss.item()*y.shape[0]\n",
    "            total += y.shape[0]\n",
    "        val_loss, val_acc = val_metrics(model, valid_dl)\n",
    "        if i % 5 == 1:\n",
    "            print(\"train loss %.3f val loss %.3f and val accuracy %.3f\" % (sum_loss/total, val_loss, val_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "def val_metrics(model, valid_dl):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    sum_loss = 0.0\n",
    "    for x, s, y in valid_dl:\n",
    "        x = x.long().cuda()\n",
    "        y = y.float().unsqueeze(1).cuda()\n",
    "        y_hat = model(x, s)\n",
    "        loss = F.binary_cross_entropy_with_logits(y_hat, y)\n",
    "        y_pred = y_hat > 0\n",
    "        correct += (y_pred.float() == y).float().sum()\n",
    "        total += y.shape[0]\n",
    "        sum_loss += loss.item()*y.shape[0]\n",
    "    return sum_loss/total, correct/total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 2000\n",
    "train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)\n",
    "valid_dl = DataLoader(valid_ds, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "33916\n"
     ]
    }
   ],
   "source": [
    "vocab_size = len(words)\n",
    "print(vocab_size)\n",
    "model = LSTMModel(vocab_size, 50, 50).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train loss 0.667 val loss 0.623 and val accuracy 0.661\n",
      "train loss 0.588 val loss 0.682 and val accuracy 0.632\n",
      "train loss 0.328 val loss 0.381 and val accuracy 0.840\n",
      "train loss 0.216 val loss 0.384 and val accuracy 0.854\n",
      "train loss 0.144 val loss 0.373 and val accuracy 0.869\n",
      "train loss 0.102 val loss 0.432 and val accuracy 0.869\n"
     ]
    }
   ],
   "source": [
    "train_epocs(model, epochs=30, lr=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train loss 0.080 val loss 0.456 and val accuracy 0.868\n",
      "train loss 0.073 val loss 0.486 and val accuracy 0.867\n",
      "train loss 0.070 val loss 0.478 and val accuracy 0.868\n",
      "train loss 0.065 val loss 0.485 and val accuracy 0.868\n",
      "train loss 0.061 val loss 0.504 and val accuracy 0.869\n",
      "train loss 0.060 val loss 0.507 and val accuracy 0.869\n"
     ]
    }
   ],
   "source": [
    "train_epocs(model, epochs=30, lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train loss 0.058 val loss 0.520 and val accuracy 0.867\n",
      "train loss 0.057 val loss 0.528 and val accuracy 0.867\n",
      "train loss 0.052 val loss 0.536 and val accuracy 0.868\n",
      "train loss 0.050 val loss 0.538 and val accuracy 0.866\n",
      "train loss 0.049 val loss 0.549 and val accuracy 0.867\n",
      "train loss 0.048 val loss 0.557 and val accuracy 0.866\n"
     ]
    }
   ],
   "source": [
    "train_epocs(model, epochs=30, lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_model(m, p): torch.save(m.state_dict(), p)\n",
    "    \n",
    "def load_model(m, p): m.load_state_dict(torch.load(p))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir $PATH/\"models\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = PATH/\"models/model-86.pth\"\n",
    "save_model(model, p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.5580890059471131, tensor(0.8661, device='cuda:0'))"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_metrics(model, valid_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_model(model, p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GRU model with dropout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GRUModel(torch.nn.Module) :\n",
    "    def __init__(self, vocab_size, embedding_dim, hidden_dim):\n",
    "        super(GRUModel, self).__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)\n",
    "        self.dropout = nn.Dropout(0.5)\n",
    "        self.gru = nn.GRU(embedding_dim, hidden_dim, batch_first=True)\n",
    "        self.linear = nn.Linear(hidden_dim, 1)\n",
    "        \n",
    "    def forward(self, x, s):\n",
    "        s, sort_index = torch.sort(s, 0,descending=True)\n",
    "        s = s.numpy().tolist()\n",
    "        x = x[sort_index]\n",
    "        x = self.embeddings(x)\n",
    "        x = self.dropout(x)\n",
    "        x_pack = pack_padded_sequence(x, list(s), batch_first=True)\n",
    "        out_pack, ht= self.gru(x_pack)\n",
    "        out = self.linear(ht[-1])\n",
    "        return torch.zeros_like(out).scatter_(0, sort_index.unsqueeze(1).cuda(), out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "33916\n"
     ]
    }
   ],
   "source": [
    "vocab_size = len(words)\n",
    "print(vocab_size)\n",
    "model2 = GRUModel(vocab_size, 50, 50).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train loss 0.671 val loss 0.742 and val accuracy 0.601\n",
      "train loss 0.447 val loss 0.440 and val accuracy 0.800\n",
      "train loss 0.216 val loss 0.400 and val accuracy 0.865\n",
      "train loss 0.129 val loss 0.460 and val accuracy 0.869\n",
      "train loss 0.086 val loss 0.524 and val accuracy 0.872\n",
      "train loss 0.052 val loss 0.638 and val accuracy 0.865\n"
     ]
    }
   ],
   "source": [
    "train_epocs(model2, epochs=30, lr=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = PATH/\"models/model-gru-87.pth\"\n",
    "save_model(model2, p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bidirectional and multiple layers GRUs / LSTMs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 7\n",
    "train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "x,s,y = next(iter(train_dl)) # here s is the length of the sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_size = len(words)\n",
    "embedding_dim = 10\n",
    "hidden_dim = 9\n",
    "embed = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)\n",
    "lstm1 = nn.LSTM(embedding_dim, hidden_dim, batch_first=True, bidirectional=True)\n",
    "lstm2 = nn.LSTM(embedding_dim, hidden_dim, num_layers=2, batch_first=True, dropout=0.1, bidirectional=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "s, index = s.sort(0, descending=True)\n",
    "x = x[index]\n",
    "x = embed(x.long())\n",
    "x_pack = pack_padded_sequence(x, list(s), batch_first=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [],
   "source": [
    "lstm_out, (ht, ct) = lstm1(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 7, 9])"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ht.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([7, 9])"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ht[-2,:,:].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "lstm_out, (ht2, ct2) = lstm2(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 7, 9])"
      ]
     },
     "execution_count": 91,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ht2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([7, 9]), torch.Size([7, 9]))"
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ht2[-2,:,:].shape, ht2[-1,:,:].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([7, 18])"
      ]
     },
     "execution_count": 93,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#concat the final forward (ht[-2,:,:]) and backward (ht[-1,:,:]) hidden layers      \n",
    "h = torch.cat((ht2[-2,:,:], ht2[-1,:,:]), dim = 1)\n",
    "h.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 311,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is slow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTMBiModel(torch.nn.Module) :\n",
    "    def __init__(self, vocab_size, embedding_dim, hidden_dim) :\n",
    "        super(LSTMBiModel,self).__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)\n",
    "        self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layers=2, batch_first=True,\n",
    "                            dropout=0.3, bidirectional=True)\n",
    "        self.linear = nn.Linear(2*hidden_dim, 1)\n",
    "        \n",
    "    def forward(self, x, s):\n",
    "        s, sort_index = torch.sort(s, 0,descending=True)\n",
    "        s = s.numpy().tolist()\n",
    "        x = x[sort_index]\n",
    "        x = self.embeddings(x)\n",
    "        x_pack = pack_padded_sequence(x, s, batch_first=True)\n",
    "        out_pack, (ht, ct) = self.lstm(x_pack)\n",
    "        h = torch.cat((ht[-2,:,:], ht[-1,:,:]), dim = 1)\n",
    "        h = self.linear(h)\n",
    "        return torch.zeros_like(h).scatter_(0, sort_index.unsqueeze(1).cuda(), h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_size = len(words)\n",
    "model3 = LSTMBiModel(vocab_size, 50, 50).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train loss 0.413 val loss 0.404 and val accuracy 0.821\n",
      "train loss 0.342 val loss 0.448 and val accuracy 0.808\n",
      "train loss 0.344 val loss 0.439 and val accuracy 0.809\n"
     ]
    }
   ],
   "source": [
    "train_epocs(model3, epochs=15, lr=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bi GRUS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GRUBiModel(torch.nn.Module) :\n",
    "    def __init__(self, vocab_size, embedding_dim, hidden_dim) :\n",
    "        super(GRUBiModel,self).__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)\n",
    "        self.gru = nn.GRU(embedding_dim, hidden_dim, num_layers=2, batch_first=True,\n",
    "                            dropout=0.3, bidirectional=True)\n",
    "        self.linear = nn.Linear(2*hidden_dim, 1)\n",
    "        \n",
    "    def forward(self, x, s):\n",
    "        x = self.embeddings(x)\n",
    "        x_pack = pack_padded_sequence(x, s, batch_first=True)\n",
    "        out_pack, ht = self.gru(x_pack)\n",
    "        h = torch.cat((ht[-2,:,:], ht[-1,:,:]), dim = 1)\n",
    "        return self.linear(h)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise:\n",
    "Start with pre-trained embeddings."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model in this notebook is adapted from this [pytorch tutorial](https://pytorch.org/tutorials/beginner/nlp/sequence_models_tutorial.html). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": "block",
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}