1 Language Modeling.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yIqzcT6vyUJI"
   },
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eQrUWcSeyUJm"
   },
   "source": [
    "<center><h3>**Welcome to the Language modeling Notebook.**</h3></center>\n",
    "\n",
    "In this assignment, you are going to train a neural network to **generate news headlines**.\n",
    "To reduce computational needs, we have reduced it to headlines about technology, and a handful of Tech giants.\n",
    "In this assignment you will:\n",
    "- Learn to preprocess raw text so it can be fed into an LSTM.\n",
    "- Make use of the LSTM library of Pytorch, to train a Language model to generate headlines\n",
    "- Use your network to generate headlines, and judge which headlines are likely or not\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lCAzgVCryUJr"
   },
   "source": [
    "**What is a language model?**\n",
    "\n",
    "Language modeling is the task of assigning a probability to sentences in a language. Besides assigning a probability to each sequence of words, the language models also assigns a probability for the likelihood of a given word (or a sequence of words) to follow a sequence of words.\n",
    "— Page 105, __[Neural Network Methods in Natural Language Processing](https://www.amazon.com/Language-Processing-Synthesis-Lectures-Technologies/dp/1627052984/)__, 2017.\n",
    "\n",
    "In terms of neural network, we are training a neural network to produce probabilities (classification) over a fixed vocabulary of words.\n",
    "Concretely, we are training a neural network to produce:\n",
    "$$ P ( w_{i+1} | w_1, w_2, w_3, ..., w_i), \\forall i \\in (1,n)$$\n",
    "\n",
    "** Why is language modeling important? **\n",
    "\n",
    "Language modeling is a core problem in NLP.\n",
    "\n",
    "Language models can either be used as a stand-alone to produce new text that matches the distribution of text the model is trained on, but can also be used at the front-end of a more sophisticated model to produce better results.\n",
    "\n",
    "Recently for example, the __[BERT](https://arxiv.org/abs/1810.04805)__ paper show-cased that pretraining a large neural network on a language modeling task can help improve state-of-the-art on many NLP tasks. \n",
    "\n",
    "How good can the generation of a Language model be?\n",
    "\n",
    "If you have not seen the post about GPT-2 by OpenAI, you should read some of the samples they generated from their language model __[here](https://blog.openai.com/better-language-models/#sample1)__.\n",
    "Because of computational restrictions, we will not achieve as good text production, but the same algorithm is at the core. They just use more data and compute."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LL29FMhfyUJv"
   },
   "source": [
    "# Library imports"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5OXaa4g1yUJx"
   },
   "source": [
    "Before starting, make sure you have all these libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 3248,
     "status": "ok",
     "timestamp": 1614531798003,
     "user": {
      "displayName": "Sean Chen",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjK8rkabs7CsaHjMrOJ7lF9JD8ld_9_f9SxLG76=s64",
      "userId": "07641202443553730242"
     },
     "user_tz": 480
    },
    "id": "g81IeWKxyYO-",
    "outputId": "6e559492-7654-4cb2-9306-0c3382b8ccd1"
   },
   "outputs": [],
   "source": [
    "!pip install segtok"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the first of the following two cells if you are running the homework locally, and run the second cell if you are running the homework in Colab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DRIVE=False\n",
    "root_folder = \"\"\n",
    "dataset_folder = \"dataset/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 3213,
     "status": "ok",
     "timestamp": 1614531798008,
     "user": {
      "displayName": "Sean Chen",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjK8rkabs7CsaHjMrOJ7lF9JD8ld_9_f9SxLG76=s64",
      "userId": "07641202443553730242"
     },
     "user_tz": 480
    },
    "id": "se7xdZD7y-CL",
    "outputId": "b8e44982-3eee-4a94-aa23-be86f9ec6061"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "root_folder = \"/content/drive/My Drive/cs182_hw3/\"\n",
    "dataset_folder = \"/content/drive/My Drive/cs182_hw3_public/dataset/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "N1twBPSLyUJz"
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pPLGvVwByUJ3"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "sys.path.append(root_folder)\n",
    "from segtok import tokenizer\n",
    "from collections import Counter\n",
    "import torch as th\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "import numpy as np\n",
    "import json\n",
    "from utils import validate_to_array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OQTfdCRdyUJ8"
   },
   "source": [
    "# Loading the datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3dNGkjKmyUJ9"
   },
   "source": [
    "Make sure the dataset files are all in the `dataset` folder of the assignment.\n",
    "\n",
    " - If you are using this notebook locally: You should run the `download_data.sh` script.\n",
    " - If you are using the Colab version of the notebook, make sure that your Google Drive is mounted, and you verify from the file explorer in Colab that the files are viewable within `/content/drive/cs182_hw3_public/dataset/`\n",
    " \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4822,
     "status": "ok",
     "timestamp": 1614531799702,
     "user": {
      "displayName": "Sean Chen",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjK8rkabs7CsaHjMrOJ7lF9JD8ld_9_f9SxLG76=s64",
      "userId": "07641202443553730242"
     },
     "user_tz": 480
    },
    "id": "iLaRmt5uyUJ_",
    "outputId": "6d4b3711-3bb0-478f-b06e-cad96ad3d36d"
   },
   "outputs": [],
   "source": [
    "# This cell loads the data for the model\n",
    "# Run this before working on loading any of the additional data\n",
    "\n",
    "with open(dataset_folder+\"headline_generation_dataset_processed.json\", \"r\") as f:\n",
    "    d_released = json.load(f)\n",
    "\n",
    "with open(dataset_folder+\"headline_generation_vocabulary.txt\", \"r\",encoding='utf8') as f:\n",
    "    vocabulary = f.read().split(\"\\n\")\n",
    "w2i = {w: i for i, w in enumerate(vocabulary)} # Word to index\n",
    "i2w = {i: w for i, w in enumerate(vocabulary)} # Index to word\n",
    "unkI, padI, startI = w2i['UNK'], w2i['PAD'], w2i['<START>']\n",
    "\n",
    "vocab_size = len(vocabulary)\n",
    "input_length = len(d_released[0]['numerized']) # The length of the first element in the dataset, they are all of the same length\n",
    "d_train = [d for d in d_released if d['cut'] == 'training']\n",
    "d_valid = [d for d in d_released if d['cut'] == 'validation']\n",
    "\n",
    "print(\"Number of training samples:\",len(d_train))\n",
    "print(\"Number of validation samples:\",len(d_valid))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u31F-vXqyUKB"
   },
   "source": [
    "Now that we have loaded the data, let's inspect one of the elements. Each sample in our dataset is has a `numerized` vector, that contains the preprocessed headline. This vector is what we will feed in to the neural network. The field `numerized` corresponds to this list of tokens. The already loaded dictionary `vocabulary` maps token lists to the actual string. Use these elements to recover `title` key of entry 1001 in the training dataset.\n",
    "\n",
    "**TODO**: Write the numerized2text function in notebook_utils and inspect element 1001 in the training dataset (`entry = d_train[1001]`).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GqsHaZNbyUKD"
   },
   "outputs": [],
   "source": [
    "def numerize_sequence(tokenized):\n",
    "    return [w2i.get(w, unkI) for w in tokenized]\n",
    "def pad_sequence(numerized, pad_index, to_length):\n",
    "    pad = numerized[:to_length]\n",
    "    padded = pad + [pad_index] * (to_length - len(pad))\n",
    "    mask = [w != pad_index for w in padded]\n",
    "    return padded, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4773,
     "status": "ok",
     "timestamp": 1614531799708,
     "user": {
      "displayName": "Sean Chen",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjK8rkabs7CsaHjMrOJ7lF9JD8ld_9_f9SxLG76=s64",
      "userId": "07641202443553730242"
     },
     "user_tz": 480
    },
    "id": "C23VMn5DyUKF",
    "outputId": "821756d4-de73-4bac-d52e-6205ecaa6ffe"
   },
   "outputs": [],
   "source": [
    "def numerized2text(numerized):\n",
    "    \"\"\" Converts an integer sequence in the vocabulary into a string corresponding to the title.\n",
    "    \n",
    "        Arguments:\n",
    "            numerized: List[int]  -- The list of vocabulary indices corresponding to the string\n",
    "        Returns:\n",
    "            title: str -- The string corresponding to the numerized input, without padding.\n",
    "    \"\"\"\n",
    "    #####\n",
    "    # BEGIN YOUR CODE HERE \n",
    "    # Recover each word from the vocabulary in the list of indices in numerized, using the vocabulary variable\n",
    "    # Hint 1: Use the string.join() function to reconstruct a single string\n",
    "    # Hint 2: The objects and/or functions defined in above cells may be useful.\n",
    "    #####\n",
    "    \n",
    "\n",
    "    \n",
    "    #####\n",
    "    # END YOUR CODE HERE\n",
    "    #####\n",
    "    \n",
    "    return converted_string\n",
    "\n",
    "entry = d_train[1001]\n",
    "print(\"Reversing the numerized: \"+numerized2text(entry['numerized']))\n",
    "validate_to_array(numerized2text,(entry['numerized'],),'numerized2text',root_folder)\n",
    "print(\"From the `title` entry: \"+ entry['title'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZrpgWj1FyUKH"
   },
   "source": [
    "In language modeling, we train a model to produce the next word in the sequence given all previously generated words. This has, in practice, two steps:\n",
    "\n",
    "\n",
    "    1. Adding a special <START> token to the start of the sequence for the input. This \"shifts\" the input to the right by one. We call this the \"source\" sequence\n",
    "    2. Making the network predict the original, unshifted version (we call this the \"target\" sequence)\n",
    "\n",
    "    \n",
    "Let's take an example. Say we want to train the network on the sentence: \"The cat is great.\"\n",
    "The input to the network will be \"`<START>` The cat is great.\" The target will be: \"The cat is great\".\n",
    "    \n",
    "Therefore the first prediction is to select the word \"The\" given the `<START>` token.\n",
    "The second prediction is to produce the word \"cat\" given the two tokens \"`<START>` The\".\n",
    "At each step, the network learns to predict the next word, given all previous ones.\n",
    "    \n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HO2ww8XPyUKK"
   },
   "source": [
    "Your next step is to write the build_batch function. Given a dataset, we select a random subset of samples, and will build the \"inputs\" and the \"targets\" of the batch, following the procedure we've described.\n",
    "\n",
    "**TODO**: write the build_batch function. We give you the structure, and you have to fill in where we have left things `your_code`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Kl0f7hKKyUKL"
   },
   "outputs": [],
   "source": [
    "def build_batch(dataset, indices):\n",
    "    \"\"\" Builds a batch of source and target elements from the dataset.\n",
    "    \n",
    "        Arguments:\n",
    "            dataset: List[db_element] -- A list of dataset elements\n",
    "            indices: List[int] -- A list of indices of the dataset to sample\n",
    "        Returns:\n",
    "            batch_input: List[List[int]] -- List of source sequences\n",
    "            batch_target: List[List[int]] -- List of target sequences\n",
    "            batch_target_mask: List[List[int]] -- List of target batch masks\n",
    "    \"\"\"\n",
    "    #####\n",
    "    # BEGIN YOUR CODE HERE \n",
    "    #####\n",
    "    \n",
    "    \n",
    "    # We get a list of indices we will choose from the dataset.\n",
    "    # indices = range(iteration*batch_size,(iteration+1)*batch_size)\n",
    "    \n",
    "    # Recover what the entries for the batch are\n",
    "    batch = your_code\n",
    "    \n",
    "    # Get the raw numerized for this input, each element of the dataset has a 'numerized' key\n",
    "    batch_numerized = your_code\n",
    "    \n",
    "    # Create an array of startI that will be concatenated at position 1 for the input.\n",
    "    # Should be of shape (batch_size, 1)\n",
    "    start_tokens = your_code\n",
    "\n",
    "    # Concatenate the start_tokens with the rest of the input\n",
    "    # The np.concatenate function should be useful\n",
    "    # The output should now be [batch_size, sequence_length+1]\n",
    "    batch_input = your_code\n",
    "\n",
    "    # Remove the last word from each element in the batch\n",
    "    # To restore the [batch_size, sequence_length] size\n",
    "    batch_input = your_code\n",
    "    \n",
    "    # The target should be the un-shifted numerized input\n",
    "    batch_target = your_code\n",
    "\n",
    "    # The target-mask is a 0 or 1 filter to note which tokens are\n",
    "    # padding or not, to give the loss, so the model doesn't get rewarded for\n",
    "    # predicting PAD tokens.\n",
    "    batch_target_mask = np.array([a['mask'] for a in batch])\n",
    "    \n",
    "    #####\n",
    "    # END YOUR CODE HERE \n",
    "    #####\n",
    "        \n",
    "    return batch_input, batch_target, batch_target_mask\n",
    "validate_to_array(build_batch,(d_train, range(100)),'build_batch',root_folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fBbG3vC5yUKN"
   },
   "source": [
    "# Creating the language model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "b3Eg8rHGyUKP"
   },
   "source": [
    "Now that we've written the data pipelining, we are ready to write the Neural network.\n",
    "\n",
    "The steps to setting up a neural network to do Language modeling are:\n",
    "- Creating the placeholders for the model, where we can feed in our inputs and targets.\n",
    "- Creating an RNN of our choice, size, and with optional parameters\n",
    "- Using the RNN on our placeholder inputs.\n",
    "- Getting the output from the RNN, and projecting it into a vocabulary sized dimension, so that we can make word predictions.\n",
    "- Setting up the loss on the outputs so that the network learns to produce the correct words.\n",
    "- Finally, choosing an optimizer, and defining a training operation: using the optimizer to minimize the loss.\n",
    "\n",
    "We provide skeleton code for the model, you can fill in the `your_code` section. If you are unfamiliar with Pytorch, we provide some idea of what functions to look for, you should use the Pytorch online documentation.\n",
    "\n",
    "**TODO**: Fill in the LanguageModel in the language_model file.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rKJ2D23GyUKR"
   },
   "outputs": [],
   "source": [
    "from language_model import LanguageModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qz7ixZGKyUKT"
   },
   "source": [
    "# Training the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9sDUBkpnyUKU"
   },
   "source": [
    "Your objective is to train the Language on the dataset you are provided to reach a **validation loss <= 5.50**\n",
    "\n",
    "**TODO**: Train your model so that it achieves a validation loss of <= 5.5. \n",
    "\n",
    "**Careful**: we will be testing this loss on an unreleased test set, so make sure to evaluate properly on a validation set and not overfit. You must save the model you want us to test under: models/final_language_model (the .index, .meta and .data files)\n",
    "\n",
    "**Advice**:\n",
    "- It should be possible to attain loss <= 5.50 with a 1-layer LSTM of size 256 or less.\n",
    "- You should not need more than 10 epochs to attain the threshold. More passes over the data can however give you a better model.\n",
    "- You can however try using:\n",
    "    - LSTM dropout (Pytorch has a layer for that)\n",
    "    - Multi-layer RNN cell (Pytorch has a layer for that)\n",
    "    - Change your optimizers, tune your learning_rate, use a learning rate schedule.\n",
    "    \n",
    "**Extra credit**:\n",
    "\n",
    "Get the loss below **validation loss <= 5.00** and get 5 points of extra-credit on this assignment. Get creative,\n",
    "\n",
    "but remember, what you do should work on our held-out test set to get the points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4699,
     "status": "ok",
     "timestamp": 1614531799716,
     "user": {
      "displayName": "Sean Chen",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjK8rkabs7CsaHjMrOJ7lF9JD8ld_9_f9SxLG76=s64",
      "userId": "07641202443553730242"
     },
     "user_tz": 480
    },
    "id": "b01AZj0RyUKW",
    "outputId": "d3628056-7dbe-4bdb-eb3d-b6cb54393d2b"
   },
   "outputs": [],
   "source": [
    "# We can create our model,\n",
    "# with parameters of our choosing.\n",
    "hidden_size = 256\n",
    "num_layers = 1\n",
    "dropout = .5\n",
    "\n",
    "# Setup the loss using cross-entropy loss.\n",
    "# The logits are the output_logits we've computed,\n",
    "# look at the pytorch docs for `CrossEntropyLoss` and `permute`\n",
    "# to align the axes correctly and to account for the masking properly.\n",
    "# The targets are the goal labels we are trying to match.\n",
    "# Note that if you directly take the mean of the loss tensor,\n",
    "# it will underestimate your loss! (why would that be?)\n",
    "# Lastly, there are a few valid forms of averaging token losses,\n",
    "# here we will take the mean of all non-mask tokens together.\n",
    "criterion = nn.CrossEntropyLoss(reduction='none')\n",
    "def loss_fn(pred, target, mask):\n",
    "    pred = your_code  # put the class probabilities in the middle\n",
    "    loss_tensor = your_code\n",
    "    loss_masked = your_code\n",
    "    loss_per_sample = your_code\n",
    "    return loss_per_sample\n",
    "\n",
    "# The build_batch function outputs numpy, but our model is built in pytorch,\n",
    "# so you need to convert numpy to pytorch.\n",
    "# You also have to cast the masks into float32, target into long, and input into long.\n",
    "# Look at the `float` and `long` function.\n",
    "batch_to_torch = lambda b_in,b_target,b_mask: (your_code,\n",
    "                                               your_code, \n",
    "                                               your_code)\n",
    "\n",
    "\n",
    "# Look at the docs for torch.optim and pick an optimizer\n",
    "# And provide it with a start learning rate.\n",
    "optimizer_class = optim.Adam\n",
    "lr = 1e-4\n",
    "epochs = 20\n",
    "batch_size = 128\n",
    "\n",
    "model_id = 'test1'\n",
    "os.makedirs(root_folder+'models/part1/',exist_ok=True)\n",
    "\n",
    "device = th.device(\"cuda\" if th.cuda.is_available() else \"cpu\")\n",
    "print(device)\n",
    "list_to_device = lambda th_obj: [tensor.to(device) for tensor in th_obj]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fjIAKqbvg9ER"
   },
   "outputs": [],
   "source": [
    "model = LanguageModel(vocab_size=vocab_size, rnn_size=hidden_size, num_layers=num_layers, dropout=dropout)\n",
    "optimizer = optimizer_class(model.parameters(), lr=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 599508,
     "status": "ok",
     "timestamp": 1614532394555,
     "user": {
      "displayName": "Sean Chen",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjK8rkabs7CsaHjMrOJ7lF9JD8ld_9_f9SxLG76=s64",
      "userId": "07641202443553730242"
     },
     "user_tz": 480
    },
    "id": "sVDHLiOmyUKX",
    "outputId": "43a34a54-ae3b-4c9d-86e3-0e3c30506f02",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Skeleton code\n",
    "# You have to write your own training process to obtain a\n",
    "# Good performing model on the validation set, and save it.\n",
    "\n",
    "model.train()\n",
    "losses = []\n",
    "accuracies = []\n",
    "from tqdm import tqdm\n",
    "for epoch in range(epochs):\n",
    "    indices = np.random.permutation(range(len(d_train)))\n",
    "    t = tqdm(range(0,(len(d_train)//batch_size)+1))\n",
    "    for i in t:\n",
    "        # Here is how you obtain a batch:\n",
    "        batch = build_batch(d_train, indices[i*batch_size:(i+1)*batch_size])\n",
    "        (batch_input, batch_target, batch_target_mask) = batch_to_torch(*batch)\n",
    "        (batch_input, batch_target, batch_target_mask) = list_to_device((batch_input, batch_target, batch_target_mask))\n",
    "        model.to(device)\n",
    "        \n",
    "        prediction = your_code\n",
    "        loss = your_code\n",
    "        losses.append(loss.item())\n",
    "        accuracy = (th.eq(prediction.argmax(dim=2,keepdim=False),batch_target).float()*batch_target_mask).sum()/batch_target_mask.sum()\n",
    "        accuracies.append(accuracy.item())\n",
    "\n",
    "        your_code\n",
    "        your_code\n",
    "        your_code\n",
    "        if i % 10 == 0:\n",
    "            t.set_description(f\"Epoch: {epoch} Iteration: {i} Loss: {np.mean(losses[-10:])} Accuracy: {np.mean(accuracies[-10:])}\")\n",
    "    # save your latest model\n",
    "    save_dict = dict(\n",
    "        kwargs = dict(\n",
    "            vocab_size=vocab_size,\n",
    "            rnn_size=hidden_size,\n",
    "            num_layers=num_layers,\n",
    "            dropout=dropout\n",
    "        ),\n",
    "        model_state_dict = model.state_dict(),\n",
    "        notes = \"\",\n",
    "        optimizer_class = optimizer_class,\n",
    "        lr = lr,\n",
    "        epochs = epochs,\n",
    "        batch_size = batch_size,\n",
    "    )\n",
    "    th.save(save_dict,root_folder+f'models/part1/model_{model_id}.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fhaj6XIdyUKZ"
   },
   "source": [
    "# Using the language model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9lgsKMqmyUKa"
   },
   "source": [
    "Congratulations, you have now trained a language model! We can now use it to evaluate likely news headlines, as well as generate our very own headlines.\n",
    "\n",
    "**TODO**: Complete the three parts below, using the model you have trained."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EnR8ma3cyUKb"
   },
   "source": [
    "## (1) Evaluation loss\n",
    "\n",
    "To evaluate the language model, we evaluate its loss (ability to predict) on unseen data that is reserved for evaluation.\n",
    "Your first evaluation is to load the model you trained, and obtain a test loss. If you are running this validation and not training, run the setup cell above the training loop first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 1003,
     "status": "ok",
     "timestamp": 1614533226460,
     "user": {
      "displayName": "Sean Chen",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjK8rkabs7CsaHjMrOJ7lF9JD8ld_9_f9SxLG76=s64",
      "userId": "07641202443553730242"
     },
     "user_tz": 480
    },
    "id": "zkL164RCyUKc",
    "outputId": "65278858-abba-4c96-ccc0-26900f798626"
   },
   "outputs": [],
   "source": [
    "model_id = \"test1\"\n",
    "save_dict = th.load(root_folder+'models/part1/'+f\"model_{model_id}.pt\",map_location='cpu')\n",
    "model = LanguageModel(**save_dict['kwargs'])\n",
    "model.load_state_dict(save_dict['model_state_dict'])\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 10236,
     "status": "ok",
     "timestamp": 1614533236326,
     "user": {
      "displayName": "Sean Chen",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjK8rkabs7CsaHjMrOJ7lF9JD8ld_9_f9SxLG76=s64",
      "userId": "07641202443553730242"
     },
     "user_tz": 480
    },
    "id": "o0B71owvyUKe",
    "outputId": "cabc3e8d-833d-48f6-9366-596300db020c"
   },
   "outputs": [],
   "source": [
    "# We will evaluate your model in the best_models folder\n",
    "# In a very similar way as the code below.\n",
    "# Make sure your validation loss is below the threshold we specified\n",
    "# and that you didn't train using the validation set\n",
    "\n",
    "batch = build_batch(d_valid, range(len(d_valid)))\n",
    "(batch_input, batch_target, batch_target_mask) = batch_to_torch(*batch)\n",
    "prediction = model(batch_input.long())\n",
    "loss = loss_fn(prediction, batch_target, batch_target_mask)\n",
    "print(\"Evaluation set loss:\", loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2W-FFdHoyUKg"
   },
   "outputs": [],
   "source": [
    "# Your best performing model should go here.\n",
    "os.makedirs(root_folder+\"best_models\",exist_ok=True)\n",
    "best_model_file = root_folder+\"best_models/part1_best_model.pt\"\n",
    "th.save(save_dict,best_model_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-oP9wBwfyUKj"
   },
   "source": [
    "## (2) Evaluation of likelihood of data\n",
    "\n",
    "One use of a language model is to see what data is more likely to have originated from the training data. Because we have trained our model on news headlines, we can see which of these headlines is more likely:\n",
    "\n",
    "``Apple to release another iPhone in September``\n",
    "\n",
    "\n",
    " ``Apple and Samsung resolve all lawsuits amicably``\n",
    " \n",
    "**TODO**: Use the model to obtain the loss the neural network assigns to each sentence.\n",
    "Because the neural network assigns probability to the words appearing in a sequence, this loss can be used as a proxy to measure how likely the sentence is to have occurred in the dataset.\n",
    "Once you have the loss for each headline, write down which sentence was judged to be more likely, and explain why/if you think this is coherent.\n",
    "\n",
    "**Your answer:**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CH-YlqE2yUKk"
   },
   "outputs": [],
   "source": [
    "def raw_sample_pred(headline, model):\n",
    "    #####\n",
    "    # BEGIN YOUR CODE HERE \n",
    "    #####\n",
    "    # From the code in the Preprocessing section at the end of the notebook\n",
    "    # Find out how to tokenize the headline\n",
    "    tokenized = your_code\n",
    "\n",
    "    # Find out how to numerize the tokenized headline\n",
    "    numerized = your_code\n",
    "\n",
    "    # Learn how to pad and obtain the mask of the sequence.\n",
    "    padded, mask = your_code\n",
    "\n",
    "    # Obtain the predicted headline and target headline\n",
    "    input_headline = your_code\n",
    "    pred_headline = your_code\n",
    "    target_headline = your_code\n",
    "    mask = your_code\n",
    "\n",
    "    #####\n",
    "    # END YOUR CODE HERE \n",
    "    #####\n",
    "\n",
    "    return pred_headline,target_headline,mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 689,
     "status": "ok",
     "timestamp": 1614533244496,
     "user": {
      "displayName": "Sean Chen",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjK8rkabs7CsaHjMrOJ7lF9JD8ld_9_f9SxLG76=s64",
      "userId": "07641202443553730242"
     },
     "user_tz": 480
    },
    "id": "LexHAsuiyUKm",
    "outputId": "fdfbbbe8-7da4-4f0a-8323-52194d30988d"
   },
   "outputs": [],
   "source": [
    "model.eval()\n",
    "\n",
    "headline1 = \"Apple to release new iPhone in July\"\n",
    "headline2 = \"Apple and Samsung resolve all lawsuits\"\n",
    "\n",
    "headlines = [headline1.lower(), headline2.lower()] # Our LSTM is trained on lower-cased headlines\n",
    "for headline in headlines:\n",
    "    pred_headline,target_headline,mask = raw_sample_pred(headline, model)\n",
    "    loss = your_code # Obtain the loss\n",
    "    \n",
    "    print(\"----------------------------------------\")\n",
    "    print(\"Headline:\", headline)\n",
    "    print(\"Loss of the headline:\", loss)\n",
    "validate_to_array(raw_sample_pred,zip(headlines,[model]*2),'raw_sample_pred',root_folder,multi=True)\n",
    "# Important check: one headline should be more likely (and have lower loss)\n",
    "# Than the other headline. You should know which headline should have lower loss."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kpn8XlLXyUKn"
   },
   "source": [
    "## (3) Generation of headlines\n",
    "\n",
    "We can use our language model to generate text according to the distribution of our training data.\n",
    "The way generation works is the following:\n",
    "\n",
    "We seed the model with a beginning of sequence, and obtain the distribution for the next word.\n",
    "We select the most likely word (argmax) and add it to our sequence of words.\n",
    "Now our sequence is one word longer, and we can feed it in again as an input, for the network to produce the next sentence.\n",
    "We do this a fixed number of times (up to 20 words), and obtain automatically generated headlines!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hycdvFs6yUKo"
   },
   "source": [
    "We have provided a few headline starters that should produce interesting generated headlines.\n",
    "\n",
    "**TODO:** Get creative and find at least 2 more headline_starters that produce interesting headlines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GCYvwLw1yUKp"
   },
   "outputs": [],
   "source": [
    "def generate_sentence(headline_starter, model):\n",
    "    # Tokenize and numerize the headline. Put the numerized headline\n",
    "    # beginning in `current_build`\n",
    "    tokenized = tokenizer.word_tokenizer(headline_starter.lower())\n",
    "    current_build = [startI] + numerize_sequence(tokenized)\n",
    "\n",
    "    while len(current_build) < input_length:\n",
    "        # Pad the current_build into a input_length vector.\n",
    "        # We do this so that it can be processed by our LanguageModel class\n",
    "        current_padded, _m = pad_sequence(current_build, padI, input_length)\n",
    "\n",
    "        # Obtain the logits for the current padded sequence\n",
    "        # This involves obtaining the output_logits from our model,\n",
    "        # and not the loss like we have done so far\n",
    "        logits = your_code\n",
    "        logits_np = logits.detach().cpu().numpy()\n",
    "\n",
    "        # Obtain the row of logits that interest us, the logits for the last non-pad\n",
    "        # inputs\n",
    "        last_logits = your_code\n",
    "\n",
    "        # Find the highest scoring words in the last_logits\n",
    "        # array, or sample from the softmax.\n",
    "        # The np.argmax function may be useful for first option,\n",
    "        # sp.special.softmax and np.random.choice may be useful for second option.\n",
    "        # Append this word to our current build\n",
    "        current_build.append(your_code)\n",
    "\n",
    "    # Go from the current_build of word_indices\n",
    "    # To the headline (string) produced. This should involve\n",
    "    # the vocabulary, and a string merger.\n",
    "    produced_sentence = your_code\n",
    "    return produced_sentence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 3069,
     "status": "ok",
     "timestamp": 1614533251841,
     "user": {
      "displayName": "Sean Chen",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjK8rkabs7CsaHjMrOJ7lF9JD8ld_9_f9SxLG76=s64",
      "userId": "07641202443553730242"
     },
     "user_tz": 480
    },
    "id": "xVKeMmSAyUKq",
    "outputId": "75862bb0-787d-4aa4-9df5-4dfbc0c35e2d",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model.eval()\n",
    "# Here are some headline starters.\n",
    "# They're all about tech companies, because\n",
    "# That is what is in our dataset\n",
    "headline_starters = [\"apple has released\", \"google has released\", \"amazon\", \"tesla to\"]\n",
    "for headline_starter in headline_starters:\n",
    "    print(\"===================\")\n",
    "    print(\"Generating headline starting with: \"+headline_starter)\n",
    "\n",
    "    produced_sentence = generate_sentence(headline_starter, model)\n",
    "    print(produced_sentence)\n",
    "validate_to_array(generate_sentence,zip(headline_starters,[model]*len(headline_starters)),\"generate_sentence\",root_folder,multi=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JQ6ni5GZyUKr"
   },
   "source": [
    "## All done\n",
    "\n",
    "You are done with the first part of the HW.\n",
    "\n",
    "Next notebook deals with Summarization of text!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "id": "QoPgKURxyUKs"
   },
   "source": [
    "# Preprocessing (read only)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "FfOq00MZyUKt"
   },
   "source": [
    "**You can skip this section, however you may find these functions useful in later sections of this notebook**\n",
    "\n",
    "We have provided this code so you see how the dataset was generated. You will have to come back some of these functions later in the assignment, so feel free to read through, to get familiar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true,
    "id": "uS8FUmo9yUKt"
   },
   "outputs": [],
   "source": [
    "# You do not need to run this\n",
    "# This is to show you how the dataset was created\n",
    "# You should read to understand, so you can preprocess text\n",
    "# In the same way, in the evaluation section\n",
    "\n",
    "for a in dataset:\n",
    "    a['tokenized'] = tokenizer.word_tokenizer(a['title'].lower())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true,
    "id": "g_9okOHoyUKu",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# You do not need to run this\n",
    "# This is to show you how the dataset was created\n",
    "# You should read to understand, so you can preprocess text\n",
    "# In the same way, in the evaluation section\n",
    "\n",
    "word_counts = Counter()\n",
    "for a in dataset:\n",
    "    word_counts.update(a['tokenized'])\n",
    "\n",
    "print(word_counts.most_common(30))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true,
    "id": "OhNeO1LdyUKv"
   },
   "outputs": [],
   "source": [
    "# You do not need to run this\n",
    "# This is to show you how the dataset was created\n",
    "# You should read to understand, so you can preprocess text\n",
    "# In the same way, in the evaluation section\n",
    "\n",
    "# Creating the vocab\n",
    "vocab_size = 20000\n",
    "special_words = [\"<START>\", \"UNK\", \"PAD\"]\n",
    "vocabulary = special_words + [w for w, c in word_counts.most_common(vocab_size-len(special_words))]\n",
    "w2i = {w: i for i, w in enumerate(vocabulary)}\n",
    "\n",
    "# Numerizing and padding\n",
    "input_length = 20\n",
    "unkI, padI, startI = w2i['UNK'], w2i['PAD'], w2i['<START>']\n",
    "\n",
    "for a in dataset:\n",
    "    a['numerized'] = numerize_sequence(a['tokenized']) # Change words to IDs\n",
    "    a['numerized'], a['mask'] = pad_sequence(a['numerized'], padI, input_length) # Append appropriate PAD tokens\n",
    "    \n",
    "# Compute fraction of words that are UNK:\n",
    "word_counters = Counter([w for a in dataset for w in a['input'] if w != padI])\n",
    "\n",
    "print(\"Fraction of UNK words:\", float(word_counters[unkI]) / sum(word_counters.values()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true,
    "id": "rpeDdDrUyUKw"
   },
   "outputs": [],
   "source": [
    "# You do not need to run this\n",
    "# This is to show you how the dataset was created\n",
    "# You should read to understand, so you can preprocess text\n",
    "# In the same way, in the evaluation section\n",
    "\n",
    "d_released_processed   = [d for d in dataset if d['cut'] != 'testing']\n",
    "d_unreleased_processed = [d for d in dataset if d['cut'] == 'testing']\n",
    "\n",
    "with open(\"dataset/headline_generation_dataset_processed.json\", \"w\") as f:\n",
    "    json.dump(d_released_processed, f)\n",
    "\n",
    "# This file is purposefully left out of the assignment, we will use it to evaluate your model.\n",
    "with open(\"dataset/headline_generation_dataset_unreleased_processed.json\", \"w\") as f:\n",
    "    json.dump(d_unreleased_processed, f)\n",
    "    \n",
    "with open(\"dataset/headline_generation_vocabulary.txt\", \"w\") as f:\n",
    "    f.write(\"\\n\".join(vocabulary).encode('utf8'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lS--8I0KlDvO"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "anaconda-cloud": {},
  "colab": {
   "collapsed_sections": [],
   "name": "1 Language Modeling.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}