P05-C01-simple-rnn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  Recurrent Neural Networks (RNNs) for Language Modeling\n",
    "\n",
    "In previous tutorials, we worked with *feedforward* neural networks. They're called feedforward networks because each layer feeds into the next layer in a chain connecting the inputs to the outputs.\n",
    "\n",
    "\n",
    "![](img/multilayer-perceptron.png)\n",
    "\n",
    "At each iteration $t$, we feed in a new example $x_t$, by setting the values of the input nodes (orange). We then *feed the activation forward* by successively calculating the activations of each higher layer in the network. Finally, we read the outputs from the topmost layer. \n",
    "\n",
    "So when we feed the next example $x_{t+1}$, we overwrite all of the previous activations. If consecutive inputs to our network have no special relationship to each other (say, images uploaded by unrelated users), then this is perfectly acceptable behavior. But what if our inputs exhibit a seqeuntial relationship?\n",
    "\n",
    "Say for example that you want to predict the next character in a string of text. We might devide to feed each character into the neural network with the goal of predicting the succeeding character. \n",
    "\n",
    "![](img/recurrent-motivation.png)\n",
    "\n",
    "In the above example, the neural network forgets the previous context every time you feed a new input. How is the neural network supposed to know that \"e\" is followed by a space? It's hard to see why that should be so probable if you didn't know that the \"e\" was the final letter in the word \"Time\". \n",
    "\n",
    "Recurrent neural networks provide a slick way to incorporate sequential structure. At each time step $t$, each hidden layer $h_t$ (typically) will receive input from both the current input $x_t$ and from *that same hidden layer* at the previous time step $h_{t-1}$\n",
    "\n",
    "![](img/recurrent-lm.png)\n",
    "\n",
    "Now, when our net is trying to predict what comes after the \"e\" in time, it has access to its previous *beliefs*, and by extension, the entire history of inputs. Zooming back in to see how the nodes in a basic RNN are connected, you'll see that each node in the hidden layer is connected to each node at the hidden layer at the next time step:\n",
    "\n",
    "![](img/simple-rnn.png)\n",
    "\n",
    "Even though the neural network contains loops (the hidden layer is connected to itself), because this connection spans a time step our network is still technically a feedforward network. Thus we can still train by backpropagration just as we normally would with an MLP. Typically the loss function will be an average of the losses at each time step.\n",
    "\n",
    "In this tutorial, we're going to roll up our sleeves and write a simple RNN in MXNet using nothing but ``mxnet.ndarray`` and ``mxnet.autograd``. In practice, unless you're trying to develop fundamentlally new recurrent layers, you'll want to use the prebuilt layers that call down to extremely optimized primitives. You'll also want to rely some pre-built batching code because batching sequences can be a pain. But we think in general, if you're going to work with this stuff, and have a modicum of self respect, you'll want to implement from scratch and understand how it works at a reasonably low level. \n",
    "\n",
    "Let's go ahead and import our dependencies and specify our context. If you've been following along without a GPU until now, this might be where you'll want to get your hands on some faster hardware. GPU instances are available by the hour through Amazon Web Services. A single GPU via a [p2 instance](https://aws.amazon.com/ec2/instance-types/p2/) (NVIDIA K80s) or even an older g2 instance will be perfetly adequate for this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import mxnet as mx\n",
    "from mxnet import nd, autograd\n",
    "import numpy as np\n",
    "mx.random.seed(1)\n",
    "ctx = mx.gpu(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset: \"The Time Machine\" \n",
    "\n",
    "Now mess with some data. I grabbed a copy of the ``Time Machine``, mostly because it's available freely thanks to the good people at [Project Gutenberg](http://www.gutenberg.org) and a lot of people are tired of seeing RNNs generate Shakespeare. In case you prefer to torturing Shakespeare to torturing H.G. Wells, I've also included Andrej Karpathy's tinyshakespeare.txt in the data folder. Let's get started by reading in the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(\"data/nlp/timemachine.txt\") as f:\n",
    "    time_machine = f.read()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And you'll probably want to get a taste for what the text looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Project Gutenberg's The Time Machine, by H. G. (Herbert George) Wells\n",
      "\n",
      "This eBook is for the use of anyone anywhere at no cost and with\n",
      "almost no restrictions whatsoever.  You may copy it, give it away or\n",
      "re-use it under the terms of the Project Gutenberg License included\n",
      "with this eBook or online at www.gutenberg.net\n",
      "\n",
      "\n",
      "Title: The Time Machine\n",
      "\n",
      "Author: H. G. (Herbert George) Wells\n",
      "\n",
      "Release Date: October 2, 2004 [EBook #35]\n",
      "[Last updated: October 3, 2014]\n",
      "\n",
      "Language: English\n",
      "\n",
      "\n",
      "*** START OF THIS PR\n"
     ]
    }
   ],
   "source": [
    "print(time_machine[0:500])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tidying up\n",
    "\n",
    "I went through and discovered that the last 38083 characters consist entirely of legalese from the Gutenberg gang. So let's chop that off lest our language model learn to generate such boring drivel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "End of Project Gutenberg's The Time Machine, by H. G. (Herbert George) Wells\n",
      "\n",
      "*** END OF THIS PROJECT GUTENBERG EBOOK THE TIME MACHINE ***\n",
      "\n",
      "***** This file should be named 35.txt or 35.zip *****\n",
      "This and all associated files of various formats will be found in:\n",
      "        http://www.gutenberg.net/3/35/\n",
      "\n",
      "\n",
      "\n",
      "        Updated editions will replace the previous one--the old editions\n",
      "        will be renamed.\n",
      "\n",
      "        Creating the works from public domain print editions means that no\n",
      "        one owns a United States copyright in these works, so the Foundation\n",
      "        (and you!) c\n"
     ]
    }
   ],
   "source": [
    "print(time_machine[-38075:-37500])\n",
    "time_machine = time_machine[:-38083]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numerical representations of characters\n",
    "\n",
    "When we create numerical representations of characters, we'll use one-hot representations. A one-hot is a vector that taked value 1 in the index corresponding to a character, and 0 elsewhere. Because this vector is as long as the vocab, let's get a definitive list of characters in this dataset so that our representation is not longer than necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['W', '8', 'c', 'T', '?', '\"', 'b', 'R', '[', 'q', 'M', '1', 'n', 'Q', 'L', 's', \"'\", 'f', 'K', ':', 'r', 'k', '4', ';', ']', 'J', 'p', '_', 'F', 'z', 'u', 'd', '(', 'v', '0', 'O', 'S', 'H', 'B', '2', 'i', 'V', '9', 'G', '3', 'a', '*', 'j', 'N', 't', 'h', '-', 'e', 'X', 'w', ',', 'Y', '5', 'I', 'P', ')', 'm', 'l', ' ', 'E', '!', 'x', 'y', '.', 'U', 'D', 'o', '#', 'A', '\\n', 'g', 'C']\n",
      "Length of vocab: 77\n"
     ]
    }
   ],
   "source": [
    "character_list = list(set(time_machine))\n",
    "vocab_size = len(character_list)\n",
    "print(character_list)\n",
    "print(\"Length of vocab: %s\" % vocab_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll often want to access the index corresponding to each character quickly so let's store this as a dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'W': 0, '8': 1, '?': 4, '\"': 5, '*': 46, 'u': 30, 'b': 6, 'R': 7, 'q': 9, 'M': 10, 'n': 12, 'Q': 13, 'Y': 56, 'L': 14, 's': 15, '5': 57, \"'\": 16, 'f': 17, 'c': 2, ':': 19, 'r': 20, 'k': 21, '4': 22, ';': 23, 'T': 3, ']': 24, 'J': 25, 'p': 26, '_': 27, 'F': 28, 'z': 29, 'd': 31, '(': 32, '0': 34, 'C': 76, 'O': 35, 'K': 18, 'S': 36, 'H': 37, 'B': 38, '2': 39, 'V': 41, 'U': 69, '9': 42, 'G': 43, '3': 44, 'a': 45, 'j': 47, 'N': 48, 't': 49, 'h': 50, 'e': 52, 'X': 53, '1': 11, 'w': 54, ',': 55, '-': 51, 'A': 73, 'I': 58, 'P': 59, ')': 60, 'm': 61, ' ': 63, 'E': 64, '!': 65, 'x': 66, 'y': 67, '.': 68, 'l': 62, 'i': 40, 'D': 70, 'v': 33, 'o': 71, '#': 72, '[': 8, '\\n': 74, 'g': 75}\n"
     ]
    }
   ],
   "source": [
    "character_dict = {}\n",
    "for e, char in enumerate(character_list):\n",
    "    character_dict[char] = e\n",
    "print(character_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "time_numerical = [character_dict[char] for char in time_machine]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "179533\n",
      "[59, 20, 71, 47, 52, 2, 49, 63, 43, 30, 49, 52, 12, 6, 52, 20, 75, 16, 15, 63]\n",
      "Project Gutenberg's The Time Machine, b\n"
     ]
    }
   ],
   "source": [
    "#########################\n",
    "#  Check that the length is right\n",
    "#########################\n",
    "print(len(time_numerical))\n",
    "\n",
    "#########################\n",
    "#  Check that the format looks right\n",
    "#########################\n",
    "print(time_numerical[:20])\n",
    "\n",
    "#########################\n",
    "#  Convert back to text\n",
    "#########################\n",
    "print(\"\".join([character_list[idx] for idx in time_numerical[:39]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## One-hot representations\n",
    "\n",
    "We can use NDArray's one_hot() render a one-hot representation of each character. But frack it, since this is the from scratch tutorial, let's right this ourselves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def one_hots(numerical_list, vocab_size=vocab_size):\n",
    "    result = []\n",
    "    for idx in numerical_list:\n",
    "        array = nd.zeros(shape=(1, vocab_size), ctx=ctx)\n",
    "        array[0, idx] = 1.0\n",
    "        result.append(array)\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[\n",
      "[[ 0.  0.  0.  0.  0.  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.  0.  0.  0.  0.  0.\n",
      "   0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.\n",
      "   0.  0.  0.  0.  0.  1.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.\n",
      "   0.  0.  0.  0.  0.]]\n",
      "<NDArray 1x77 @gpu(0)>, \n",
      "[[ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.\n",
      "   0.  0.  1.  0.  0.  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.  0.  0.  0.  0.  0.\n",
      "   0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.\n",
      "   0.  0.  0.  0.  0.]]\n",
      "<NDArray 1x77 @gpu(0)>]\n"
     ]
    }
   ],
   "source": [
    "print(one_hots(time_numerical[:2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That looks about right. Now let's write a function to convert our one-hots back to readable text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def textify(vector_list):\n",
    "    result = \"\"\n",
    "    for vector in vector_list:\n",
    "        vector = vector[0]\n",
    "        result += character_list[int(nd.argmax(vector, axis=0).asscalar())]\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Project Gutenberg's The Time Machine, by\""
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "textify(one_hots(time_numerical[0:40]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing the data for training\n",
    "\n",
    "Great, it's not the most efficient implementation, but we know how it works. So we're already doing better than the majority of people with job titles in machine learning. Now, let's chop up our dataset in to sequences we could feed into our model.\n",
    "\n",
    "You might think we could just feed in the entire dataset as one gigantic input and backpropagate across the entire sequence. When you try to backpropagate across thousands of steps a few things go wrong:\n",
    "(1) The time it takes to compute a single gradient update will be unreasonably long\n",
    "(2) The gradient across thousands of recurrent steps has a tendency to either blow up, causing NaN errors due to losing precision, or to vanish.\n",
    "\n",
    "Thus we're going to look at feeding in our data in reasonably short sequences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Project Gutenberg's The Time Machine, by H. G. (Herbert George) \""
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seq_length = 64\n",
    "dataset = [one_hots(time_numerical[i*seq_length:(i+1)*seq_length]) for i in range(int(np.floor((len(time_numerical)-1)/seq_length)))]\n",
    "textify(dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've chopped our dataset into sequences of length ``seq_length``, at every time step, our input is a single one-hot vector. This means that our computation of the hidden layer would consist of matrix-vector multiplications, which are not expecially efficient on GPU. To take advantage of the available computing resources, we'll want to feed through a batch of sequences at the same time. The following code may look tricky but it's just some plumbing to make the data look like this.\n",
    "\n",
    "![](img/recurrent-batching.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2805\n",
      "87\n"
     ]
    }
   ],
   "source": [
    "print(len(dataset))\n",
    "sequences_per_batch_row = int(np.floor(len(dataset))/batch_size)\n",
    "print(sequences_per_batch_row)\n",
    "data_rows = [dataset[i*sequences_per_batch_row:i*sequences_per_batch_row+sequences_per_batch_row] \n",
    "            for i in range(batch_size)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Let's sanity check that everything went the way we hop. For each data_row, the second sequence should follow the first:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "***Batch 0:***\n",
      " Project Gutenberg's The Time Machine, by H. G. (Herbert George) Wells\n",
      "\n",
      "This eBook is for the use of anyone anywhere at no cost a \n",
      "\n",
      "\n",
      "***Batch 1:***\n",
      " vement of the barometer. Yesterday it was so high, yesterday night\n",
      "it fell, then this morning it rose again, and so gently upwar \n",
      "\n",
      "\n",
      "***Batch 2:***\n",
      " andlesticks upon the mantel and several in sconces, so that\n",
      "the room was brilliantly illuminated. I sat in a low arm-chair\n",
      "neare \n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i in range(3):\n",
    "    print(\"***Batch %s:***\\n %s \\n\\n\" % (i, textify(data_rows[i][0]) + textify(data_rows[i][1])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's stack these data_rows together to form our batches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def stack_the_datasets(datasets):\n",
    "    full_dataset = []\n",
    "    # iterate over the sequences\n",
    "    for s in range(len(datasets[0])):\n",
    "        sequence = []\n",
    "        # iterate over the elements of the seqeunce\n",
    "        for elem in range(len(datasets[0][0])):\n",
    "            sequence.append(nd.concatenate([ds[s][elem].reshape((1,-1)) for ds in datasets], axis=0))\n",
    "        full_dataset.append(sequence)\n",
    "    return(full_dataset)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_data = stack_the_datasets(data_rows)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And let's check that the data stacking procedure worked as expected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(32, 77)\n",
      "Seq 0, Batch 0 \"Project Gutenberg's The Time Machine, by H. G. (Herbert George) \"\n",
      "Seq 1, Batch 0 \"Wells\n",
      "\n",
      "This eBook is for the use of anyone anywhere at no cost a\"\n"
     ]
    }
   ],
   "source": [
    "print(training_data[0][0].shape)\n",
    "print(\"Seq 0, Batch 0 \\\"%s\\\"\" % textify([training_data[0][i][0].reshape((1,-1)) for i in range(seq_length)]))\n",
    "print(\"Seq 1, Batch 0 \\\"%s\\\"\" % textify([training_data[1][i][0].reshape((1,-1)) for i in range(seq_length)]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing our labels\n",
    "\n",
    "Now let's repurpose the same batching code to create our label batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(32, 77)\n"
     ]
    }
   ],
   "source": [
    "labels = [one_hots(time_numerical[i*seq_length+1:(i+1)*seq_length+1]) for i in range(int(np.floor((len(time_numerical)-1)/seq_length)))]\n",
    "label_rows = [labels[i*sequences_per_batch_row:i*sequences_per_batch_row+sequences_per_batch_row] for i in range(batch_size)]\n",
    "training_labels = stack_the_datasets(label_rows)\n",
    "print(training_labels[0][0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A final sanity check\n",
    "\n",
    "Remember that our target at every time step is to predict the next character in the sequence. So our labels should look just like our inputs but offset by one character. Let's look at corresponding inputs and outputs to make sure everything lined up as expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "andlesticks upon the mantel and several in sconces, so that\n",
      "the \n",
      "ndlesticks upon the mantel and several in sconces, so that\n",
      "the r\n"
     ]
    }
   ],
   "source": [
    "print(textify([training_data[0][i][2].reshape((1,-1)) for i in range(seq_length)]))\n",
    "print(textify([training_labels[0][i][2].reshape((1,-1)) for i in range(seq_length)]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recurrent neural networks\n",
    "\n",
    "[Explain RNN updates]\n",
    "\n",
    "![](img/simple-rnn.png)\n",
    "\n",
    "Recall that the update for an ordinary hidden layer in a neural network with activation function $phi$ is given by \n",
    "$$h = \\phi(X  W) + b$$\n",
    "\n",
    "To make this a recurrent neural network, we're simply going to add a weight sum of theprevious hidden state $h_{t-1}$:\n",
    "\n",
    "$$h_t = \\phi(X_t  W_{xh} + h_{t-1} W_{hh} + b_h )$$\n",
    "\n",
    "Then at every time set $t$, we'll calculate the output as:\n",
    "$$\\hat{y}_t = \\text{softmax}(h_t W_{hy} + b_y)$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Allocate parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "num_inputs = 77\n",
    "num_hidden = 256\n",
    "num_outputs = 77\n",
    "\n",
    "########################\n",
    "#  Weights connecting the inputs to the hidden layer\n",
    "########################\n",
    "Wxh = nd.random_normal(shape=(num_inputs,num_hidden), ctx=ctx) * .01\n",
    "\n",
    "########################\n",
    "#  Recurrent weights connecting the hidden layer across time steps\n",
    "########################\n",
    "Whh = nd.random_normal(shape=(num_hidden,num_hidden), ctx=ctx)* .01\n",
    "\n",
    "########################\n",
    "#  Bias vector for hidden layer\n",
    "########################\n",
    "bh = nd.random_normal(shape=num_hidden, ctx=ctx) * .01\n",
    "\n",
    "\n",
    "########################\n",
    "# Weights to the output nodes\n",
    "########################\n",
    "Why = nd.random_normal(shape=(num_hidden,num_inputs), ctx=ctx) * .01\n",
    "by = nd.random_normal(shape=num_inputs, ctx=ctx) * .01\n",
    "\n",
    "# NOTE: to keep notation consistent, \n",
    "# we should really use capital letters \n",
    "# for hidden layers and outputs, \n",
    "# since we are doing batchwise computations] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attach the gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "params = [Wxh, Whh, bh, Why, by]\n",
    "\n",
    "for param in params:\n",
    "    param.attach_grad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Softmax Activation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def softmax(y_linear, temperature=1.0):\n",
    "    lin = (y_linear-nd.max(y_linear)) / temperature\n",
    "    exp = nd.exp(lin)\n",
    "    partition =nd.sum(exp, axis=0, exclude=True).reshape((-1,1))\n",
    "    return exp / partition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "[[ 0.50049996  0.49949998]\n",
       " [ 0.49949998  0.50049996]]\n",
       "<NDArray 2x2 @cpu(0)>"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "####################\n",
    "# With a temperature of 1 (always 1 during training), we get back some set of proabilities\n",
    "####################\n",
    "softmax(nd.array([[1,-1],[-1,1]]), temperature=1000.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "[[ 0.50049996  0.49949998]\n",
       " [ 0.49949998  0.50049996]]\n",
       "<NDArray 2x2 @cpu(0)>"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "####################\n",
    "# If we set a high temperature, we can get more entropic (*noisier*) probabilities\n",
    "####################\n",
    "softmax(nd.array([[1,-1],[-1,1]]), temperature=1000.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "[[ 1.  0.]\n",
       " [ 0.  1.]]\n",
       "<NDArray 2x2 @cpu(0)>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "####################\n",
    "# Often we want to sample with low temperatures to produce sharp proababilities\n",
    "####################\n",
    "softmax(nd.array([[10,-10],[-10,10]]), temperature=.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def simple_rnn(inputs, state, temperature=1.0):\n",
    "    outputs = []\n",
    "    h = state\n",
    "    for X in inputs:\n",
    "        h_linear = nd.dot(X, Wxh) + nd.dot(h, Whh) + bh\n",
    "        h = nd.tanh(h_linear)\n",
    "        yhat_linear = nd.dot(h, Why) + by\n",
    "        yhat = softmax(yhat_linear, temperature=temperature) \n",
    "        outputs.append(yhat)\n",
    "    return (outputs, h)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cross-entropy loss function\n",
    "\n",
    "At every time step our task is to predict the next character, given the string up to that point. This is the familiar multi-task classification that we introduced for handwritten digit classification. Accordingly, we'll rely on the same loss function, cross-entropy. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# def cross_entropy(yhat, y):\n",
    "#     return - nd.sum(y * nd.log(yhat))\n",
    "\n",
    "def cross_entropy(yhat, y):\n",
    "    return - nd.mean(nd.sum(y * nd.log(yhat), axis=0, exclude=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "[ 0.53647929]\n",
       "<NDArray 1 @cpu(0)>"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cross_entropy(nd.array([.2,.5,.3]), nd.array([1.,0,0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Averaging the loss over the sequence\n",
    "\n",
    "Because the unfolded RNN has multiple outputs (one at every time step) we can calculate a loss at every time step. The weights corresponding to the net at time step $t$ influence both the loss at time step $t$ and the loss at time step $t+1$. To combine our losses into a single global loss, we'll take the average of the losses at each time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def average_ce_loss(outputs, labels):\n",
    "    assert(len(outputs) == len(labels))\n",
    "    total_loss = nd.array([0.], ctx=ctx)\n",
    "    for (output, label) in zip(outputs,labels):\n",
    "        total_loss = total_loss + cross_entropy(output, label)\n",
    "#         print(total_loss.shape)\n",
    "#     loss_list = [cross_entropy(outputs[i], labels[i]) for (i, _) in enumerate(outputs)]\n",
    "    return total_loss / len(outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def SGD(params, lr):    \n",
    "    for param in params:\n",
    "        param[:] = param - lr * param.grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating text by sampling\n",
    "\n",
    "We have now defined a model that takes a sequence of real inputs from our training data and tries to predict the next character at every time step. You might wonder, what can we do with this model? Why should I care about predicting the next character in a sequence of text?\n",
    "\n",
    "This capability is exciting because given such a model, we can now generate strings of plausible text. The generation procedure goes as follows. Say our string begins with the character \"T\". We can feed the letter \"T\" and get a conditional probability distribution over the next character $P(x_2|x_1=\\text{\"T\"})$. We can the sample from this distribution, e.g. producing an \"i\", and then assign $x_2=\\text{\"i\"}$, feeding this to the network at the next time step.\n",
    "\n",
    "[Add a nice graphic to illustrate sampling]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def sample(prefix, num_chars, temperature=1.0):\n",
    "    #####################################\n",
    "    # Initialize the string that we'll return to the supplied prefix\n",
    "    #####################################\n",
    "    string = prefix\n",
    "\n",
    "    #####################################\n",
    "    # Prepare the prefix as a sequence of one-hots for ingestion by RNN\n",
    "    #####################################\n",
    "    prefix_numerical = [character_dict[char] for char in prefix]\n",
    "    input = one_hots(prefix_numerical)\n",
    "    \n",
    "    #####################################\n",
    "    # Set the initial state of the hidden representation ($h_0$) to the zero vector\n",
    "    #####################################    \n",
    "    sample_state = nd.zeros(shape=(1, num_hidden), ctx=ctx)  \n",
    "\n",
    "    #####################################\n",
    "    # For num_chars iterations,\n",
    "    #     1) feed in the current input\n",
    "    #     2) sample next character from from output distribution\n",
    "    #     3) add sampled character to the decoded string\n",
    "    #     4) prepare the sampled character as a one_hot (to be the next input)\n",
    "    #####################################    \n",
    "    for i in range(num_chars):\n",
    "        outputs, sample_state = simple_rnn(input, sample_state, temperature=temperature)\n",
    "        choice = np.random.choice(77, p=outputs[-1][0].asnumpy())\n",
    "        string += character_list[choice]\n",
    "        input = one_hots([choice])\n",
    "    return string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 2000\n",
    "moving_loss = 0.\n",
    "\n",
    "learning_rate = .5\n",
    "\n",
    "# state = nd.zeros(shape=(batch_size, num_hidden), ctx=ctx)\n",
    "for e in range(epochs):\n",
    "    ############################\n",
    "    # Attenuate the learning rate by a factor of 2 every 100 epochs.\n",
    "    ############################\n",
    "    if ((e+1) % 100 == 0):\n",
    "        learning_rate = learning_rate / 2.0\n",
    "    state = nd.zeros(shape=(batch_size, num_hidden), ctx=ctx)\n",
    "    for i, (data, label) in enumerate(zip(training_data, training_labels)):\n",
    "        data_one_hot = data\n",
    "        label_one_hot = label\n",
    "        with autograd.record():\n",
    "            outputs, state = simple_rnn(data_one_hot, state)\n",
    "            loss = average_ce_loss(outputs, label_one_hot)\n",
    "            loss.backward()\n",
    "        SGD(params, learning_rate)\n",
    "\n",
    "        ##########################\n",
    "        #  Keep a moving average of the losses\n",
    "        ##########################\n",
    "        if (i == 0) and (e == 0):\n",
    "            moving_loss = np.mean(loss.asnumpy()[0])\n",
    "        else:\n",
    "            moving_loss = .99 * moving_loss + .01 * np.mean(loss.asnumpy()[0])\n",
    "      \n",
    "    print(\"Epoch %s. Loss: %s\" % (e, moving_loss)) \n",
    "    print(sample(\"The Time Ma\", 1024, temperature=.1))\n",
    "    print(sample(\"The Medical Man rose, came to the lamp,\", 1024, temperature=.1))\n",
    "            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusions\n",
    "\n",
    "Once you start running this code, it will spit out a sample at the end of each epoch. I'll leave this output cell blank so you don't see megabytes of text, but here are some patterns that I observed when I ran this code. \n",
    "\n",
    "The network seems to first work out patterns with no sequential relationship and then slowly incorporate longer and longer windows of context. After just 1 epoch, my RNN generated this:\n",
    "\n",
    ">         e       e e ee    e   eee     e     e ee   e  e      ee     e e   ee  e   e            ee    e   e   e     e  e   e     e          e   e ee e    aee    e e               ee  e     e   ee ee   e ee     e e       e e e        ete    e   e e   e e   e       ee  n eee    ee e     eeee  e e    e         e  e  e ee    e  e   e    e       e  e  eee ee      e         e            e       e    e e    ee   ee e e e   e  e  e e  e t       e  ee         e eee  e  e      e ee    e    e       e                e      eee   e  e  e   eeeee      e     eeee e e   ee ee     ee     a    e e eee           ee  e e   e e   aee           e      e     e e               eee       e           e         e     e    e e   e      e   e e   e    e    e ee e      ee                 e  e  e   e    e  e   e                    e      e   e        e     ee  e    e    ee n  e   ee   e  e         e  e         e      e    t    ee  ee  ee   eee  et     e        e     e e              ee   e  e  e     e  e  e e       e              e       e\"\n",
    "\n",
    "It's learned that spaces and \"e\"s (to my knowledge, there's no aesthetically pleasing way to spell the plural form of the letter \"e\") are the most common characters.\n",
    "\n",
    "I little bit later on it spits out strings like:\n",
    "\n",
    "> the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the\n",
    "\n",
    "At this point it's learned that after the space usually comes a nonspace character, and perhaps that \"t\" is the most common character to immediately follow a space, \"h\" to follow a \"t\" and \"e\" to follow \"th\". However it doesn't appear to be looking far enough back to realize that the word \"the\" should be very unlikely immesiately after the word \"the\"... \n",
    "\n",
    "By the 175th epoch, the model appears to be putting together a fairly large vocabulary although it the words together in ways that on might charitably describe as \"creative\".\n",
    "\n",
    "> the little people had been as I store of the sungher had leartered along the realing of the stars of the little past and stared at the thing that I had the sun had to the stars of the sunghed a stirnt a moment the sun had come and fart as the stars of the sunghed a stirnt a moment the sun had to the was completely and of the little people had been as I stood and all amations of the staring and some of the really\n",
    "\n",
    "In subsequent tutorials we'll explore sophisticated techniques for evaluating and improving language models. We'll also take a look at some related but more complicate problems like language translations and image captioning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "For whinges or inquiries, [open an issue on  GitHub.](https://github.com/zackchase/mxnet-the-straight-dope)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}