2_4_LSTMs/2. LSTM Training, Part of Speech Tagging.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LSTM for Part-of-Speech Tagging\n",
    "\n",
    "In this section, we will use an LSTM to predict part-of-speech tags for words. What exactly is part-of-speech tagging?\n",
    "\n",
    "Part of speech tagging is the process of determining the *category* of a word from the words in its surrounding context. You can think of part of speech tagging as a way to go from words to their [Mad Libs](https://en.wikipedia.org/wiki/Mad_Libs#Format) categories. Mad Libs are incomplete short stories that have many words replaced by blanks. Each blank has a specified word-category, such as `\"noun\"`, `\"verb\"`, `\"adjective\"`, and so on. One player asks another to fill in these blanks (prompted only by the word-category) until they have created a complete, silly story of their own. Here is an example of such categories:\n",
    "\n",
    "```text\n",
    "Today, you'll be learning how to [verb]. It may be a [adjective] process, but I think it will be rewarding! \n",
    "If you want to take a break you should [verb] and treat yourself to some [plural noun].\n",
    "```\n",
    "... and a set of possible words that fall into those categories:\n",
    "```text\n",
    "Today, you'll be learning how to code. It may be a challenging process, but I think it will be rewarding! \n",
    "If you want to take a break you should stretch and treat yourself to some puppies.\n",
    "```\n",
    "\n",
    "\n",
    "### Why Tag Speech?\n",
    "\n",
    "Tagging parts of speech is often used to help disambiguate natural language phrases because it can be done quickly and with high accuracy. It can help answer: what subject is someone talking about? Tagging can be used for many NLP tasks like creating new sentences using a sequence of tags that make sense together, filling in a Mad Libs style game, and determining correct pronunciation during speech synthesis. It is also used in information retrieval, and for word disambiguation (ex. determining when someone says *right* like the direction versus *right* like \"that's right!\").\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparing the Data\n",
    "\n",
    "Now, we know that neural networks do not do well with words as input and so our first step will be to prepare our training data and map each word to a numerical value. \n",
    "\n",
    "We start by creating a small set of training data, you can see that this is a few simple sentences broken down into a list of words and their corresponding word-tags. Note that the sentences are turned into lowercase words using `lower()` and then split into separate words using `split()`, which splits the sentence by whitespace characters.\n",
    "\n",
    "#### Words to indices\n",
    "\n",
    "Then, from this training data, we create a dictionary that maps each unique word in our vocabulary to a numerical value; a unique index `idx`. We do the same for each word-tag, for example: a noun will be represented by the number `1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import resources\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# training sentences and their corresponding word-tags\n",
    "training_data = [\n",
    "    (\"The cat ate the cheese\".lower().split(), [\"DET\", \"NN\", \"V\", \"DET\", \"NN\"]),\n",
    "    (\"She read that book\".lower().split(), [\"NN\", \"V\", \"DET\", \"NN\"]),\n",
    "    (\"The dog loves art\".lower().split(), [\"DET\", \"NN\", \"V\", \"NN\"]),\n",
    "    (\"The elephant answers the phone\".lower().split(), [\"DET\", \"NN\", \"V\", \"DET\", \"NN\"])\n",
    "]\n",
    "\n",
    "# create a dictionary that maps words to indices\n",
    "word2idx = {}\n",
    "for sent, tags in training_data:\n",
    "    for word in sent:\n",
    "        if word not in word2idx:\n",
    "            word2idx[word] = len(word2idx)\n",
    "\n",
    "# create a dictionary that maps tags to indices\n",
    "tag2idx = {\"DET\": 0, \"NN\": 1, \"V\": 2}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, print out the created dictionary to see the words and their numerical values! \n",
    "\n",
    "You should see every word in our training set and its index value. Note that the word \"the\" only appears once because our vocabulary only includes *unique* words."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'the': 0, 'cat': 1, 'ate': 2, 'cheese': 3, 'she': 4, 'read': 5, 'that': 6, 'book': 7, 'dog': 8, 'loves': 9, 'art': 10, 'elephant': 11, 'answers': 12, 'phone': 13}\n"
     ]
    }
   ],
   "source": [
    "# print out the created dictionary\n",
    "print(word2idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# a helper function for converting a sequence of words to a Tensor of numerical values\n",
    "# will be used later in training\n",
    "def prepare_sequence(seq, to_idx):\n",
    "    '''This function takes in a sequence of words and returns a \n",
    "    corresponding Tensor of numerical values (indices for each word).'''\n",
    "    idxs = [to_idx[w] for w in seq]\n",
    "    idxs = np.array(idxs)\n",
    "    return torch.from_numpy(idxs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0,  8, 12,  0, 13])\n"
     ]
    }
   ],
   "source": [
    "# check out what prepare_sequence does for one of our training sentences:\n",
    "example_input = prepare_sequence(\"The dog answers the phone\".lower().split(), word2idx)\n",
    "print(example_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Creating the Model\n",
    "\n",
    "Our model will assume a few things:\n",
    "1. Our input is broken down into a sequence of words, so a sentence will be [w1, w2, ...]\n",
    "2. These words come from a larger list of words that we already know (a vocabulary)\n",
    "3. We have a limited set of tags, `[NN, V, DET]`, which mean: a noun, a verb, and a determinant (words like \"the\" or \"that\"), respectively\n",
    "4. We want to predict\\* a tag for each input word\n",
    "\n",
    "\\* To do the prediction, we will pass an LSTM over a test sentence and apply a softmax function to the hidden state of the LSTM; the result is a vector of tag scores from which we can get the predicted tag for a word based on the *maximum* value in this distribution of tag scores. \n",
    "\n",
    "Mathematically, we can represent any tag prediction $\\hat{y}_i$ as: \n",
    "\n",
    "\\begin{align}\\hat{y}_i = \\text{argmax}_j \\  (\\log \\text{Softmax}(Ah_i + b))_j\\end{align}\n",
    "\n",
    "Where $A$ is a learned weight and $b$, a learned bias term, and the hidden state at timestep $i$ is $h_i$. \n",
    "\n",
    "\n",
    "### Word embeddings\n",
    "\n",
    "We know that an LSTM takes in an expected input size and hidden_dim, but sentences are rarely of a consistent size, so how can we define the input of our LSTM?\n",
    "\n",
    "Well, at the very start of this net, we'll create an `Embedding` layer that takes in the size of our vocabulary and returns a vector of a specified size, `embedding_dim`, for each word in an input sequence of words. It's important that this be the first layer in this net. You can read more about this embedding layer in [the PyTorch documentation](https://pytorch.org/tutorials/beginner/nlp/word_embeddings_tutorial.html#word-embeddings-in-pytorch).\n",
    "\n",
    "Pictured below is the expected architecture for this tagger model.\n",
    "\n",
    "<img src='images/speech_tagger.png' height=60% width=60% >\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTMTagger(nn.Module):\n",
    "\n",
    "    def __init__(self, embedding_dim, hidden_dim, vocab_size, tagset_size):\n",
    "        ''' Initialize the layers of this model.'''\n",
    "        super(LSTMTagger, self).__init__()\n",
    "        \n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        # embedding layer that turns words into a vector of a specified size\n",
    "        self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)\n",
    "\n",
    "        # the LSTM takes embedded word vectors (of a specified size) as inputs \n",
    "        # and outputs hidden states of size hidden_dim\n",
    "        self.lstm = nn.LSTM(embedding_dim, hidden_dim)\n",
    "\n",
    "        # the linear layer that maps the hidden state output dimension \n",
    "        # to the number of tags we want as output, tagset_size (in this case this is 3 tags)\n",
    "        self.hidden2tag = nn.Linear(hidden_dim, tagset_size)\n",
    "        \n",
    "        # initialize the hidden state (see code below)\n",
    "        self.hidden = self.init_hidden()\n",
    "\n",
    "        \n",
    "    def init_hidden(self):\n",
    "        ''' At the start of training, we need to initialize a hidden state;\n",
    "           there will be none because the hidden state is formed based on perviously seen data.\n",
    "           So, this function defines a hidden state with all zeroes and of a specified size.'''\n",
    "        # The axes dimensions are (n_layers, batch_size, hidden_dim)\n",
    "        return (torch.zeros(1, 1, self.hidden_dim),\n",
    "                torch.zeros(1, 1, self.hidden_dim))\n",
    "\n",
    "    def forward(self, sentence):\n",
    "        ''' Define the feedforward behavior of the model.'''\n",
    "        # create embedded word vectors for each word in a sentence\n",
    "        embeds = self.word_embeddings(sentence)\n",
    "        \n",
    "        # get the output and hidden state by passing the lstm over our word embeddings\n",
    "        # the lstm takes in our embeddings and hiddent state\n",
    "        lstm_out, self.hidden = self.lstm(\n",
    "            embeds.view(len(sentence), 1, -1), self.hidden)\n",
    "        \n",
    "        # get the scores for the most likely tag for a word\n",
    "        tag_outputs = self.hidden2tag(lstm_out.view(len(sentence), -1))\n",
    "        tag_scores = F.log_softmax(tag_outputs, dim=1)\n",
    "        \n",
    "        return tag_scores\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define how the model trains\n",
    "\n",
    "To train the model, we have to instantiate it and define the loss and optimizers that we want to use.\n",
    "\n",
    "First, we define the size of our word embeddings. The `EMBEDDING_DIM` defines the size of our word vectors for our simple vocabulary and training set; we will keep them small so we can see how the weights change as we train.\n",
    "\n",
    "**Note: the embedding dimension for a complex dataset will usually be much larger, around 64, 128, or 256 dimensional.**\n",
    "\n",
    "\n",
    "#### Loss and Optimization\n",
    "\n",
    "Since our LSTM outputs a series of tag scores with a softmax layer, we will use `NLLLoss`. In tandem with a softmax layer, NLL Loss creates the kind of cross entropy loss that we typically use for analyzing a distribution of class scores. We'll use standard gradient descent optimization, but you are encouraged to play around with other optimizers!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the embedding dimension defines the size of our word vectors\n",
    "# for our simple vocabulary and training set, we will keep these small\n",
    "EMBEDDING_DIM = 6\n",
    "HIDDEN_DIM = 6\n",
    "\n",
    "# instantiate our model\n",
    "model = LSTMTagger(EMBEDDING_DIM, HIDDEN_DIM, len(word2idx), len(tag2idx))\n",
    "\n",
    "# define our loss and optimizer\n",
    "loss_function = nn.NLLLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just to check that our model has learned something, let's first look at the scores for a sample test sentence *before* our model is trained. Note that the test sentence *must* be made of words from our vocabulary otherwise its words cannot be turned into indices.\n",
    "\n",
    "The scores should be Tensors of length 3 (for each of our tags) and there should be scores for each word in the input sentence.\n",
    "\n",
    "For the test sentence, \"The cheese loves the elephant\", we know that this has the tags (DET, NN, V, DET, NN) or `[0, 1, 2, 0, 1]`, but our network does not yet know this. In fact, in this case, our model starts out with a hidden state of all zeroes and so all the scores and the predicted tags should be low, random, and about what you'd expect for a network that is not yet trained!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.9905, -1.2271, -1.0923],\n",
      "        [-0.9351, -1.2513, -1.1354],\n",
      "        [-0.8954, -1.2777, -1.1620],\n",
      "        [-0.9653, -1.2382, -1.1111],\n",
      "        [-0.9125, -1.2902, -1.1293]], grad_fn=<LogSoftmaxBackward>)\n",
      "\n",
      "\n",
      "Predicted tags: \n",
      " tensor([0, 0, 0, 0, 0])\n"
     ]
    }
   ],
   "source": [
    "test_sentence = \"The cheese loves the elephant\".lower().split()\n",
    "\n",
    "# see what the scores are before training\n",
    "# element [i,j] of the output is the *score* for tag j for word i.\n",
    "# to check the initial accuracy of our model, we don't need to train, so we use model.eval()\n",
    "inputs = prepare_sequence(test_sentence, word2idx)\n",
    "inputs = inputs\n",
    "tag_scores = model(inputs)\n",
    "print(tag_scores)\n",
    "\n",
    "# tag_scores outputs a vector of tag scores for each word in an inpit sentence\n",
    "# to get the most likely tag index, we grab the index with the maximum score!\n",
    "# recall that these numbers correspond to tag2idx = {\"DET\": 0, \"NN\": 1, \"V\": 2}\n",
    "_, predicted_tags = torch.max(tag_scores, 1)\n",
    "print('\\n')\n",
    "print('Predicted tags: \\n',predicted_tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Train the Model\n",
    "\n",
    "Loop through all our training data for multiple epochs (again we are using a small epoch value for this simple training data). This loop:\n",
    "\n",
    "1. Prepares our model for training by zero-ing the gradients\n",
    "2. Initializes the hidden state of our LSTM\n",
    "3. Prepares our data for training\n",
    "4. Runs a forward pass on our inputs to get tag_scores\n",
    "5. Calculates the loss between tag_scores and the true tag\n",
    "6. Updates the weights of our model using backpropagation\n",
    "\n",
    "In this example, we are printing out the average epoch loss, every 20 epochs; you should see it decrease over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 20, loss: 1.04330\n",
      "Epoch: 40, loss: 0.93115\n",
      "Epoch: 60, loss: 0.72896\n",
      "Epoch: 80, loss: 0.59111\n",
      "Epoch: 100, loss: 0.46961\n",
      "Epoch: 120, loss: 0.35279\n",
      "Epoch: 140, loss: 0.25012\n",
      "Epoch: 160, loss: 0.17300\n",
      "Epoch: 180, loss: 0.12352\n",
      "Epoch: 200, loss: 0.09230\n",
      "Epoch: 220, loss: 0.07173\n",
      "Epoch: 240, loss: 0.05752\n",
      "Epoch: 260, loss: 0.04733\n",
      "Epoch: 280, loss: 0.03976\n",
      "Epoch: 300, loss: 0.03398\n"
     ]
    }
   ],
   "source": [
    "# normally these epochs take a lot longer \n",
    "# but with our toy data (only 3 sentences), we can do many epochs in a short time\n",
    "n_epochs = 300\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    \n",
    "    epoch_loss = 0.0\n",
    "    \n",
    "    # get all sentences and corresponding tags in the training data\n",
    "    for sentence, tags in training_data:\n",
    "        \n",
    "        # zero the gradients\n",
    "        model.zero_grad()\n",
    "\n",
    "        # zero the hidden state of the LSTM, this detaches it from its history\n",
    "        model.hidden = model.init_hidden()\n",
    "\n",
    "        # prepare the inputs for processing by out network, \n",
    "        # turn all sentences and targets into Tensors of numerical indices\n",
    "        sentence_in = prepare_sequence(sentence, word2idx)\n",
    "        targets = prepare_sequence(tags, tag2idx)\n",
    "\n",
    "        # forward pass to get tag scores\n",
    "        tag_scores = model(sentence_in)\n",
    "\n",
    "        # compute the loss, and gradients \n",
    "        loss = loss_function(tag_scores, targets)\n",
    "        epoch_loss += loss.item()\n",
    "        loss.backward()\n",
    "        \n",
    "        # update the model parameters with optimizer.step()\n",
    "        optimizer.step()\n",
    "        \n",
    "    # print out avg loss per 20 epochs\n",
    "    if(epoch%20 == 19):\n",
    "        print(\"Epoch: %d, loss: %1.5f\" % (epoch+1, epoch_loss/len(training_data)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing\n",
    "\n",
    "See how your model performs *after* training. Compare this output with the scores from before training, above.\n",
    "\n",
    "Again, for the test sentence, \"The cheese loves the elephant\", we know that this has the tags (DET, NN, V, DET, NN) or `[0, 1, 2, 0, 1]`. Let's see if our model has learned to find these tags!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-1.9509e-02, -3.9655e+00, -7.9270e+00],\n",
      "        [-5.6050e+00, -7.7664e-03, -5.5073e+00],\n",
      "        [-5.8641e+00, -1.7819e+00, -1.8773e-01],\n",
      "        [-1.5017e-02, -4.2895e+00, -6.7309e+00],\n",
      "        [-6.0840e+00, -1.0051e-02, -4.8637e+00]], grad_fn=<LogSoftmaxBackward>)\n",
      "\n",
      "\n",
      "Predicted tags: \n",
      " tensor([0, 1, 2, 0, 1])\n"
     ]
    }
   ],
   "source": [
    "test_sentence = \"The cheese loves the elephant\".lower().split()\n",
    "\n",
    "# see what the scores are after training\n",
    "inputs = prepare_sequence(test_sentence, word2idx)\n",
    "inputs = inputs\n",
    "tag_scores = model(inputs)\n",
    "print(tag_scores)\n",
    "\n",
    "# print the most likely tag index, by grabbing the index with the maximum score!\n",
    "# recall that these numbers correspond to tag2idx = {\"DET\": 0, \"NN\": 1, \"V\": 2}\n",
    "_, predicted_tags = torch.max(tag_scores, 1)\n",
    "print('\\n')\n",
    "print('Predicted tags: \\n',predicted_tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Great job!\n",
    "\n",
    "To improve this model, see if you can add sentences to this model and create a more robust speech tagger. Try to initialize the hidden state in a different way or play around with the optimizers and see if you can decrease model loss even faster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}