Part1_Sequence.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\newcommand{\\mat}[1]{\\boldsymbol {#1}}\n",
    "\\newcommand{\\mattr}[1]{\\boldsymbol {#1}^\\top}\n",
    "\\newcommand{\\matinv}[1]{\\boldsymbol {#1}^{-1}}\n",
    "\\newcommand{\\vec}[1]{\\boldsymbol {#1}}\n",
    "\\newcommand{\\vectr}[1]{\\boldsymbol {#1}^\\top}\n",
    "\\newcommand{\\rvar}[1]{\\mathrm {#1}}\n",
    "\\newcommand{\\rvec}[1]{\\boldsymbol{\\mathrm{#1}}}\n",
    "\\newcommand{\\diag}{\\mathop{\\mathrm {diag}}}\n",
    "\\newcommand{\\set}[1]{\\mathbb {#1}}\n",
    "\\newcommand{\\norm}[1]{\\left\\lVert#1\\right\\rVert}\n",
    "\\newcommand{\\pderiv}[2]{\\frac{\\partial #1}{\\partial #2}}\n",
    "\\newcommand{\\bb}[1]{\\boldsymbol{#1}}\n",
    "$$\n",
    "# Part 1: Sequence Models\n",
    "<a id=part1></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this part we will learn about working with text sequences using recurrent neural networks.\n",
    "We'll go from a raw text file all the way to a fully trained GRU-RNN model and generate works of art!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import unittest\n",
    "import os\n",
    "import sys\n",
    "import pathlib\n",
    "import urllib\n",
    "import shutil\n",
    "import re\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "test = unittest.TestCase()\n",
    "plt.rcParams.update({'font.size': 12})\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print('Using device:', device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text generation with a char-level RNN\n",
    "<a id=part1_1></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Obtaining the corpus\n",
    "<a id=part1_2></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's begin by downloading a corpus containing all the works of William Shakespeare.\n",
    "Since he was very prolific, this corpus is fairly large and will provide us with enough data for\n",
    "obtaining impressive results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CORPUS_URL = 'https://github.com/cedricdeboom/character-level-rnn-datasets/raw/master/datasets/shakespeare.txt'\n",
    "DATA_DIR = pathlib.Path.home().joinpath('.pytorch-datasets')\n",
    "\n",
    "def download_corpus(out_path=DATA_DIR, url=CORPUS_URL, force=False):\n",
    "    pathlib.Path(out_path).mkdir(exist_ok=True)\n",
    "    out_filename = os.path.join(out_path, os.path.basename(url))\n",
    "    \n",
    "    if os.path.isfile(out_filename) and not force:\n",
    "        print(f'Corpus file {out_filename} exists, skipping download.')\n",
    "    else:\n",
    "        print(f'Downloading {url}...')\n",
    "        with urllib.request.urlopen(url) as response, open(out_filename, 'wb') as out_file:\n",
    "            shutil.copyfileobj(response, out_file)\n",
    "        print(f'Saved to {out_filename}.')\n",
    "    return out_filename\n",
    "    \n",
    "corpus_path = download_corpus()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the text into memory and print a snippet:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(corpus_path, 'r') as f:\n",
    "    corpus = f.read()\n",
    "\n",
    "print(f'Corpus length: {len(corpus)} chars')\n",
    "print(corpus[7:1234])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Preprocessing\n",
    "<a id=part1_3></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first thing we'll need is to map from each unique character in the corpus to an index that will represent it in our learning process.\n",
    "\n",
    "**TODO**: Implement the `char_maps()` function in the `hw3/charnn.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hw3.charnn as charnn\n",
    "\n",
    "char_to_idx, idx_to_char = charnn.char_maps(corpus)\n",
    "print(char_to_idx)\n",
    "\n",
    "test.assertEqual(len(char_to_idx), len(idx_to_char))\n",
    "test.assertSequenceEqual(list(char_to_idx.keys()), list(idx_to_char.values()))\n",
    "test.assertSequenceEqual(list(char_to_idx.values()), list(idx_to_char.keys()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Seems we have some strange characters in the corpus that are very rare and are probably due to mistakes.\n",
    "To reduce the length of each tensor we'll need to later represent our chars, it's best to remove them.\n",
    "\n",
    "**TODO**: Implement the `remove_chars()` function in the `hw3/charnn.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus, n_removed = charnn.remove_chars(corpus, ['}','$','_','<','\\ufeff'])\n",
    "print(f'Removed {n_removed} chars')\n",
    "\n",
    "# After removing the chars, re-create the mappings\n",
    "char_to_idx, idx_to_char = charnn.char_maps(corpus)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next thing we need is an **embedding** of the chracters.\n",
    "An embedding is a representation of each token from the sequence as a tensor.\n",
    "For a char-level RNN, our tokens will be chars and we can thus use the simplest possible embedding: encode each char as a **one-hot** tensor. In other words, each char will be represented\n",
    "as a tensor whos length is the total number of unique chars (`V`) which contains all zeros except at the index\n",
    "corresponding to that specific char.\n",
    "\n",
    "**TODO**: Implement the functions `chars_to_onehot()` and `onehot_to_chars()` in the `hw3/charnn.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wrap the actual embedding functions for calling convenience\n",
    "def embed(text):\n",
    "    return charnn.chars_to_onehot(text, char_to_idx)\n",
    "\n",
    "def unembed(embedding):\n",
    "    return charnn.onehot_to_chars(embedding, idx_to_char)\n",
    "\n",
    "text_snippet = corpus[3104:3148]\n",
    "print(text_snippet)\n",
    "print(embed(text_snippet[0:3]))\n",
    "\n",
    "test.assertEqual(text_snippet, unembed(embed(text_snippet)))\n",
    "test.assertEqual(embed(text_snippet).dtype, torch.int8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset Creation\n",
    "<a id=part1_4></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We wish to train our model to generate text by constantly predicting what the next char should be based on the past.\n",
    "To that end we'll need to train our recurrent network in a way similar to a classification task. At each timestep, we input a char and set the expected output (label) to be the next char in the original sequence.\n",
    "\n",
    "We will split our corpus into shorter sequences of length `S` chars (try to think why; see question below).\n",
    "Each **sample** we provide our model with will therefore be a tensor of shape `(S,V)` where `V` is the embedding dimension. Our model will operate sequentially on each char in the sequence.\n",
    "For each sample, we'll also need a **label**. This is simple another sequence, shifted by one char so that the label of each char is the next char in the corpus.\n",
    "\n",
    "**TODO**: Implement the `chars_to_labelled_samples()` function in the `hw3/charnn.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create dataset of sequences\n",
    "seq_len = 64\n",
    "vocab_len = len(char_to_idx)\n",
    "\n",
    "# Create labelled samples\n",
    "samples, labels = charnn.chars_to_labelled_samples(corpus, char_to_idx, seq_len, device)\n",
    "print(f'samples shape: {samples.shape}')\n",
    "print(f'labels shape: {labels.shape}')\n",
    "\n",
    "# Test shapes\n",
    "num_samples = (len(corpus) - 1) // seq_len\n",
    "test.assertEqual(samples.shape, (num_samples, seq_len, vocab_len))\n",
    "test.assertEqual(labels.shape, (num_samples, seq_len))\n",
    "\n",
    "# Test content\n",
    "for _ in range(1000):\n",
    "    # random sample\n",
    "    i = np.random.randint(num_samples, size=(1,))[0]\n",
    "    # Compare to corpus\n",
    "    test.assertEqual(unembed(samples[i]), corpus[i*seq_len:(i+1)*seq_len], msg=f\"content mismatch in sample {i}\")\n",
    "    # Compare to labels\n",
    "    sample_text = unembed(samples[i])\n",
    "    label_text = str.join('', [idx_to_char[j.item()] for j in labels[i]])\n",
    "    test.assertEqual(sample_text[1:], label_text[0:-1], msg=f\"label mismatch in sample {i}\")\n",
    "    \n",
    "print(f'sample 100 as text:\\n{unembed(samples[100])}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As usual, instead of feeding one sample as a time into our model's forward we'll work with **batches** of samples. This means that at every timestep, our model will operate on a batch of chars that are from **different sequences**.\n",
    "Effectively this will allow us to parallelize training our model by dong matrix-matrix multiplications\n",
    "instead of matrix-vector during the forward pass.\n",
    "\n",
    "Let's use the standard PyTorch `Dataset`/`DataLoader` combo. Luckily for the dataset we can use a built-in\n",
    "class, `TensorDataset` to return tuples of `(sample, label)` from the `samples` and `labels` tensors we created above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.utils.data\n",
    "\n",
    "# Create DataLoader returning batches of samples.\n",
    "batch_size = 32\n",
    "\n",
    "ds_corpus = torch.utils.data.TensorDataset(samples, labels)\n",
    "dl_corpus = torch.utils.data.DataLoader(ds_corpus, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what that gives us:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'num batches: {len(dl_corpus)}')\n",
    "\n",
    "x0, y0 = next(iter(dl_corpus))\n",
    "print(f'shape of a batch sample: {x0.shape}')\n",
    "print(f'shape of a batch label: {y0.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Implementation\n",
    "<a id=part1_5></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, our data set is ready so we can focus on our model.\n",
    "\n",
    "We'll implement here is a multilayer gated recurrent unit (GRU) model, with dropout.\n",
    "This model is a type of RNN which performs similar to the well-known LSTM model,\n",
    "but it's somewhat easier to train because it has less parameters.\n",
    "We'll modify the regular GRU slightly by applying dropout to\n",
    "the hidden states passed between layers of the model.\n",
    "\n",
    "The model accepts an input $\\mat{X}\\in\\set{R}^{S\\times V}$ containing a sequence of embedded chars.\n",
    "It returns an output $\\mat{Y}\\in\\set{R}^{S\\times V}$ of predictions for the next char and the final hidden state\n",
    "$\\mat{H}\\in\\set{R}^{L\\times H}$. Here $S$ is the sequence length, $V$ is the vocabulary size (number of unique chars), $L$ is the number of layers in the model and $H$ is the hidden dimension."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mathematically, the model's forward function at layer $k\\in[1,L]$ and timestep $t\\in[1,S]$ can be described as\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\vec{z_t}^{[k]} &= \\sigma\\left(\\vec{x}^{[k]}_t {\\mattr{W}_{\\mathrm{xz}}}^{[k]} +\n",
    "    \\vec{h}_{t-1}^{[k]} {\\mattr{W}_{\\mathrm{hz}}}^{[k]} + \\vec{b}_{\\mathrm{z}}^{[k]}\\right) \\\\\n",
    "\\vec{r_t}^{[k]} &= \\sigma\\left(\\vec{x}^{[k]}_t {\\mattr{W}_{\\mathrm{xr}}}^{[k]} +\n",
    "    \\vec{h}_{t-1}^{[k]} {\\mattr{W}_{\\mathrm{hr}}}^{[k]} + \\vec{b}_{\\mathrm{r}}^{[k]}\\right) \\\\\n",
    "\\vec{g_t}^{[k]} &= \\tanh\\left(\\vec{x}^{[k]}_t {\\mattr{W}_{\\mathrm{xg}}}^{[k]} +\n",
    "    (\\vec{r_t}^{[k]}\\odot\\vec{h}_{t-1}^{[k]}) {\\mattr{W}_{\\mathrm{hg}}}^{[k]} + \\vec{b}_{\\mathrm{g}}^{[k]}\\right) \\\\\n",
    "\\vec{h_t}^{[k]} &= \\vec{z}^{[k]}_t \\odot \\vec{h}^{[k]}_{t-1} + \\left(1-\\vec{z}^{[k]}_t\\right)\\odot \\vec{g_t}^{[k]}\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The input to each layer is,\n",
    "$$\n",
    "\\mat{X}^{[k]} =\n",
    "\\begin{bmatrix}\n",
    "    {\\vec{x}_1}^{[k]} \\\\ \\vdots \\\\ {\\vec{x}_S}^{[k]}\n",
    "\\end{bmatrix} \n",
    "=\n",
    "\\begin{cases}\n",
    "    \\mat{X} & \\mathrm{if} ~k = 1~ \\\\\n",
    "    \\mathrm{dropout}_p \\left(\n",
    "    \\begin{bmatrix}\n",
    "        {\\vec{h}_1}^{[k-1]} \\\\ \\vdots \\\\ {\\vec{h}_S}^{[k-1]}\n",
    "    \\end{bmatrix} \\right) & \\mathrm{if} ~1 < k \\leq L+1~\n",
    "\\end{cases}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output of the entire model is then,\n",
    "$$\n",
    "\\mat{Y} = \\mat{X}^{[L+1]} {\\mattr{W}_{\\mathrm{hy}}} + \\mat{B}_{\\mathrm{y}}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and the final hidden state is\n",
    "$$\n",
    "\\mat{H} = \n",
    "\\begin{bmatrix}\n",
    "    {\\vec{h}_S}^{[1]} \\\\ \\vdots \\\\ {\\vec{h}_S}^{[L]}\n",
    "\\end{bmatrix}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notes:\n",
    "- $t\\in[1,S]$ is the timestep, i.e. the current position within the sequence of each sample.\n",
    "- $\\vec{x}_t^{[k]}$ is the input of layer $k$ at timestep $t$, respectively.\n",
    "- The outputs of the **last layer** $\\vec{y}_t^{[L]}$, are the predicted next characters for every input char.\n",
    "  These are similar to class scores in classification tasks.\n",
    "- The hidden states at the **last timestep**, $\\vec{h}_S^{[k]}$, are the final hidden state returned from the model.\n",
    "- $\\sigma(\\cdot)$ is the sigmoid function, i.e. $\\sigma(\\vec{z}) = 1/(1+e^{-\\vec{z}})$ which returns values in $(0,1)$.\n",
    "- $\\tanh(\\cdot)$ is the hyperbolic tangent, i.e. $\\tanh(\\vec{z}) = (e^{2\\vec{z}}-1)/(e^{2\\vec{z}}+1)$ which returns values in $(-1,1)$.\n",
    "- $\\vec{h_t}^{[k]}$ is the hidden state of layer $k$ at time $t$. This can be thought of as the memory of that layer.\n",
    "- $\\vec{g_t}^{[k]}$ is the candidate hidden state for time $t+1$.\n",
    "- $\\vec{z_t}^{[k]}$ is known as the update gate. It combines the previous state with the input to determine how much the current state will be combined with the new candidate state. For example, if $\\vec{z_t}^{[k]}=\\vec{1}$ then the current input has no effect on the output.\n",
    "- $\\vec{r_t}^{[k]}$ is known as the reset gate. It combines the previous state with the input to determine how much of the previous state will affect the current state candidate. For example if $\\vec{r_t}^{[k]}=\\vec{0}$ the previous state has no effect on the current candidate state."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a graphical representation of the GRU's forward pass at each timestep. The $\\vec{\\tilde{h}}$ in the image is our $\\vec{g}$ (candidate next state).\n",
    "\n",
    "<img src=\"imgs/gru_cell.png\" width=\"400\"/>\n",
    "\n",
    "You can see how the reset and update gates allow the model to completely ignore it's previous state, completely ignore it's input, or any mixture of those states (since the gates are actually continuous and between $(0,1)$)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a graphical representation of the entire model.\n",
    "You can ignore the $c_t^{[k]}$ (cell state) variables (which are relevant for LSTM models).\n",
    "Our model has only the hidden state, $h_t^{[k]}$. Also notice that we added dropout between layers (the up arrows).\n",
    "\n",
    "<img src=\"imgs/lstm_model.png\"/>\n",
    "\n",
    "The purple tensors are inputs (a sequence and initial hidden state per layer), and the green tensors are outputs (another sequence and final hidden state per layer). Each blue block implements the above forward equations.\n",
    "Blocks that are on the same vertical level are at the same layer, and therefore share parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**TODO**: Implement the `MultilayerGRU` class in the `hw3/charnn.py` module.\n",
    "\n",
    "Notes:\n",
    "- You'll need to handle input **batches** now.\n",
    "  The math is identical to the above, but all the tensors will have an extra batch\n",
    "  dimension as their first dimension.\n",
    "- Use the diagram above to help guide your implementation.\n",
    "  It will help you visualize what shapes to returns where, etc.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "in_dim = vocab_len\n",
    "h_dim = 256\n",
    "n_layers = 2\n",
    "model = charnn.MultilayerGRU(in_dim, h_dim, out_dim=in_dim, n_layers=n_layers)\n",
    "model = model.to(device)\n",
    "print(model)\n",
    "\n",
    "# Test forward pass\n",
    "y, h = model(x0.to(dtype=torch.float))\n",
    "print(f'y.shape={y.shape}')\n",
    "print(f'h.shape={h.shape}')\n",
    "\n",
    "test.assertEqual(y.shape, (batch_size, seq_len, vocab_len))\n",
    "test.assertEqual(h.shape, (batch_size, n_layers, h_dim))\n",
    "test.assertEqual(len(list(model.parameters())), 9 * n_layers + 2) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating text by sampling\n",
    "<a id=part1_6></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have a model, we can implement **text generation** based on it.\n",
    "The idea is simple:\n",
    "At each timestep our model receives one char $x_t$ from the input sequence and outputs scores $y_t$\n",
    "for what the next char should be.\n",
    "We'll convert these scores into a probability over each of the possible chars.\n",
    "In other words, for each input char $x_t$ we create a probability distribution for the next char\n",
    "conditioned on the current one and the state of the model (representing all previous inputs):\n",
    "$$p(x_{t+1}|x_t; \\vec{h}_t).$$\n",
    "\n",
    "Once we have such a distribution, we'll sample a char from it.\n",
    "This will be the first char of our generated sequence.\n",
    "Now we can feed this new char into the model, create another distribution, sample the next char and so on.\n",
    "Note that it's crucial to propagate the hidden state when sampling.\n",
    "\n",
    "The important point however is how to create the distribution from the scores.\n",
    "One way, as we saw in previous ML tasks, is to use the softmax function.\n",
    "However, a drawback of softmax is that it can generate very diffuse (more uniform) distributions if the score values are very similar. When sampling, we would prefer to control the distributions and make them less uniform to increase the chance of sampling the char(s) with the highest scores compared to the others.\n",
    "\n",
    "To control the variance of the distribution, a common trick is to add a hyperparameter $T$, known as the \n",
    "*temperature* to the softmax function. The class scores are simply scaled by $T$ before softmax is applied:\n",
    "$$\n",
    "\\mathrm{softmax}_T(\\vec{y}) = \\frac{e^{\\vec{y}/T}}{\\sum_k e^{y_k/T}}\n",
    "$$\n",
    "\n",
    "A low $T$ will result in less uniform distributions and vice-versa."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**TODO**: Implement the `hot_softmax()` function in the `hw3/charnn.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scores = y[0,0,:].detach()\n",
    "_, ax = plt.subplots(figsize=(15,5))\n",
    "\n",
    "for t in reversed([0.3, 0.5, 1.0, 100]):\n",
    "    ax.plot(charnn.hot_softmax(scores, temperature=t).cpu().numpy(), label=f'T={t}')\n",
    "ax.set_xlabel('$x_{t+1}$')\n",
    "ax.set_ylabel('$p(x_{t+1}|x_t)$')\n",
    "ax.legend()\n",
    "\n",
    "uniform_proba = 1/len(char_to_idx)\n",
    "uniform_diff = torch.abs(charnn.hot_softmax(scores, temperature=100) - uniform_proba)\n",
    "test.assertTrue(torch.all(uniform_diff < 1e-4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**TODO**: Implement the `generate_from_model()` function in the `hw3/charnn.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for _ in range(3):\n",
    "    text = charnn.generate_from_model(model, \"foobar\", 50, (char_to_idx, idx_to_char), T=0.5)\n",
    "    print(text)\n",
    "    test.assertEqual(len(text), 50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training\n",
    "<a id=part1_7></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To train such a model, we'll calculate the loss at each time step by comparing the predicted char to\n",
    "the actual char from our label. We can use cross entropy since per char it's similar to a classification problem.\n",
    "We'll then sum the losses over the sequence and back-propagate the gradients though time.\n",
    "Notice that the back-propagation algorithm will \"visit\" each layer's parameter tensors multiple times,\n",
    "so we'll accumulate gradients in parameters of the blocks. Luckily `autograd` will handle this part for us."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As usual, the first step of training will be to try and **overfit** a large model (many parameters) to a tiny dataset.\n",
    "Again, this is to ensure the model and training code are implemented correctly, i.e. that the model can learn.\n",
    "\n",
    "For a generative model such as this, overfitting is slightly trickier than for for classification.\n",
    "What we'll aim to do is to get our model to **memorize** a specific sequence of chars, so that when given the first\n",
    "char in the sequence it will immediately spit out the rest of the sequence verbatim.\n",
    "\n",
    "Let's create a tiny dataset to memorize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pick a tiny subset of the dataset\n",
    "subset_start, subset_end = 1001, 1005\n",
    "ds_corpus_ss = torch.utils.data.Subset(ds_corpus, range(subset_start, subset_end))\n",
    "dl_corpus_ss = torch.utils.data.DataLoader(ds_corpus_ss, batch_size=1, shuffle=False)\n",
    "\n",
    "# Convert subset to text\n",
    "subset_text = ''\n",
    "for i in range(subset_end - subset_start):\n",
    "    subset_text += unembed(ds_corpus_ss[i][0])\n",
    "print(f'Text to \"memorize\":\\n\\n{subset_text}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's implement the first part of our training code.\n",
    "\n",
    "**TODO**: Implement the `train_epoch()` and `train_batch()` methods of the `RNNTrainer` class in the `hw3/training.py` module. \n",
    "Note: Think about how to correctly handle the hidden state of the model between batches and epochs\n",
    "(for this specific task, i.e. text generation)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from hw3.training import RNNTrainer\n",
    "\n",
    "torch.manual_seed(42)\n",
    "\n",
    "lr = 0.01\n",
    "num_epochs = 500\n",
    "\n",
    "in_dim = vocab_len\n",
    "h_dim = 128\n",
    "n_layers = 2\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "model = charnn.MultilayerGRU(in_dim, h_dim, out_dim=in_dim, n_layers=n_layers).to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "trainer = RNNTrainer(model, loss_fn, optimizer, device)\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    epoch_result = trainer.train_epoch(dl_corpus_ss, verbose=False)\n",
    "    \n",
    "    # Every X epochs, we'll generate a sequence starting from the first char in the first sequence\n",
    "    # to visualize how/if/what the model is learning.\n",
    "    if epoch == 0 or (epoch+1) % 25 == 0:\n",
    "        avg_loss = np.mean(epoch_result.losses)\n",
    "        accuracy = np.mean(epoch_result.accuracy)\n",
    "        print(f'\\nEpoch #{epoch+1}: Avg. loss = {avg_loss:.3f}, Accuracy = {accuracy:.2f}%')\n",
    "        \n",
    "        generated_sequence = charnn.generate_from_model(model, subset_text[0],\n",
    "                                                        seq_len*(subset_end-subset_start),\n",
    "                                                        (char_to_idx,idx_to_char), T=0.1)\n",
    "        # Stop if we've successfully memorized the small dataset.\n",
    "        print(generated_sequence)\n",
    "        if generated_sequence == subset_text:\n",
    "            break\n",
    "\n",
    "# Test successful overfitting\n",
    "test.assertGreater(epoch_result.accuracy, 99)\n",
    "test.assertEqual(generated_sequence, subset_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "OK, so training works - we can memorize a short sequence. Next on the agenda is to split our full dataset into a training and test sets of batched sequences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Full dataset definition\n",
    "vocab_len = len(char_to_idx)\n",
    "seq_len = 64\n",
    "batch_size = 256\n",
    "train_test_ratio = 0.9\n",
    "num_samples = (len(corpus) - 1) // seq_len\n",
    "num_train = int(train_test_ratio * num_samples)\n",
    "\n",
    "samples, labels = charnn.chars_to_labelled_samples(corpus, char_to_idx, seq_len, device)\n",
    "\n",
    "ds_train = torch.utils.data.TensorDataset(samples[:num_train], labels[:num_train])\n",
    "dl_train = torch.utils.data.DataLoader(ds_train, batch_size=batch_size, shuffle=False, drop_last=True)\n",
    "\n",
    "ds_test = torch.utils.data.TensorDataset(samples[num_train:], labels[num_train:])\n",
    "dl_test = torch.utils.data.DataLoader(ds_test, batch_size=batch_size, shuffle=False, drop_last=True)\n",
    "\n",
    "print(f'Train: {len(dl_train):3d} batches, {len(dl_train)*batch_size*seq_len:7d} chars')\n",
    "print(f'Test:  {len(dl_test):3d} batches, {len(dl_test)*batch_size*seq_len:7d} chars')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll now train a much larger model on our large dataset.\n",
    "You'll need a **GPU** for this part.\n",
    "\n",
    "The code blocks below will train the model and save checkpoints containing the training state and the best model parameters to a file. This allows you to stop training and resume it later from where you left.\n",
    "\n",
    "Note that you can use the `main.py` script provided within the assignment folder to run this notebook from the command line as if it were a python script by using the `run-nb` subcommand. This allows you to train your model using this notebook without starting jupyter. You can combine this with `srun` or `sbatch` to run the notebook with a GPU on the course servers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Full training definition\n",
    "lr = 0.001\n",
    "num_epochs = 50\n",
    "\n",
    "in_dim = out_dim = vocab_len\n",
    "hidden_dim = 512\n",
    "n_layers = 3\n",
    "dropout = 0.5\n",
    "checkpoint_file = 'checkpoints/rnn'\n",
    "max_batches = 300\n",
    "early_stopping = 5\n",
    "\n",
    "model = charnn.MultilayerGRU(in_dim, hidden_dim, out_dim, n_layers, dropout)\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=2, verbose=True)\n",
    "trainer = RNNTrainer(model, loss_fn, optimizer, device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**TODO**:\n",
    "- Implement the `fit()` method of the `Trainer` class. You can reuse the implementation from HW2, but make sure to implement early stopping and checkpoints.\n",
    "- Implement the `test_epoch()` and `test_batch()` methods of the `RNNTrainer` class in the `hw3/training.py` module.\n",
    "- Run the following block to train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cs236605.plot import plot_fit\n",
    "\n",
    "def post_epoch_fn(epoch, test_res, train_res, verbose):\n",
    "    # Update learning rate\n",
    "    scheduler.step(test_res.accuracy)\n",
    "    # Sample from model to show progress\n",
    "    if verbose:\n",
    "        start_seq = \"ACT I.\"\n",
    "        generated_sequence = charnn.generate_from_model(\n",
    "            model, start_seq, 100, (char_to_idx,idx_to_char), T=0.5\n",
    "        )\n",
    "        print(generated_sequence)\n",
    "\n",
    "# Train, unless final checkpoint is found\n",
    "checkpoint_file_final = f'{checkpoint_file}_final.pt'\n",
    "if os.path.isfile(checkpoint_file_final):\n",
    "    print(f'*** Loading final checkpoint file {checkpoint_file_final} instead of training')\n",
    "    saved_state = torch.load(checkpoint_file_final, map_location=device)\n",
    "    model.load_state_dict(saved_state['model_state'])\n",
    "else:\n",
    "    try:\n",
    "        # Print pre-training sampling\n",
    "        print(charnn.generate_from_model(model, \"ACT I.\", 100, (char_to_idx,idx_to_char), T=0.5))\n",
    "\n",
    "        fit_res = trainer.fit(dl_train, dl_test, num_epochs, max_batches=max_batches,\n",
    "                              post_epoch_fn=post_epoch_fn, early_stopping=early_stopping,\n",
    "                              checkpoints=checkpoint_file, print_every=1)\n",
    "        \n",
    "        fig, axes = plot_fit(fit_res)\n",
    "    except KeyboardInterrupt as e:\n",
    "        print('\\n *** Training interrupted by user')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating a work of art\n",
    "<a id=part1_8></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Armed with our fully trained model, let's generate the next Hamlet! You should experiment with modifying the sampling temperature and see what happens.\n",
    "\n",
    "**TODO**: Specify the generation parameters in the `part1_generation_params()` function within the `hw3/answers.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hw3.answers\n",
    "\n",
    "start_seq, temperature = hw3.answers.part1_generation_params()\n",
    "\n",
    "generated_sequence = charnn.generate_from_model(\n",
    "    model, start_seq, 10000, (char_to_idx,idx_to_char), T=temperature\n",
    ")\n",
    "\n",
    "print(generated_sequence)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Questions\n",
    "<a id=part1_9></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**TODO** Answer the following questions. Write your answers in the appropriate variables in the module `hw3/answers.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cs236605.answers import display_answer\n",
    "import hw3.answers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question 1\n",
    "Why do we split the corpus into sequences instead of training on the whole text?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_answer(hw3.answers.part1_q1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question 2\n",
    "How is it possible that the generated text clearly shows memory longer than the sequence length?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_answer(hw3.answers.part1_q2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question 3\n",
    "Why are we not shuffling the order of batches when training?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_answer(hw3.answers.part1_q3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question 4\n",
    "1. Why do we lower the temperature for sampling (compared to the default of $1.0$ when training)?\n",
    "2. What happens when the temperature is very high and why?\n",
    "3. What happens when the temperature is very low and why?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_answer(hw3.answers.part1_q4)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}