06_language_model_rnn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neural Language Model - RNN (Recurrent Neural Network)\n",
    "\n",
    "RNN-based architectures (such as, LSTM, GRU, etc) is widely used in today's NLP.\n",
    "\n",
    "Recall that language model created in [previous example](./05_language_model_basic.ipynb) won't care the long past context.<br>\n",
    "For example, when the following sentence is given, \n",
    "\n",
    "\"In the United States, the president has now been\"\n",
    "\n",
    "it won't care the context \"In the United States\" when it refers only the last 5 words.<br>\n",
    "There might then be inconsistency in the sentence between former part and latter part.\n",
    "\n",
    "Let me assume another sentence \"it's vulgar and mean, but I liked it.\".<br>\n",
    "This sentence includes some negative phrases (\"vulgar\", \"mean\"), but the overall sentence has positive sentiment. This example shows that it's needed for precise predictions to understand not only individual phrases, but also the context in which they occur.\n",
    "\n",
    "In recurrent architecture, past context (called states) is inherited to the next prediction by the state memory $ s $ (which is trained by input and previous state), and this connection continues in the chain as follows. (See the following diagram.)<br>\n",
    "In this network, the next state $s_{i+1}$ is predicted by input $x_i$ and previous state $s_i$ in the network $R$ (which is called a recurrent unit) and this will be connected from beginning to the end of sequence. The output $y$ in each recurrent unit is generated by the state $s$ and the function $f(\\cdot)$. The output $y$ is then used for prediction in each unit.\n",
    "\n",
    "> Note : In simple RNN and GRU, $ f(\\cdot) $ in the following diagram is identity function.\n",
    "\n",
    "Recurrent Neural Network (RNN) will then be able to represent arbitrary size of sequence.\n",
    "\n",
    "![recurrent architecture](images/rnn_architecture.png)\n",
    "\n",
    "There are a lot of variants (including today's state-of-the-art model) in recurrent architecture.\n",
    "\n",
    "In **bidirectional RNN (BiRNN)**, the states in both directions (forward states and backward states) are maintained and trained as follows.\n",
    "\n",
    "![bidirectional rnn](images/bidirectional_rnn.png)\n",
    "\n",
    "Imagine that you predict the word [jumped] in the sentence, \"the brown fox [xxxxx] over the dog\". In this example, the latter context (\"over the dog\") is also important in the prediction.<br>\n",
    "The bidirectional RNN (BiRNN) is very effective, also in tagging tasks.\n",
    "\n",
    "In **deep RNN** (see below), the output is more deeply learned by multi-layered architecture. (See the following picture.)\n",
    "\n",
    "![deep rnn](images/deep_rnn.png)\n",
    "\n",
    "One of successful architecture in RNN is recurrent **gated architecture**.<br>\n",
    "With simple RNN, it will suffer from vanishing gradient problems, with which a lot of layers will rapidly lead the gradients of loss to zeros. (It will then eventually become hard to train the long past context in sequence.)<br>\n",
    "Briefly saying, gated architecture will avoid this problem by using gate vector $ g $ and new memory $ s^{\\prime} $ as follows :\n",
    "\n",
    "$ s^{\\prime} = g \\cdot x + (1 - g) \\cdot s $\n",
    "\n",
    "where $ \\cdot $ is inner product operation and $1$ is vector $(1,1,\\ldots,1)$.\n",
    "\n",
    "This computation will read the entries of input $ x $ which correspond to 1 values in $ g $, and read the entries of state $ s $ which correspond to 0 values in $ g $.<br>\n",
    "$ g $ is then also controlled and trained by input and previous memory state.\n",
    "\n",
    "**LSTM (Long Short Term Memory)** and **GRU (Gated Recurrent Unit)** are widely used gated architectures in language tasks.<br>\n",
    "I'll show you GRU in the following diagram.\n",
    "\n",
    "![gru architecture](images/gru_gate.png)\n",
    "\n",
    "$$ R : r_i = \\sigma(W_{rx} x_i + W_{rs} s_{i-1}) $$\n",
    "$$ Z : z_i = \\sigma(W_{zx} x_i + W_{zs} s_{i-1}) $$\n",
    "\n",
    "$$ \\tilde{S} : \\tilde{s}_i = tanh(W_{sx} x_i + W_{ss} (r_i \\cdot s_{i-1})) $$\n",
    "$$ S : s_i = (1 - z_i) \\cdot s_{i-1} + z_i \\cdot \\tilde{s}_i $$\n",
    "\n",
    "where $ \\sigma(\\cdot) $ is sigmoid activation and $ tanh(\\cdot) $ is tanh activation. (See [here](https://tsmatz.wordpress.com/2017/08/30/regression-in-machine-learning-math-for-beginners/) for sigmoid and tanh operation.)\n",
    "\n",
    "> Note : The bias term is often included, such as $ Z : z_i = \\sigma(W_{zx} x_i + W_{zs} s_{i-1} + b_z) $.\n",
    "\n",
    "In GRU architecture, the new state candidate $ \\tilde{s}_i $ is computed by using the controlled parameter $ r_i $. (And $ r_i $ is also trained by inputs.)<br>\n",
    "The updated final state $ s_i $ is then determined based on the weight between previous state $ s_{i-1} $ and state candidate $ \\tilde{s}_i $, by using controlled parameter $ z_i $. (And $ z_i $ is also trained by inputs.)\n",
    "\n",
    "In this example, we will train 2 language models with simple RNN (Simple Recurrent Neural Network) and GRU (Gated Recurrent Unit) architecture in word's prediction task.\n",
    "\n",
    "*back to [index](https://github.com/tsmatz/nlp-tutorials/)*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install required packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch pandas numpy nltk"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, I have used short description text in news papers dataset, since it's formal-styled concise sentence (not including slangs and it's today's modern English).<br>\n",
    "Before starting, please download [News_Category_Dataset_v2.json](https://www.kaggle.com/datasets/rmisra/news-category-dataset/versions/2) (collected by HuffPost) in Kaggle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0         She left her husband. He killed their children...\n",
       "1                                  Of course it has a song.\n",
       "2         The actor and his longtime girlfriend Anna Ebe...\n",
       "3         The actor gives Dems an ass-kicking for not fi...\n",
       "4         The \"Dietland\" actress said using the bags is ...\n",
       "                                ...                        \n",
       "200848    Verizon Wireless and AT&T are already promotin...\n",
       "200849    Afterward, Azarenka, more effusive with the pr...\n",
       "200850    Leading up to Super Bowl XLVI, the most talked...\n",
       "200851    CORRECTION: An earlier version of this story i...\n",
       "200852    The five-time all-star center tore into his te...\n",
       "Name: short_description, Length: 200853, dtype: object"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.read_json(\"News_Category_Dataset_v2.json\",lines=True)\n",
    "train_data = df[\"short_description\"]\n",
    "train_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get the better performance (accuracy), we standarize the input text as follows.\n",
    "- Make all words to lowercase in order to reduce words\n",
    "- Make \"-\" (hyphen) to space\n",
    "- Remove all punctuation except \" ' \" (e.g, don't, isn't) and \"&\" (e.g, AT&T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0         she left her husband he killed their children ...\n",
       "1                                   of course it has a song\n",
       "2         the actor and his longtime girlfriend anna ebe...\n",
       "3         the actor gives dems an ass kicking for not fi...\n",
       "4         the dietland actress said using the bags is a ...\n",
       "                                ...                        \n",
       "200848    verizon wireless and at&t are already promotin...\n",
       "200849    afterward azarenka more effusive with the pres...\n",
       "200850    leading up to super bowl xlvi the most talked ...\n",
       "200851    correction an earlier version of this story in...\n",
       "200852    the five time all star center tore into his te...\n",
       "Name: short_description, Length: 200853, dtype: object"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data = train_data.str.lower()\n",
    "train_data = train_data.str.replace(\"-\", \" \", regex=True)\n",
    "train_data = train_data.str.replace(r\"[^'\\&\\w\\s]\", \"\", regex=True)\n",
    "train_data = train_data.str.strip()\n",
    "train_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we add ```<start>``` and ```<end>``` tokens in each sequence as follows, because these are important information for learning the ordered sequence.\n",
    "\n",
    "```this is a pen``` --> ```<start> this is a pen <end>```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<start> she left her husband he killed their children just another day in america <end>'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data = [\" \".join([\"<start>\", x, \"<end>\"]) for x in train_data]\n",
    "# print first row\n",
    "train_data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate sequence inputs\n",
    "\n",
    "We will generate the sequence of word's indices (i.e, tokenize) from text.\n",
    "\n",
    "![Index vectorize](images/index_vectorize2.png)\n",
    "\n",
    "First we create a list of vocabulary (```vocab```)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nltk.tokenize import SpaceTokenizer\n",
    "\n",
    "###\n",
    "# define Vocab\n",
    "###\n",
    "class Vocab:\n",
    "    def __init__(self, list_of_sentence, tokenization, special_token, max_tokens=None):\n",
    "        # count vocab frequency\n",
    "        vocab_freq = {}\n",
    "        tokens = tokenization(list_of_sentence)\n",
    "        for t in tokens:\n",
    "            for vocab in t:\n",
    "                if vocab not in vocab_freq:\n",
    "                    vocab_freq[vocab] = 0 \n",
    "                vocab_freq[vocab] += 1\n",
    "        # sort by frequency\n",
    "        vocab_freq = {k: v for k, v in sorted(vocab_freq.items(), key=lambda i: i[1], reverse=True)}\n",
    "        # create vocab list\n",
    "        self.vocabs = [special_token] + list(vocab_freq.keys())\n",
    "        if max_tokens:\n",
    "            self.vocabs = self.vocabs[:max_tokens]\n",
    "        self.stoi = {v: i for i, v in enumerate(self.vocabs)}\n",
    "\n",
    "    def _get_tokens(self, list_of_sentence):\n",
    "        for sentence in list_of_sentence:\n",
    "            tokens = tokenizer.tokenize(sentence)\n",
    "            yield tokens\n",
    "\n",
    "    def get_itos(self):\n",
    "        return self.vocabs\n",
    "\n",
    "    def get_stoi(self):\n",
    "        return self.stoi\n",
    "\n",
    "    def append_token(self, token):\n",
    "        self.vocabs.append(token)\n",
    "        self.stoi = {v: i for i, v in enumerate(self.vocabs)}\n",
    "\n",
    "    def __call__(self, list_of_tokens):\n",
    "        def get_token_index(token):\n",
    "            if token in self.stoi:\n",
    "                return self.stoi[token]\n",
    "            else:\n",
    "                return 0\n",
    "        return [get_token_index(t) for t in list_of_tokens]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.vocabs)\n",
    "\n",
    "###\n",
    "# generate Vocab\n",
    "###\n",
    "max_word = 50000\n",
    "\n",
    "# create tokenizer\n",
    "tokenizer = SpaceTokenizer()\n",
    "\n",
    "# define tokenization function\n",
    "def yield_tokens(data):\n",
    "    for text in data:\n",
    "        tokens = tokenizer.tokenize(text)\n",
    "        yield tokens\n",
    "\n",
    "# build vocabulary list\n",
    "vocab = Vocab(\n",
    "    train_data,\n",
    "    tokenization=yield_tokens,\n",
    "    special_token=\"<unk>\",\n",
    "    max_tokens=max_word,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The generated token index is ```0, 1, ... , vocab_size - 1```.<br>\n",
    "Now I set ```vocab_size``` (here 50000) as a token id in padded positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "pad_index = vocab.__len__()\n",
    "vocab.append_token(\"<pad>\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get list for both index-to-word and word-to-index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "itos = vocab.get_itos()\n",
    "stoi = vocab.get_stoi()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The number of token index is 50001.\n",
      "The padded index is 50000.\n"
     ]
    }
   ],
   "source": [
    "# test\n",
    "print(\"The number of token index is {}.\".format(vocab.__len__()))\n",
    "print(\"The padded index is {}.\".format(stoi[\"<pad>\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we build a collator function, which is used for pre-processing in data loader.\n",
    "\n",
    "In this collator, first we create a list of word's indices as follows.\n",
    "\n",
    "```<start> this is pen <end>``` --> ```[2, 7, 5, 14, 1]```\n",
    "\n",
    "Next we separate into features (x) and labels (y).<br>\n",
    "In this task, we predict the next word in the sequence, and we then create the following features (x) and labels (y) in each row.\n",
    "\n",
    "<u>before</u> :\n",
    "\n",
    "```[2, 7, 5, 14, 1]```\n",
    "\n",
    "<u>after</u> :\n",
    "\n",
    "```x : [2, 7, 5, 14, 1]```\n",
    "\n",
    "```y : [7, 5, 14, 1, -100]```\n",
    "\n",
    "> Note : Here I set -100 as an unknown label id, because PyTorch cross-entropy function (```torch.nn.functional.cross_entropy()```) has a property ```ignore_index``` which default value is -100.\n",
    "\n",
    "Finally we pad the inputs as follows.<br>\n",
    "The padded index in features is ```pad_index``` and the padded index in label is -100. (See above note.)\n",
    "\n",
    "```x : [2, 7, 5, 14, 1, N, ... , N]```\n",
    "\n",
    "```y : [7, 5, 14, 1, -100, -100, ... , -100]```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "max_seq_len = 256\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "def collate_batch(batch):\n",
    "    label_list, feature_list = [], []\n",
    "    for text in batch:\n",
    "        # tokenize to a list of word's indices\n",
    "        tokens = vocab(tokenizer.tokenize(text))\n",
    "        # separate into features and labels\n",
    "        y = tokens[1:]\n",
    "        y.append(-100)\n",
    "        x = tokens\n",
    "        # limit length to max_seq_len\n",
    "        y = y[:max_seq_len]\n",
    "        x = x[:max_seq_len]\n",
    "        # pad features and labels\n",
    "        y += [-100] * (max_seq_len - len(y))\n",
    "        x += [pad_index] * (max_seq_len - len(x))\n",
    "        # add to list\n",
    "        label_list.append(y)\n",
    "        feature_list.append(x)\n",
    "    # convert to tensor\n",
    "    label_list = torch.tensor(label_list, dtype=torch.int64).to(device)\n",
    "    feature_list = torch.tensor(feature_list, dtype=torch.int64).to(device)\n",
    "    return label_list, feature_list\n",
    "\n",
    "dataloader = DataLoader(\n",
    "    train_data,\n",
    "    batch_size=16,\n",
    "    shuffle=True,\n",
    "    collate_fn=collate_batch\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "label shape in batch : torch.Size([16, 256])\n",
      "feature shape in batch : torch.Size([16, 256])\n",
      "***** label sample *****\n",
      "tensor([   5,  277, 5532,  629,    5, 5773,   14,   44,    3,  132,  767,    1,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100], device='cuda:0')\n",
      "***** features sample *****\n",
      "tensor([    2,     5,   277,  5532,   629,     5,  5773,    14,    44,     3,\n",
      "          132,   767,     1, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000, 50000,\n",
      "        50000, 50000, 50000, 50000, 50000, 50000], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "# test\n",
    "for labels, features in dataloader:\n",
    "    break\n",
    "\n",
    "print(\"label shape in batch : {}\".format(labels.size()))\n",
    "print(\"feature shape in batch : {}\".format(features.size()))\n",
    "print(\"***** label sample *****\")\n",
    "print(labels[0])\n",
    "print(\"***** features sample *****\")\n",
    "print(features[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build network\n",
    "\n",
    "Now we build a model for this next word's prediction using simple RNN architecture.\n",
    "\n",
    "![RNN network](images/rnn_network.png)\n",
    "\n",
    "In PyTorch, you can use ```torch.nn.RNN``` module for processing simple RNN, and we also use this built-in module in this example.\n",
    "\n",
    "In the following example, the shape of RNN input is expected to be ```(batch_size, sequence_length, input_dimension)```.<br>\n",
    "However, to tell which time steps in each sequence should be processed in RNN (i.e, for RNN masking), we wrap this tensor as a packed sequence with ```torch.nn.utils.rnn.pack_padded_sequence()``` before passing into RNN module.<br>\n",
    "For example, when batch size is 4 and we generate a packed sequence with ```lengths=[5, 3, 3, 2]``` in ```torch.nn.utils.rnn.pack_padded_sequence()```, the processed sequence# in each time-step will then be :\n",
    "\n",
    "```\n",
    "time-step 1 : {1, 2, 3, 4}\n",
    "time-step 2 : {1, 2, 3, 4}\n",
    "time-step 3 : {1, 2, 3}\n",
    "time-step 4 : {1}\n",
    "time-step 5 : {1}\n",
    "```\n",
    "\n",
    "As a result, it's processed with new batch size ```[4, 4, 3, 1, 1]```. (See below picture.)\n",
    "\n",
    "![packed sequence](images/rnn_packed_sequence.png)\n",
    "\n",
    "> Note : When the length is not sorted, first all sequences in batch are sorted by descending length of sequence, and planned to run batches to meet each time-steps. (When it's unpacked, the order is returned to the original position.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "embedding_dim = 64\n",
    "rnn_units = 512\n",
    "\n",
    "class SimpleRnnModel(nn.Module):\n",
    "    def __init__(self, vocab_size, seq_len, embedding_dim, rnn_units, padding_idx):\n",
    "        super().__init__()\n",
    "\n",
    "        self.seq_len = seq_len\n",
    "        self.padding_idx = padding_idx\n",
    "\n",
    "        self.embedding = nn.Embedding(\n",
    "            vocab_size,\n",
    "            embedding_dim,\n",
    "            padding_idx=padding_idx,\n",
    "        )\n",
    "        self.rnn = nn.RNN(\n",
    "            input_size=embedding_dim,\n",
    "            hidden_size=rnn_units,\n",
    "            num_layers=1,\n",
    "            batch_first=True,\n",
    "        )\n",
    "        self.classify = nn.Linear(rnn_units, vocab_size)\n",
    "\n",
    "    def forward(self, inputs, states=None, return_final_state=False):\n",
    "        # embedding\n",
    "        #   --> (batch_size, seq_len, embedding_dim)\n",
    "        outs = self.embedding(inputs)\n",
    "        # build \"lengths\" property to pack inputs (see above)\n",
    "        lengths = (inputs != self.padding_idx).int().sum(dim=1, keepdim=False)\n",
    "        # pack inputs for RNN\n",
    "        packed_inputs = torch.nn.utils.rnn.pack_padded_sequence(\n",
    "            outs,\n",
    "            lengths.cpu(),\n",
    "            batch_first=True,\n",
    "            enforce_sorted=False,\n",
    "        )\n",
    "        # apply RNN\n",
    "        if states is None:\n",
    "            packed_outs, final_state = self.rnn(packed_inputs)\n",
    "        else:\n",
    "            packed_outs, final_state = self.rnn(packed_inputs, states)\n",
    "        # unpack results\n",
    "        #   --> (batch_size, seq_len, rnn_units)\n",
    "        outs, _ = torch.nn.utils.rnn.pad_packed_sequence(\n",
    "            packed_outs,\n",
    "            batch_first=True,\n",
    "            padding_value=0.0,\n",
    "            total_length=self.seq_len,\n",
    "        )\n",
    "        # apply feed-forward to classify\n",
    "        #   --> (batch_size, seq_len, vocab_size)\n",
    "        logits = self.classify(outs)\n",
    "        # return results\n",
    "        if return_final_state:\n",
    "            return logits, final_state  # This is used in prediction\n",
    "        else:\n",
    "            return logits               # This is used in training\n",
    "\n",
    "model = SimpleRnnModel(\n",
    "    vocab_size=vocab.__len__(),\n",
    "    seq_len=max_seq_len,\n",
    "    embedding_dim=embedding_dim,\n",
    "    rnn_units=rnn_units,\n",
    "    padding_idx=pad_index).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now run training with above model.\n",
    "\n",
    "As I have mentioned above, the loss on label id=-100 is ignored in ```cross_entropy()``` function. The padded position and the end of sequence will then be ignored in optimization.\n",
    "\n",
    "> Note : Because the default value of  ```ignore_index``` property in ```cross_entropy()``` function is -100. (You can change this default value.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 - loss: 5.9431 - accuracy: 0.11630\n",
      "Epoch 2 - loss: 6.4704 - accuracy: 0.1789\n",
      "Epoch 3 - loss: 6.2977 - accuracy: 0.0833\n",
      "Epoch 4 - loss: 5.6396 - accuracy: 0.1762\n",
      "Epoch 5 - loss: 5.4232 - accuracy: 0.1679\n"
     ]
    }
   ],
   "source": [
    "from torch.nn import functional as F\n",
    "\n",
    "num_epochs = 5\n",
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)\n",
    "for epoch in range(num_epochs):\n",
    "    for labels, seqs in dataloader:\n",
    "        # optimize\n",
    "        optimizer.zero_grad()\n",
    "        logits = model(seqs)\n",
    "        loss = F.cross_entropy(logits.transpose(1,2), labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        # calculate accuracy\n",
    "        pred_labels = logits.argmax(dim=2)\n",
    "        num_correct = (pred_labels == labels).float().sum()\n",
    "        num_total = (labels != -100).float().sum()\n",
    "        accuracy = num_correct / num_total\n",
    "        print(\"Epoch {} - loss: {:2.4f} - accuracy: {:2.4f}\".format(epoch+1, loss.item(), accuracy), end=\"\\r\")\n",
    "    print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate Text (Simple RNN)\n",
    "\n",
    "Here I simply generate several text with trained model.\n",
    "\n",
    "The metrics to evaluate text generation task is not so easy. (Because simply checking an exact match to a reference text is not optimal.)<br>\n",
    "Use some common metrics available in these cases, such as, BLEU or ROUGE.\n",
    "\n",
    "> Note : Here I use greedy search and this will sometimes lead to wrong sequence. For drawbacks and solutions, see note in [this example](./05_language_model_basic.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<start> prime minister theresa <unk> said the <unk> was a hero in the world of the arctic monkeys <end>\n",
      "<start> chairman of the <unk> ' s widow ' s chief of staff reince priebus said the former chief of staff reince priebus said he ' s advocating to be a source of the past <end>\n",
      "<start> he was expected to be a politician <end>\n"
     ]
    }
   ],
   "source": [
    "end_index = stoi[\"<end>\"]\n",
    "max_output = 128\n",
    "\n",
    "def pred_output(text):\n",
    "    generated_text = \"<start> \" + text\n",
    "    _, inputs = collate_batch([generated_text])\n",
    "    mask = (inputs != pad_index).int()\n",
    "    last_idx = mask[0].sum() - 1\n",
    "    final_states = None\n",
    "    outputs, final_states = model(inputs, final_states, return_final_state=True)\n",
    "    pred_index = outputs[0][last_idx].argmax()\n",
    "    for loop in range(max_output):\n",
    "        generated_text += \" \"\n",
    "        next_word = itos[pred_index]\n",
    "        generated_text += next_word\n",
    "        if pred_index.item() == end_index:\n",
    "            break\n",
    "        _, inputs = collate_batch([next_word])\n",
    "        outputs, final_states = model(inputs, final_states, return_final_state=True)\n",
    "        pred_index = outputs[0][0].argmax()\n",
    "    return generated_text\n",
    "\n",
    "print(pred_output(\"prime\"))\n",
    "print(pred_output(\"chairman\"))\n",
    "print(pred_output(\"he was expected\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train with GRU\n",
    "\n",
    "Next we train the same task with gated architecture, GRU (gated recurrent unit).<br>\n",
    "As I have mentioned above, GRU layer has following architecture.\n",
    "\n",
    "![gru architecture](images/gru_gate.png)\n",
    "\n",
    "$$ R : r_i = \\sigma(W_{rx} x_i + W_{rs} s_{i-1}) $$\n",
    "$$ Z : z_i = \\sigma(W_{zx} x_i + W_{zs} s_{i-1}) $$\n",
    "\n",
    "$$ \\tilde{S} : \\tilde{s}_i = tanh(W_{sx} x_i + W_{ss} (r_i \\cdot s_{i-1})) $$\n",
    "$$ S : s_i = (1 - z_i) \\cdot s_{i-1} + z_i \\cdot \\tilde{s}_i $$\n",
    "\n",
    "In this example, we use built-in layer ```torch.nn.GRU``` in PyTorch.\n",
    "\n",
    "> Note : In the following example, we use bias term in GRU layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_dim = 64\n",
    "rnn_units = 512\n",
    "\n",
    "class GruModel(nn.Module):\n",
    "    def __init__(self, vocab_size, seq_len, embedding_dim, rnn_units, padding_idx):\n",
    "        super().__init__()\n",
    "\n",
    "        self.seq_len = seq_len\n",
    "        self.padding_idx = padding_idx\n",
    "\n",
    "        self.embedding = nn.Embedding(\n",
    "            vocab_size,\n",
    "            embedding_dim,\n",
    "            padding_idx=padding_idx,\n",
    "        )\n",
    "        self.rnn = nn.GRU(\n",
    "            input_size=embedding_dim,\n",
    "            hidden_size=rnn_units,\n",
    "            num_layers=1,\n",
    "            batch_first=True,\n",
    "        )\n",
    "        self.classify = nn.Linear(rnn_units, vocab_size)\n",
    "\n",
    "    def forward(self, inputs, states=None, return_final_state=False):\n",
    "        # embedding\n",
    "        #   --> (batch_size, seq_len, embedding_dim)\n",
    "        outs = self.embedding(inputs)\n",
    "        # build \"lengths\" property to pack inputs (see above)\n",
    "        lengths = (inputs != self.padding_idx).int().sum(dim=1, keepdim=False)\n",
    "        # pack inputs for RNN\n",
    "        packed_inputs = torch.nn.utils.rnn.pack_padded_sequence(\n",
    "            outs,\n",
    "            lengths.cpu(),\n",
    "            batch_first=True,\n",
    "            enforce_sorted=False,\n",
    "        )\n",
    "        # apply RNN\n",
    "        if states is None:\n",
    "            packed_outs, final_state = self.rnn(packed_inputs)\n",
    "        else:\n",
    "            packed_outs, final_state = self.rnn(packed_inputs, states)\n",
    "        # unpack results\n",
    "        #   --> (batch_size, seq_len, rnn_units)\n",
    "        outs, _ = torch.nn.utils.rnn.pad_packed_sequence(\n",
    "            packed_outs,\n",
    "            batch_first=True,\n",
    "            padding_value=0.0,\n",
    "            total_length=self.seq_len,\n",
    "        )\n",
    "        # apply feed-forward to classify\n",
    "        #   --> (batch_size, seq_len, vocab_size)\n",
    "        logits = self.classify(outs)\n",
    "        # return results\n",
    "        if return_final_state:\n",
    "            return logits, final_state  # This is used in prediction\n",
    "        else:\n",
    "            return logits               # This is used in training\n",
    "\n",
    "model = GruModel(\n",
    "    vocab_size=vocab.__len__(),\n",
    "    seq_len=max_seq_len,\n",
    "    embedding_dim=embedding_dim,\n",
    "    rnn_units=rnn_units,\n",
    "    padding_idx=pad_index).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 - loss: 5.7050 - accuracy: 0.16552\n",
      "Epoch 2 - loss: 5.4469 - accuracy: 0.1743\n",
      "Epoch 3 - loss: 3.1864 - accuracy: 0.4911\n",
      "Epoch 4 - loss: 5.4429 - accuracy: 0.1346\n",
      "Epoch 5 - loss: 5.3104 - accuracy: 0.2817\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 5\n",
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)\n",
    "for epoch in range(num_epochs):\n",
    "    for labels, seqs in dataloader:\n",
    "        # optimize\n",
    "        optimizer.zero_grad()\n",
    "        logits = model(seqs)\n",
    "        loss = F.cross_entropy(logits.transpose(1,2), labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        # calculate accuracy\n",
    "        pred_labels = logits.argmax(dim=2)\n",
    "        num_correct = (pred_labels == labels).float().sum()\n",
    "        num_total = (labels != -100).float().sum()\n",
    "        accuracy = num_correct / num_total\n",
    "        print(\"Epoch {} - loss: {:2.4f} - accuracy: {:2.4f}\".format(epoch+1, loss.item(), accuracy), end=\"\\r\")\n",
    "    print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate Text (GRU)\n",
    "\n",
    "Here I simply generate several text with trained model.\n",
    "\n",
    "The metrics to evaluate text generation task is not so easy. (Because simply checking an exact match to a reference text is not optimal.)<br>\n",
    "Use some common metrics available in these cases, such as, BLEU or ROUGE.\n",
    "\n",
    "> Note : Here I use greedy search and this will sometimes lead to wrong sequence. For drawbacks and solutions, see note in [this example](./05_language_model_basic.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<start> prime minister justin trudeau is a big part of the game of the republican party <end>\n",
      "<start> chairman of the house appropriations committee on the verge of the supreme court nominee <end>\n",
      "<start> he was expected to be a little girl <end>\n"
     ]
    }
   ],
   "source": [
    "end_index = stoi[\"<end>\"]\n",
    "max_output = 128\n",
    "\n",
    "def pred_output(text):\n",
    "    generated_text = \"<start> \" + text\n",
    "    _, inputs = collate_batch([generated_text])\n",
    "    mask = (inputs != pad_index).int()\n",
    "    last_idx = mask[0].sum() - 1\n",
    "    final_states = None\n",
    "    outputs, final_states = model(inputs, final_states, return_final_state=True)\n",
    "    pred_index = outputs[0][last_idx].argmax()\n",
    "    for loop in range(max_output):\n",
    "        generated_text += \" \"\n",
    "        next_word = itos[pred_index]\n",
    "        generated_text += next_word\n",
    "        if pred_index.item() == end_index:\n",
    "            break\n",
    "        _, inputs = collate_batch([next_word])\n",
    "        outputs, final_states = model(inputs, final_states, return_final_state=True)\n",
    "        pred_index = outputs[0][0].argmax()\n",
    "    return generated_text\n",
    "\n",
    "print(pred_output(\"prime\"))\n",
    "print(pred_output(\"chairman\"))\n",
    "print(pred_output(\"he was expected\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}