rnn_tv_script_generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TV Script Generation\n",
    "\n",
    "[Seinfeld](https://en.wikipedia.org/wiki/Seinfeld) TV scripts are generated utilizing RNNs.  Part of the 9 season [Seinfeld dataset](https://www.kaggle.com/thec03u5/seinfeld-chronicles#scripts.csv) scripts is utilized.  The Neural Network below generates a new ,\"fake\" TV script, based on patterns it recognizes in this training data.\n",
    "\n",
    "## Get the Data\n",
    "\n",
    "The data is located in `./data/Seinfeld_Scripts.txt`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load in data\n",
    "import helper\n",
    "data_dir = './data/Seinfeld_Scripts.txt'\n",
    "text = helper.load_data(data_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Explore the Data\n",
    "Play around with `view_line_range` to view different parts of the data. This will give you a sense of the data you'll be working with. You can see, for example, that it is all lowercase text, and each new line of dialogue is separated by a newline character `\\n`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset Stats\n",
      "Roughly the number of unique words: 46367\n",
      "Number of lines: 109233\n",
      "Average number of words in each line: 5.544240293684143\n",
      "\n",
      "The lines 0 to 10:\n",
      "jerry: do you know what this is all about? do you know, why were here? to be out, this is out...and out is one of the single most enjoyable experiences of life. people...did you ever hear people talking about we should go out? this is what theyre talking about...this whole thing, were all out now, no one is home. not one person here is home, were all out! there are people trying to find us, they dont know where we are. (on an imaginary phone) did you ring?, i cant find him. where did he go? he didnt tell me where he was going. he must have gone out. you wanna go out you get ready, you pick out the clothes, right? you take the shower, you get all ready, get the cash, get your friends, the car, the spot, the reservation...then youre standing around, what do you do? you go we gotta be getting back. once youre out, you wanna get back! you wanna go to sleep, you wanna get up, you wanna go out again tomorrow, right? where ever you are in life, its my feeling, youve gotta go. \n",
      "\n",
      "jerry: (pointing at georges shirt) see, to me, that button is in the worst possible spot. the second button literally makes or breaks the shirt, look at it. its too high! its in no-mans-land. you look like you live with your mother. \n",
      "\n",
      "george: are you through? \n",
      "\n",
      "jerry: you do of course try on, when you buy? \n",
      "\n",
      "george: yes, it was purple, i liked it, i dont actually recall considering the buttons. \n",
      "\n"
     ]
    }
   ],
   "source": [
    "view_line_range = (0, 10)\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "print('Dataset Stats')\n",
    "print('Roughly the number of unique words: {}'.format(len({word: None for word in text.split()})))\n",
    "\n",
    "lines = text.split('\\n')\n",
    "print('Number of lines: {}'.format(len(lines)))\n",
    "word_count_line = [len(line.split()) for line in lines]\n",
    "print('Average number of words in each line: {}'.format(np.average(word_count_line)))\n",
    "\n",
    "print()\n",
    "print('The lines {} to {}:'.format(*view_line_range))\n",
    "print('\\n'.join(text.split('\\n')[view_line_range[0]:view_line_range[1]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Pre-processing Functions\n",
    "The following pre-processing functions are implemented below:\n",
    "- Lookup Table\n",
    "- Tokenize Punctuation\n",
    "\n",
    "### Lookup Table\n",
    "A word embedding is created by first transforming the words into ids.  In this function, two dictionaries are created:\n",
    "- Dictionary to go from the words to an id, called `vocab_to_int`\n",
    "- Dictionary to go from the id to word, called `int_to_vocab`\n",
    "\n",
    "These dictionaries are returned in the following **tuple**: `(vocab_to_int, int_to_vocab)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tests Passed\n"
     ]
    }
   ],
   "source": [
    "import problem_unittests as tests\n",
    "from collections import Counter\n",
    "\n",
    "def create_lookup_tables(text):\n",
    "    \"\"\"\n",
    "    Create lookup tables for vocabulary\n",
    "    :param text: The text of tv scripts split into words\n",
    "    :return: A tuple of dicts (vocab_to_int, int_to_vocab)\n",
    "    \"\"\"\n",
    "    # Create Counter class for text\n",
    "    word_counts = Counter(text)\n",
    "    \n",
    "    # Sort counter descending (most used words appear first)\n",
    "    sorted_vocab = sorted(word_counts, key=word_counts.get, reverse=True)\n",
    "    \n",
    "    # Create dictionaries to return\n",
    "    int_to_vocab = {ii: word for ii, word in enumerate(sorted_vocab)}\n",
    "    vocab_to_int = {word: ii for ii, word in int_to_vocab.items()}\n",
    "            \n",
    "    # return tuple\n",
    "    return (vocab_to_int, int_to_vocab)\n",
    "\n",
    "\n",
    "tests.test_create_lookup_tables(create_lookup_tables)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tokenize Punctuation\n",
    "The script is split into a word array using spaces as delimiters.  However, punctuations like periods and exclamation marks can create multiple ids for the same word. For example, \"bye\" and \"bye!\" would generate two different word ids.\n",
    "\n",
    "The function `token_lookup` returns a dict that will be used to tokenize symbols like \"!\" into \"||Exclamation_Mark||\".  A dictionary is created for the following symbols where the symbol is the key and value is the token:\n",
    "- Period ( **.** )\n",
    "- Comma ( **,** )\n",
    "- Quotation Mark ( **\"** )\n",
    "- Semicolon ( **;** )\n",
    "- Exclamation mark ( **!** )\n",
    "- Question mark ( **?** )\n",
    "- Left Parentheses ( **(** )\n",
    "- Right Parentheses ( **)** )\n",
    "- Dash ( **-** )\n",
    "- Return ( **\\n** )\n",
    "\n",
    "This dictionary will be used to tokenize the symbols and add the delimiter (space) around it.  This separates each symbols as its own word, making it easier for the neural network to predict the next word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tests Passed\n"
     ]
    }
   ],
   "source": [
    "def token_lookup():\n",
    "    \"\"\"\n",
    "    Generate a dict to turn punctuation into a token.\n",
    "    :return: Tokenized dictionary where the key is the punctuation and the value is the token\n",
    "    \"\"\"\n",
    "    token_dictionary = {}\n",
    "    token_dictionary[\".\"] = \"||Period||\"\n",
    "    token_dictionary[\",\"] = \"||Comma||\"\n",
    "    token_dictionary[\"\\\"\"] = \"||Quotation_Mark||\"\n",
    "    token_dictionary[\";\"] = \"||Semicolon||\"\n",
    "    token_dictionary[\"!\"] = \"||Exclamation_Mark||\"\n",
    "    token_dictionary[\"?\"] = \"||Question_Mark||\"\n",
    "    token_dictionary[\"(\"] = \"||Left_Parentheses||\"\n",
    "    token_dictionary[\")\"] = \"||Right_Parentheses||\"\n",
    "    token_dictionary[\"-\"] = \"||Dash||\"\n",
    "    token_dictionary[\"\\n\"] = \"||Return||\"\n",
    "    \n",
    "    return token_dictionary\n",
    "\n",
    "tests.test_tokenize(token_lookup)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-process all the data and save it\n",
    "\n",
    "Running the code cell below will pre-process all the data and save it to file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pre-process training data\n",
    "helper.preprocess_and_save_data(data_dir, token_lookup, create_lookup_tables)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Checkpoint\n",
    "Checkpoint data in case you want to halt the Jupyter kernel and come back at a later time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import helper\n",
    "import problem_unittests as tests\n",
    "\n",
    "int_text, vocab_to_int, int_to_vocab, token_dict = helper.load_preprocess()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Neural Network\n",
    "This section contains the components necessary to build an RNN by implementing the RNN Module and forward and backpropagation functions.\n",
    "\n",
    "### Check Access to GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# Check for a GPU\n",
    "train_on_gpu = torch.cuda.is_available()\n",
    "if not train_on_gpu:\n",
    "    print('No GPU found. Please use a GPU to train your neural network.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Input\n",
    "The [TensorDataset](http://pytorch.org/docs/master/data.html#torch.utils.data.TensorDataset) is used to provide a known format to our dataset; in combination with [DataLoader](http://pytorch.org/docs/master/data.html#torch.utils.data.DataLoader), it will handle batching, shuffling, and other dataset iteration functions.\n",
    "\n",
    "The TensorDataset is created by passing in feature and target tensors. The DataLoader is then created as usual.\n",
    "```\n",
    "data = TensorDataset(feature_tensors, target_tensors)\n",
    "data_loader = torch.utils.data.DataLoader(data, \n",
    "                                          batch_size=batch_size)\n",
    "```\n",
    "\n",
    "### Batching\n",
    "The `batch_data` function below batches `words` data into chunks of size `batch_size` using the `TensorDataset` and `DataLoader` classes.\n",
    "\n",
    ">Words can be batched using the DataLoader, but `feature_tensors` and `target_tensors` will have to be created of the correct size and content for a given `sequence_length`.\n",
    "\n",
    "For example, say we have these as input:\n",
    "```\n",
    "words = [1, 2, 3, 4, 5, 6, 7]\n",
    "sequence_length = 4\n",
    "```\n",
    "\n",
    "Your first `feature_tensor` should contain the values:\n",
    "```\n",
    "[1, 2, 3, 4]\n",
    "```\n",
    "And the corresponding `target_tensor` should just be the next \"word\"/tokenized word value:\n",
    "```\n",
    "5\n",
    "```\n",
    "This should continue with the second `feature_tensor`, `target_tensor` being:\n",
    "```\n",
    "[2, 3, 4, 5]  # features\n",
    "6             # target\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "import numpy as np\n",
    "\n",
    "def batch_data(words, sequence_length, batch_size):\n",
    "    \"\"\"\n",
    "    Batch the neural network data using DataLoader\n",
    "    :param words: The word ids of the TV scripts\n",
    "    :param sequence_length: The sequence length of each batch\n",
    "    :param batch_size: The size of each batch; the number of sequences in a batch\n",
    "    :return: DataLoader with batched data\n",
    "    \"\"\"\n",
    "    # total number of batches\n",
    "    n_batches = len(words)//batch_size\n",
    "    # get rid of words that don't fit into a complete batch \n",
    "    words = words[:n_batches*batch_size]\n",
    "    # determine ending point for loop through words\n",
    "    targets_length = len(words) - sequence_length\n",
    "    features, targets = [], []\n",
    "    # iterate through words array, going one sequence at a time\n",
    "    for n in range(0, targets_length):\n",
    "        feature_batch = words[n:n+sequence_length]\n",
    "        target_batch = words[n+sequence_length]\n",
    "        \n",
    "        # Append to features, targets lists\n",
    "        features.append(feature_batch)\n",
    "        targets.append(target_batch)\n",
    "    \n",
    "    # Convert features, targets lists to torch tensors\n",
    "    features = torch.from_numpy(np.asarray(features))\n",
    "    targets = torch.from_numpy(np.asarray(targets))\n",
    "    \n",
    "    # Create tensor dataset and pass to create a shuffled data_loader\n",
    "    data = TensorDataset(features, targets)\n",
    "    data_loader = DataLoader(data, shuffle=True, batch_size=batch_size)\n",
    "\n",
    "    # return a dataloader\n",
    "    return data_loader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataloader testing\n",
    "\n",
    "Below, test text data is generated and a dataloader defined using the function above. Then, sample batch data of inputs `sample_x` and targets `sample_y` are obtained from our dataloader.\n",
    "\n",
    "The code below returns data that looks like the following:\n",
    "\n",
    "```\n",
    "torch.Size([10, 5])\n",
    "tensor([[ 28,  29,  30,  31,  32],\n",
    "        [ 21,  22,  23,  24,  25],\n",
    "        [ 17,  18,  19,  20,  21],\n",
    "        [ 34,  35,  36,  37,  38],\n",
    "        [ 11,  12,  13,  14,  15],\n",
    "        [ 23,  24,  25,  26,  27],\n",
    "        [  6,   7,   8,   9,  10],\n",
    "        [ 38,  39,  40,  41,  42],\n",
    "        [ 25,  26,  27,  28,  29],\n",
    "        [  7,   8,   9,  10,  11]])\n",
    "\n",
    "torch.Size([10])\n",
    "tensor([ 33,  26,  22,  39,  16,  28,  11,  43,  30,  12])\n",
    "```\n",
    "\n",
    "### Sizes\n",
    "`sample_x` should be of size `(batch_size, sequence_length)` or (10, 5) in this case and `sample_y` should just have one dimension: batch_size (10). \n",
    "\n",
    "### Values\n",
    "\n",
    "Notice that the targets, sample_y, are the *next* value in the ordered test_text data. So, for an input sequence `[ 28,  29,  30,  31,  32]` that ends with the value `32`, the corresponding output should be `33`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10, 5])\n",
      "tensor([[ 5,  6,  7,  8,  9],\n",
      "        [ 2,  3,  4,  5,  6],\n",
      "        [24, 25, 26, 27, 28],\n",
      "        [20, 21, 22, 23, 24],\n",
      "        [18, 19, 20, 21, 22],\n",
      "        [22, 23, 24, 25, 26],\n",
      "        [ 9, 10, 11, 12, 13],\n",
      "        [ 0,  1,  2,  3,  4],\n",
      "        [44, 45, 46, 47, 48],\n",
      "        [30, 31, 32, 33, 34]])\n",
      "\n",
      "torch.Size([10])\n",
      "tensor([10,  7, 29, 25, 23, 27, 14,  5, 49, 35])\n"
     ]
    }
   ],
   "source": [
    "# test dataloader\n",
    "\n",
    "test_text = range(50)\n",
    "t_loader = batch_data(test_text, sequence_length=5, batch_size=10)\n",
    "\n",
    "data_iter = iter(t_loader)\n",
    "sample_x, sample_y = data_iter.next()\n",
    "\n",
    "print(sample_x.shape)\n",
    "print(sample_x)\n",
    "print()\n",
    "print(sample_y.shape)\n",
    "print(sample_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## The Neural Network\n",
    "An RNN is implemented below using PyTorch's [Module class](http://pytorch.org/docs/master/nn.html#torch.nn.Module). An LSTM is used. \n",
    "\n",
    "**The output of this model is the *last* batch of word scores** after a complete sequence has been processed. That is, for each input sequence of words, the function outputs the word scores for a single, most likely, next word.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tests Passed\n"
     ]
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "class RNN(nn.Module):\n",
    "    \n",
    "    def __init__(self, vocab_size, output_size, embedding_dim, hidden_dim, n_layers, dropout=0.5):\n",
    "        \"\"\"\n",
    "        Initialize the PyTorch RNN Module\n",
    "        :param vocab_size: The number of input dimensions of the neural network (the size of the vocabulary)\n",
    "        :param output_size: The number of output dimensions of the neural network\n",
    "        :param embedding_dim: The size of embeddings, should you choose to use them        \n",
    "        :param hidden_dim: The size of the hidden layer outputs\n",
    "        :param dropout: dropout to add in between LSTM/GRU layers\n",
    "        \"\"\"\n",
    "        super(RNN, self).__init__()\n",
    "        \n",
    "        # Class variables needed throughout model\n",
    "        self.n_hidden = hidden_dim\n",
    "        self.drop_prob = dropout\n",
    "        self.n_layers = n_layers\n",
    "        self.output_size = output_size\n",
    "        \n",
    "        # Create embedding layer to reduce dimensionality\n",
    "        self.embedding = nn.Embedding(vocab_size, embedding_dim)\n",
    "\n",
    "        # LSTM instantiated\n",
    "        self.lstm = nn.LSTM(embedding_dim, hidden_dim, n_layers, dropout=dropout, batch_first=True)\n",
    "        \n",
    "        # Dropout layer to use before last, fully-connected layer\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        # last, fully-connected layer\n",
    "        self.fc = nn.Linear(hidden_dim, output_size)\n",
    "    \n",
    "    \n",
    "    def forward(self, nn_input, hidden):\n",
    "        \"\"\"\n",
    "        Forward propagation of the neural network\n",
    "        :param nn_input: The input to the neural network\n",
    "        :param hidden: The hidden state        \n",
    "        :return: Two Tensors, the output of the neural network and the latest hidden state\n",
    "        \"\"\"\n",
    "        \n",
    "        # First dimension of nn_input is batch_size. \n",
    "        batch_size = nn_input.shape[0]\n",
    "        nn_input = nn_input.long()\n",
    "        \n",
    "        # Convert nn_input to embeddings\n",
    "        embedded_input = self.embedding(nn_input)\n",
    "        \n",
    "        # return one batch of output word scores and the hidden state\n",
    "        r_output, hidden = self.lstm(embedded_input, hidden)\n",
    "        \n",
    "        # Stack up LSTM outputs using view\n",
    "        r_output = r_output.contiguous().view(-1, self.n_hidden)\n",
    "        \n",
    "        ## pass r_output through the dropout layer\n",
    "        out = self.dropout(r_output)\n",
    "        \n",
    "        # send out through the last fully-connected layer\n",
    "        out = self.fc(out)\n",
    "        \n",
    "        # reshape into (batch_size, seq_length, output_size)\n",
    "        output = out.view(batch_size, -1, self.output_size)\n",
    "        # get last batch\n",
    "        out = output[:, -1]\n",
    "        # return the final output and the hidden state\n",
    "        return out, hidden\n",
    "    \n",
    "    \n",
    "    def init_hidden(self, batch_size):\n",
    "        '''\n",
    "        Initialize the hidden state of an LSTM/GRU\n",
    "        :param batch_size: The batch_size of the hidden state\n",
    "        :return: hidden state of dims (n_layers, batch_size, hidden_dim)\n",
    "        '''\n",
    "        # initialize all hidden states to zeros \n",
    "        weight = next(self.parameters()).data\n",
    "        \n",
    "        if (train_on_gpu):\n",
    "            hidden = (weight.new(self.n_layers, batch_size, self.n_hidden).zero_().cuda(),\n",
    "                  weight.new(self.n_layers, batch_size, self.n_hidden).zero_().cuda())\n",
    "        else:\n",
    "            hidden = (weight.new(self.n_layers, batch_size, self.n_hidden).zero_(),\n",
    "                      weight.new(self.n_layers, batch_size, self.n_hidden).zero_())\n",
    "        \n",
    "        return hidden\n",
    "\n",
    "tests.test_rnn(RNN, train_on_gpu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Forward and Backpropagation\n",
    "\n",
    "The RNN class is utilized to apply forward and back propagation. This function will be called, iteratively, in the training loop as follows:\n",
    "```\n",
    "loss = forward_back_prop(decoder, decoder_optimizer, criterion, inp, target)\n",
    "```\n",
    "\n",
    "It returns the average loss over a batch and the hidden state returned by a call to `RNN(inp, hidden)`. \n",
    "\n",
    "**If a GPU is available, you should move your data to that GPU device, here.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tests Passed\n"
     ]
    }
   ],
   "source": [
    "def forward_back_prop(rnn, optimizer, criterion, inp, target, hidden):\n",
    "    \"\"\"\n",
    "    Forward and backward propagation on the neural network\n",
    "    :param rnn: The PyTorch Module that holds the neural network\n",
    "    :param optimizer: The PyTorch optimizer for the neural network\n",
    "    :param criterion: The PyTorch loss function\n",
    "    :param inp: A batch of input to the neural network\n",
    "    :param target: The target output for the batch of input\n",
    "    :return: The loss and the latest hidden state Tensor\n",
    "    \"\"\"\n",
    "    \n",
    "    # set rnn to train mode\n",
    "    rnn.train()\n",
    "    \n",
    "    # Creating new variables for the hidden state\n",
    "    hidden = tuple([each.data for each in hidden])\n",
    "    \n",
    "    # move data to GPU, if available\n",
    "    if(train_on_gpu):\n",
    "        rnn.cuda()\n",
    "        inp, target = inp.cuda(), target.cuda()\n",
    "        \n",
    "    # perform backpropagation and optimization\n",
    "    # reset accumulated gradients\n",
    "    rnn.zero_grad()\n",
    "    \n",
    "    # get the output from the model\n",
    "    output, h = rnn(inp, hidden)\n",
    "    \n",
    "    # calculate the loss and perform backprop\n",
    "    loss = criterion(output, target.long())\n",
    "    loss.backward()\n",
    "    \n",
    "    # Set loss to float\n",
    "    loss = float(loss.detach().cpu().numpy())\n",
    "    \n",
    "    # Set clip to 5 to avoid exploding gradient\n",
    "    nn.utils.clip_grad_norm_(rnn.parameters(), 5)\n",
    "    optimizer.step()\n",
    "\n",
    "    # return the loss over a batch and the hidden state produced by our model\n",
    "    return loss, h\n",
    "\n",
    "# Note that these tests aren't completely extensive.\n",
    "tests.test_forward_back_prop(RNN, forward_back_prop, train_on_gpu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neural Network Training\n",
    "\n",
    "With the structure of the network complete and data ready to be fed in the neural network, it's time to train it.\n",
    "\n",
    "### Train Loop\n",
    "\n",
    "The training loop is implemented in the `train_decoder` function. This function will train the network over all the batches for the number of epochs given. The model progress will be shown every number of batches. This number is set with the `show_every_n_batches` parameter. This parameter along with other parameters are set in the next section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_rnn(rnn, batch_size, optimizer, criterion, n_epochs, show_every_n_batches=100):\n",
    "    batch_losses = []\n",
    "    \n",
    "    rnn.train()\n",
    "\n",
    "    print(\"Training for %d epoch(s)...\" % n_epochs)\n",
    "    for epoch_i in range(1, n_epochs + 1):\n",
    "        \n",
    "        # initialize hidden state\n",
    "        hidden = rnn.init_hidden(batch_size)\n",
    "        \n",
    "        for batch_i, (inputs, labels) in enumerate(train_loader, 1):\n",
    "            \n",
    "            # make sure you iterate over completely full batches, only\n",
    "            n_batches = len(train_loader.dataset)//batch_size\n",
    "            if(batch_i > n_batches):\n",
    "                break\n",
    "            \n",
    "            # forward, back prop\n",
    "            loss, hidden = forward_back_prop(rnn, optimizer, criterion, inputs, labels, hidden)          \n",
    "            # record loss\n",
    "            batch_losses.append(loss)\n",
    "\n",
    "            # printing loss stats\n",
    "            if batch_i % show_every_n_batches == 0:\n",
    "                print('Epoch: {:>4}/{:<4}  Loss: {}\\n'.format(\n",
    "                    epoch_i, n_epochs, np.average(batch_losses)))\n",
    "                batch_losses = []\n",
    "\n",
    "    # returns a trained rnn\n",
    "    return rnn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperparameters\n",
    "\n",
    "Set and train the neural network with the following parameters:\n",
    "- Set `sequence_length` to the length of a sequence.\n",
    "- Set `batch_size` to the batch size.\n",
    "- Set `num_epochs` to the number of epochs to train for.\n",
    "- Set `learning_rate` to the learning rate for an Adam optimizer.\n",
    "- Set `vocab_size` to the number of unique tokens in our vocabulary.\n",
    "- Set `output_size` to the desired size of the output.\n",
    "- Set `embedding_dim` to the embedding dimension; smaller than the vocab_size.\n",
    "- Set `hidden_dim` to the hidden dimension of your RNN.\n",
    "- Set `n_layers` to the number of layers/cells in your RNN.\n",
    "- Set `show_every_n_batches` to the number of batches at which the neural network should print progress.\n",
    "\n",
    "If the network isn't getting the desired results, tweak these parameters and/or the layers in the `RNN` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data params\n",
    "# Sequence Length\n",
    "sequence_length =  50 # of words in a sequence\n",
    "# Batch Size\n",
    "batch_size = 512\n",
    "\n",
    "# data loader - do not change\n",
    "train_loader = batch_data(int_text, sequence_length, batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training parameters\n",
    "# Number of Epochs\n",
    "num_epochs = 25\n",
    "# Learning Rate\n",
    "learning_rate = 0.003\n",
    "\n",
    "# Model parameters\n",
    "# Vocab size\n",
    "vocab_size = len(vocab_to_int)\n",
    "# Output size\n",
    "output_size = vocab_size\n",
    "# Embedding Dimension\n",
    "embedding_dim = 400\n",
    "# Hidden Dimension\n",
    "hidden_dim = 512\n",
    "# Number of RNN Layers\n",
    "n_layers = 2\n",
    "\n",
    "# Show stats for every n number of batches\n",
    "show_every_n_batches = 500"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train\n",
    "In the next cell, the neural network on the pre-processed data.  If you have a hard time getting a good loss, you may consider changing your hyperparameters. In general, you may get better results with larger hidden and n_layer dimensions, but larger models take a longer time to train. \n",
    "> **You should aim for a loss less than 3.5.** \n",
    "\n",
    "You should also experiment with different sequence lengths, which determine the size of the long range dependencies that a model can learn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RNN(\n",
      "  (embedding): Embedding(21388, 400)\n",
      "  (lstm): LSTM(400, 512, num_layers=2, batch_first=True, dropout=0.5)\n",
      "  (dropout): Dropout(p=0.5, inplace=False)\n",
      "  (fc): Linear(in_features=512, out_features=21388, bias=True)\n",
      ")\n",
      "Training for 25 epoch(s)...\n",
      "Epoch:    1/25    Loss: 4.889556662559509\n",
      "\n",
      "Epoch:    1/25    Loss: 4.4328992567062375\n",
      "\n",
      "Epoch:    1/25    Loss: 4.311304751396179\n",
      "\n",
      "Epoch:    2/25    Loss: 4.165332223400574\n",
      "\n",
      "Epoch:    2/25    Loss: 4.09853772687912\n",
      "\n",
      "Epoch:    2/25    Loss: 4.103515904903412\n",
      "\n",
      "Epoch:    3/25    Loss: 3.9996245521926364\n",
      "\n",
      "Epoch:    3/25    Loss: 3.966014447212219\n",
      "\n",
      "Epoch:    3/25    Loss: 3.9797478790283205\n",
      "\n",
      "Epoch:    4/25    Loss: 3.893775116898592\n",
      "\n",
      "Epoch:    4/25    Loss: 3.8794877009391784\n",
      "\n",
      "Epoch:    4/25    Loss: 3.880654238700867\n",
      "\n",
      "Epoch:    5/25    Loss: 3.804357359444725\n",
      "\n",
      "Epoch:    5/25    Loss: 3.804048528671265\n",
      "\n",
      "Epoch:    5/25    Loss: 3.823679150104523\n",
      "\n",
      "Epoch:    6/25    Loss: 3.7554800780678566\n",
      "\n",
      "Epoch:    6/25    Loss: 3.737756031990051\n",
      "\n",
      "Epoch:    6/25    Loss: 3.7597178235054014\n",
      "\n",
      "Epoch:    7/25    Loss: 3.702487986419204\n",
      "\n",
      "Epoch:    7/25    Loss: 3.6965535655021666\n",
      "\n",
      "Epoch:    7/25    Loss: 3.7116192841529845\n",
      "\n",
      "Epoch:    8/25    Loss: 3.654285172540971\n",
      "\n",
      "Epoch:    8/25    Loss: 3.6581213545799254\n",
      "\n",
      "Epoch:    8/25    Loss: 3.6693184714317324\n",
      "\n",
      "Epoch:    9/25    Loss: 3.6268853819482882\n",
      "\n",
      "Epoch:    9/25    Loss: 3.619745151042938\n",
      "\n",
      "Epoch:    9/25    Loss: 3.6363104457855226\n",
      "\n",
      "Epoch:   10/25    Loss: 3.586610614207753\n",
      "\n",
      "Epoch:   10/25    Loss: 3.5850555024147033\n",
      "\n",
      "Epoch:   10/25    Loss: 3.6156932082176207\n",
      "\n",
      "Epoch:   11/25    Loss: 3.561538617459708\n",
      "\n",
      "Epoch:   11/25    Loss: 3.565244808673859\n",
      "\n",
      "Epoch:   11/25    Loss: 3.580439038276672\n",
      "\n",
      "Epoch:   12/25    Loss: 3.5403549085583603\n",
      "\n",
      "Epoch:   12/25    Loss: 3.5403131518363953\n",
      "\n",
      "Epoch:   12/25    Loss: 3.562540476799011\n",
      "\n",
      "Epoch:   13/25    Loss: 3.509363783194147\n",
      "\n",
      "Epoch:   13/25    Loss: 3.522044922828674\n",
      "\n",
      "Epoch:   13/25    Loss: 3.5491084880828856\n",
      "\n",
      "Epoch:   14/25    Loss: 3.498345553633655\n",
      "\n",
      "Epoch:   14/25    Loss: 3.505973186969757\n",
      "\n",
      "Epoch:   14/25    Loss: 3.5222069334983828\n",
      "\n",
      "Epoch:   15/25    Loss: 3.484870947646065\n",
      "\n",
      "Epoch:   15/25    Loss: 3.4856020855903624\n",
      "\n",
      "Epoch:   15/25    Loss: 3.493420084476471\n",
      "\n",
      "Epoch:   16/25    Loss: 3.4595963189154664\n",
      "\n",
      "Epoch:   16/25    Loss: 3.4709648141860963\n",
      "\n",
      "Epoch:   16/25    Loss: 3.484079159736633\n",
      "\n",
      "Epoch:   17/25    Loss: 3.4489560683890073\n",
      "\n",
      "Epoch:   17/25    Loss: 3.454232624053955\n",
      "\n",
      "Epoch:   17/25    Loss: 3.4739846868515016\n",
      "\n",
      "Epoch:   18/25    Loss: 3.4334903442264406\n",
      "\n",
      "Epoch:   18/25    Loss: 3.4428554577827453\n",
      "\n",
      "Epoch:   18/25    Loss: 3.4622821516990663\n",
      "\n",
      "Epoch:   19/25    Loss: 3.4232874692049426\n",
      "\n",
      "Epoch:   19/25    Loss: 3.4231496391296385\n",
      "\n",
      "Epoch:   19/25    Loss: 3.4527434048652648\n",
      "\n",
      "Epoch:   20/25    Loss: 3.414407213368075\n",
      "\n",
      "Epoch:   20/25    Loss: 3.4103554792404176\n",
      "\n",
      "Epoch:   20/25    Loss: 3.4340201807022095\n",
      "\n",
      "Epoch:   21/25    Loss: 3.395741096231458\n",
      "\n",
      "Epoch:   21/25    Loss: 3.4008945899009704\n",
      "\n",
      "Epoch:   21/25    Loss: 3.4240234546661377\n",
      "\n",
      "Epoch:   22/25    Loss: 3.394712853206475\n",
      "\n",
      "Epoch:   22/25    Loss: 3.38656849193573\n",
      "\n",
      "Epoch:   22/25    Loss: 3.417389041900635\n",
      "\n",
      "Epoch:   23/25    Loss: 3.3801809126227167\n",
      "\n",
      "Epoch:   23/25    Loss: 3.3795588817596434\n",
      "\n",
      "Epoch:   23/25    Loss: 3.4059394087791444\n",
      "\n",
      "Epoch:   24/25    Loss: 3.3703922099269192\n",
      "\n",
      "Epoch:   24/25    Loss: 3.372167342662811\n",
      "\n",
      "Epoch:   24/25    Loss: 3.396554504394531\n",
      "\n",
      "Epoch:   25/25    Loss: 3.3637131782356704\n",
      "\n",
      "Epoch:   25/25    Loss: 3.3581013140678406\n",
      "\n",
      "Epoch:   25/25    Loss: 3.3842420926094055\n",
      "\n",
      "Model Trained and Saved\n"
     ]
    }
   ],
   "source": [
    "# create model and move to gpu if available\n",
    "rnn = RNN(vocab_size, output_size, embedding_dim, hidden_dim, n_layers, dropout=0.5)\n",
    "print(rnn)\n",
    "if train_on_gpu:\n",
    "    rnn.cuda()\n",
    "\n",
    "# defining loss and optimization functions for training\n",
    "optimizer = torch.optim.Adam(rnn.parameters(), lr=learning_rate)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "# training the model\n",
    "trained_rnn = train_rnn(rnn, batch_size, optimizer, criterion, num_epochs, show_every_n_batches)\n",
    "\n",
    "# saving the trained model\n",
    "helper.save_model('./save/trained_rnn', trained_rnn)\n",
    "print('Model Trained and Saved')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rationale for Model Hyperparameters\n",
    "\n",
    "Sequence length was kept to 50 as to more accurately dialogue sequence in Seinfeld. \n",
    "\n",
    "For the hidden dimension, 256, 512, and 1024 were tried. 256 didn't seem to be high enough to get the model to converge below 3.5. 1024 drastically increased the training since learning more parameters and therefore slow to converge. 512 was sufficient to get the model to converge in a reasonable amount of time.\n",
    "\n",
    "n_layers was held steady at 2 and worked well with 512 hidden_dim.\n",
    "\n",
    "Batch size was set to 512 to fully utilize the GPU the model was trained on (Tesla V100)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Checkpoint\n",
    "\n",
    "After running the above training cell, your model will be saved by name, `trained_rnn`, and if you save your notebook progress, **you can pause here and come back to this code at another time**. You can resume your progress by running the next cell, which will load in our word:id dictionaries _and_ load in your saved model by name!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import helper\n",
    "import problem_unittests as tests\n",
    "\n",
    "_, vocab_to_int, int_to_vocab, token_dict = helper.load_preprocess()\n",
    "trained_rnn = helper.load_model('./save/trained_rnn')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate TV Script\n",
    "With the network trained and saved, a new \"fake\" Seinfeld TV script is generated in this section.\n",
    "\n",
    "### Generate Text\n",
    "To generate the text, the network needs to start with a single word and repeat its predictions until it reaches a set length. The `generate` function will be utilized to do this. It takes a word id to start with, `prime_id`, and generates a set length of text, `predict_len`. Also note that it uses topk sampling to introduce some randomness in choosing the most likely next word, given an output set of word scores!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "def generate(rnn, prime_id, int_to_vocab, token_dict, pad_value, predict_len=100):\n",
    "    \"\"\"\n",
    "    Generate text using the neural network\n",
    "    :param decoder: The PyTorch Module that holds the trained neural network\n",
    "    :param prime_id: The word id to start the first prediction\n",
    "    :param int_to_vocab: Dict of word id keys to word values\n",
    "    :param token_dict: Dict of puncuation tokens keys to puncuation values\n",
    "    :param pad_value: The value used to pad a sequence\n",
    "    :param predict_len: The length of text to generate\n",
    "    :return: The generated text\n",
    "    \"\"\"\n",
    "    rnn.eval()\n",
    "    \n",
    "    # create a sequence (batch_size=1) with the prime_id\n",
    "    current_seq = np.full((1, sequence_length), pad_value)\n",
    "    current_seq[-1][-1] = prime_id\n",
    "    predicted = [int_to_vocab[prime_id]]\n",
    "    \n",
    "    for _ in range(predict_len):\n",
    "        if train_on_gpu:\n",
    "            current_seq = torch.LongTensor(current_seq).cuda()\n",
    "        else:\n",
    "            current_seq = torch.LongTensor(current_seq)\n",
    "        \n",
    "        # initialize the hidden state\n",
    "        hidden = rnn.init_hidden(current_seq.size(0))\n",
    "        \n",
    "        # get the output of the rnn\n",
    "        output, _ = rnn(current_seq, hidden)\n",
    "        \n",
    "        # get the next word probabilities\n",
    "        p = F.softmax(output, dim=1).data\n",
    "        if(train_on_gpu):\n",
    "            p = p.cpu() # move to cpu\n",
    "         \n",
    "        # use top_k sampling to get the index of the next word\n",
    "        top_k = 5\n",
    "        p, top_i = p.topk(top_k)\n",
    "        top_i = top_i.numpy().squeeze()\n",
    "        \n",
    "        # select the likely next word index with some element of randomness\n",
    "        p = p.numpy().squeeze()\n",
    "        word_i = np.random.choice(top_i, p=p/p.sum())\n",
    "        \n",
    "        # retrieve that word from the dictionary\n",
    "        word = int_to_vocab[word_i]\n",
    "        predicted.append(word)     \n",
    "        \n",
    "        if(train_on_gpu):\n",
    "            current_seq = current_seq.cpu() # move to cpu\n",
    "        # the generated word becomes the next \"current sequence\" and the cycle can continue\n",
    "        if train_on_gpu:\n",
    "            current_seq = current_seq.cpu()\n",
    "        current_seq = np.roll(current_seq, -1, 1)\n",
    "        current_seq[-1][-1] = word_i\n",
    "    \n",
    "    gen_sentences = ' '.join(predicted)\n",
    "    \n",
    "    # Replace punctuation tokens\n",
    "    for key, token in token_dict.items():\n",
    "        ending = ' ' if key in ['\\n', '(', '\"'] else ''\n",
    "        gen_sentences = gen_sentences.replace(' ' + token.lower(), key)\n",
    "    gen_sentences = gen_sentences.replace('\\n ', '\\n')\n",
    "    gen_sentences = gen_sentences.replace('( ', '(')\n",
    "    \n",
    "    # return all the sentences\n",
    "    return gen_sentences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate a New Script\n",
    "It's time to generate the text. Set `gen_length` to the length of TV script you want to generate and set `prime_word` to one of the following to start the prediction:\n",
    "- \"jerry\"\n",
    "- \"elaine\"\n",
    "- \"george\"\n",
    "- \"kramer\"\n",
    "\n",
    "You can set the prime word to _any word_ in our dictionary, but it's best to start with a name for generating a TV script. (You can also start with any other names you find in the original text file!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "jerry:\n",
      "\n",
      "jerry: well, i'm sorry.\n",
      "\n",
      "george:(to george) what do you want?\n",
      "\n",
      "jerry: i can't believe it, but i just don't want it.\n",
      "\n",
      "george: well, i don't know what i want to say. i'm not really sure you are.\n",
      "\n",
      "jerry: i don't know. i don't know.\n",
      "\n",
      "kramer: oh, well, what?\n",
      "\n",
      "jerry: oh, you think you should have said it.\n",
      "\n",
      "george: well i think she had a big deal, and i didn't know you were gonna go to the flea market...\n",
      "\n",
      "george: you know, i have a big salad.\n",
      "\n",
      "jerry: oh, i can't believe this.(she takes a bite)\n",
      "\n",
      "elaine: i know, i think we are. i think i could be able to get it.\n",
      "\n",
      "elaine: well, i don't know why you have to go back in the apartment.\n",
      "\n",
      "jerry: i don't know, but, i was just curious, it might be nice to see you, george.\n",
      "\n",
      "jerry:(to jerry) i can't believe it.\n",
      "\n",
      "george: i know what i was saying, but i don't think so.\n",
      "\n",
      "jerry: oh, yeah, i know.(he leaves)\n",
      "\n",
      "kramer: hey.\n",
      "\n",
      "kramer: hey, hey, hey, what is it?\n",
      "\n",
      "jerry: oh, i don't wanna see.\n",
      "\n",
      "jerry:(looking in a little tone) hey, what do you want?\n",
      "\n",
      "jerry: oh, no.\n",
      "\n",
      "elaine: oh, i got a lot of trouble.\n",
      "\n",
      "kramer: yeah.\n",
      "\n",
      "george:(looking in the direction) oh, i don't want to talk to him.\n",
      "\n",
      "george: well, i don't wanna get out of my mind.\n",
      "\n",
      "jerry:(interrupting) you can't go back to florida, but i can't take it!\n",
      "\n",
      "jerry: well...\n",
      "\n",
      "elaine:(interrupting) well, it's not a problem. i mean i got it.\n",
      "\n",
      "george: i\n"
     ]
    }
   ],
   "source": [
    "# run the cell multiple times to get different results!\n",
    "gen_length = 400 # modify the length to your preference\n",
    "prime_word = 'jerry' # name for starting the script\n",
    "\n",
    "pad_word = helper.SPECIAL_WORDS['PADDING']\n",
    "generated_script = generate(trained_rnn, vocab_to_int[prime_word + ':'], int_to_vocab, token_dict, vocab_to_int[pad_word], gen_length)\n",
    "print(generated_script)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save your favorite scripts\n",
    "\n",
    "Once you have a script that you like (or find interesting), save it to a text file!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save script to a text file\n",
    "f =  open(\"generated_script_1.txt\",\"w\")\n",
    "f.write(generated_script)\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The TV Script is Not Perfect\n",
    "It's ok if the TV script doesn't make perfect sense. It should look like alternating lines of dialogue, here is one such example of a few generated lines.\n",
    "\n",
    "### Example generated script\n",
    "\n",
    ">jerry: what about me?\n",
    ">\n",
    ">jerry: i don't have to wait.\n",
    ">\n",
    ">kramer:(to the sales table)\n",
    ">\n",
    ">elaine:(to jerry) hey, look at this, i'm a good doctor.\n",
    ">\n",
    ">newman:(to elaine) you think i have no idea of this...\n",
    ">\n",
    ">elaine: oh, you better take the phone, and he was a little nervous.\n",
    ">\n",
    ">kramer:(to the phone) hey, hey, jerry, i don't want to be a little bit.(to kramer and jerry) you can't.\n",
    ">\n",
    ">jerry: oh, yeah. i don't even know, i know.\n",
    ">\n",
    ">jerry:(to the phone) oh, i know.\n",
    ">\n",
    ">kramer:(laughing) you know...(to jerry) you don't know.\n",
    "\n",
    "You can see that there are multiple characters that say (somewhat) complete sentences, but it doesn't have to be perfect! It takes quite a while to get good results, and often, you'll have to use a smaller vocabulary (and discard uncommon words), or get more data.  The Seinfeld dataset is about 3.4 MB, which is big enough for our purposes; for script generation you'll want more than 1 MB of text, generally. "
   ]
  }
 ],
 "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.8.5"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}