RNN_language_models.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOm0hkd7i1hAAKqA5Ugxwpj",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/pmadhyastha/INM434/blob/main/RNN_language_models.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Neural sequence models"
      ],
      "metadata": {
        "id": "utEgUjj9oeM7"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C8yFO1GSnJJH"
      },
      "outputs": [],
      "source": [
        "__author__ = \"Pranava Madhyastha\" \n",
        "__version__ = \"INM434/IN3045 City, University of London, Spring 2023\""
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Simple character level RNN language model from scratch. \n",
        "\n",
        "We will begin by examining the internals of a character based RNN language model. The code below is implementing a minimal character-level Vanilla RNN model. It is a simple neural network architecture that is trained to predict the next character in a sequence of characters.\n",
        "\n",
        "\n",
        "Let's begin with some preliminaries, we will first try with a very simple sentence: `this is the NLP lab for this NLP module`. \n",
        "\n",
        "We will first compute some statistics on the given text data \n",
        "\n",
        "### TODO: The text data can either be supplied as a string or read from a file - supply a toy text file and try running the code.\n",
        "\n",
        "The code below only uses the numpy module. \n",
        "\n",
        "The code converts the text data to a list of characters, removes any duplicate characters, and assigns the resulting list to the variable chars. The number of unique characters in the text data is then computed using the len() function and assigned to the variable num_chars. Finally, the size of the text data (i.e., the total number of characters in the text data) is computed and assigned to the variable txt_data_size.\n",
        "\n"
      ],
      "metadata": {
        "id": "krKF02Evo5Fe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "\n",
        "\n",
        "# load text data\n",
        "txt_data = \"this is the NLP lab for this NLP module \" \n",
        "\n",
        "# or open some file with text based data. \n",
        "# txt_data = open('some_file.txt', 'r').read() \n",
        "\n",
        "chars = list(set(txt_data)) \n",
        "\n",
        "num_chars = len(chars) \n",
        "txt_data_size = len(txt_data)\n",
        "\n",
        "print(\"unique characters : \", num_chars) \n",
        "print(\"txt_data_size : \", txt_data_size)"
      ],
      "metadata": {
        "id": "S9XJJnymoqHS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The char_to_int and int_to_char dictionaries are created using dictionary comprehensions. These dictionaries map each character in chars to a unique integer, and vice versa.\n",
        "The input data is encoded by replacing each character in txt_data with its corresponding integer using a list comprehension. The resulting integer_encoded list contains the encoded data."
      ],
      "metadata": {
        "id": "QNEOkg9brAmm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# One-hot encode\n",
        "char_to_int = {c: i for i, c in enumerate(chars)}\n",
        "int_to_char = {i: c for i, c in enumerate(chars)}\n",
        "\n",
        "print(\"Character to integer mapping:\", char_to_int)\n",
        "print(\"----------------------------------------------------\")\n",
        "print(\"Integer to character mapping:\", int_to_char)\n",
        "print(\"----------------------------------------------------\")\n",
        "\n",
        "# Integer encode input data\n",
        "integer_encoded = [char_to_int[c] for c in txt_data]\n",
        "print(\"Integer encoded input data:\", integer_encoded)\n",
        "print(\"----------------------------------------------------\")\n",
        "print(\"Data length:\", len(integer_encoded))"
      ],
      "metadata": {
        "id": "nka54hkZo1vf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will now set the hyperparameters of a neural network that will be trained on some text data, and initializes the model parameters (i.e., weights and biases) with random values.\n"
      ],
      "metadata": {
        "id": "zd8bot4pr9l7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define hyperparameters\n",
        "iteration = 5000\n",
        "sequence_length = 10\n",
        "batch_size = round((txt_data_size / sequence_length) + 0.5) # Divide the data into batches\n",
        "hidden_size = 100 # Number of neurons in the hidden layer\n",
        "learning_rate = 0.1 # Learning rate for optimization algorithm\n",
        "\n",
        "# Initialize model parameters\n",
        "W_xh = np.random.randn(hidden_size, num_chars) * 0.01 # Weight matrix for input to hidden layer\n",
        "W_hh = np.random.randn(hidden_size, hidden_size) * 0.01 # Weight matrix for hidden to hidden layer\n",
        "W_hy = np.random.randn(num_chars, hidden_size) * 0.01 # Weight matrix for hidden to output layer\n",
        "\n",
        "b_h = np.zeros((hidden_size, 1)) # Bias vector for hidden layer\n",
        "b_y = np.zeros((num_chars, 1)) # Bias vector for output layer\n",
        "\n",
        "h_prev = np.zeros((hidden_size, 1)) # Previous hidden state, initialized as all zeros\n"
      ],
      "metadata": {
        "id": "AAe5fe0Gre4V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will now define the forward loop for the RNN. This function implements forward propagation through an RNN. It takes in a sequence of inputs, a sequence of targets, and a previous hidden state, and outputs the loss, the probabilities for the next characters, the hidden states, and the one-hot encoded input vectors.\n",
        "\n",
        "The function initializes empty dictionaries to store the input vectors, hidden states, unnormalized log probabilities, and probabilities for each time step. It then initializes the hidden state at time step -1 to the previous hidden state, and initializes the loss to 0.\n",
        "\n",
        "For each time step, the function one-hot encodes the current input character, computes the hidden state using the previous hidden state and the current input character, computes the unnormalized log probabilities for the next characters using the hidden state and the output weights, and computes the probabilities for the next characters using softmax.\n",
        "\n",
        "The function also computes the loss using the cross-entropy loss formula, which measures the difference between the predicted probabilities and the actual targets.\n",
        "\n",
        "Finally, the function returns the loss, the probabilities for the next characters, the hidden states, and the one-hot encoded input vectors."
      ],
      "metadata": {
        "id": "PyQoGbMnsXG-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def forwardprop(inputs, targets, h_prev):\n",
        "        \n",
        "    # initialize variables\n",
        "    xs, hs, ys, ps = {}, {}, {}, {} # create empty dictionaries to store values\n",
        "    hs[-1] = np.copy(h_prev) # copy previous hidden state vector to -1 key value\n",
        "    loss = 0 # initialize loss variable\n",
        "    \n",
        "    # loop through the sequence\n",
        "    for t in range(len(inputs)): # t is a \"time step\" and is used as a key in the dictionaries\n",
        "        \n",
        "        xs[t] = np.zeros((num_chars,1)) # initialize input vector\n",
        "        xs[t][inputs[t]] = 1 # set the index of the current character to 1, one-hot encoding\n",
        "        hs[t] = np.tanh(np.dot(W_xh, xs[t]) + np.dot(W_hh, hs[t-1]) + b_h) # compute hidden state\n",
        "        ys[t] = np.dot(W_hy, hs[t]) + b_y # compute unnormalized log probabilities for next characters\n",
        "        ps[t] = np.exp(ys[t]) / np.sum(np.exp(ys[t])) # compute probabilities for next characters using softmax\n",
        "        \n",
        "        loss += -np.log(ps[t][targets[t],0]) # compute loss using cross-entropy loss formula\n",
        "\n",
        "    return loss, ps, hs, xs\n"
      ],
      "metadata": {
        "id": "44_TFCC2sDnR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will now define the backprop function: the purpose of backpropagation is to compute the gradients of the loss function with respect to the parameters of the RNN, which can then be used to update those parameters via gradient descent."
      ],
      "metadata": {
        "id": "VBHz9IrptH3i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "def backprop(ps, inputs, hs, xs):\n",
        "\n",
        "    dWxh, dWhh, dWhy = np.zeros_like(W_xh), np.zeros_like(W_hh), np.zeros_like(W_hy) # make all zero matrices.\n",
        "    dbh, dby = np.zeros_like(b_h), np.zeros_like(b_y)\n",
        "    dhnext = np.zeros_like(hs[0]) # (hidden_size,1) \n",
        "\n",
        "    # reversed\n",
        "    for t in reversed(range(len(inputs))):\n",
        "        dy = np.copy(ps[t]) # shape (num_chars,1).  \"dy\" means \"dloss/dy\"\n",
        "        dy[targets[t]] -= 1 # backprop into y. After taking the soft max in the input vector, subtract 1 from the value of the element corresponding to the correct label.\n",
        "        dWhy += np.dot(dy, hs[t].T)\n",
        "        dby += dy \n",
        "        dh = np.dot(W_hy.T, dy) + dhnext # backprop into h. \n",
        "        dhraw = (1 - hs[t] * hs[t]) * dh # backprop through tanh nonlinearity #tanh'(x) = 1-tanh^2(x)\n",
        "        dbh += dhraw\n",
        "        dWxh += np.dot(dhraw, xs[t].T)\n",
        "        dWhh += np.dot(dhraw, hs[t-1].T)\n",
        "        dhnext = np.dot(W_hh.T, dhraw)\n",
        "    for dparam in [dWxh, dWhh, dWhy, dbh, dby]: \n",
        "        np.clip(dparam, -5, 5, out=dparam) # clip to mitigate exploding gradients.  \n",
        "    \n",
        "    return dWxh, dWhh, dWhy, dbh, dby\n",
        "\n"
      ],
      "metadata": {
        "id": "TEg6Bbtusn8S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let us now train the RNN: In the code below: the outer loop of the code iterates over a given number of training iterations, and within each iteration, the training data is divided into batches of a specified size, and the model is trained on each batch. The Adagrad optimisation algorithm is used to update the model parameters after each batch. Here's a bit more about the adagrad algorithm: https://optimization.cbe.cornell.edu/index.php?title=AdaGrad. "
      ],
      "metadata": {
        "id": "Im5-r2hJuVW4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "data_pointer = 0\n",
        "\n",
        "# memory variables for optimiser\n",
        "mWxh, mWhh, mWhy = np.zeros_like(W_xh), np.zeros_like(W_hh), np.zeros_like(W_hy)\n",
        "mbh, mby = np.zeros_like(b_h), np.zeros_like(b_y) \n",
        "\n",
        "\n",
        "for iteration_index in range(iteration):\n",
        "    h_prev = np.zeros((hidden_size,1)) # reset RNN memory\n",
        "    data_pointer = 0 # go from start of data\n",
        "    \n",
        "    for batch_index in range(batch_size):\n",
        "        \n",
        "        inputs = [char_to_int[ch] for ch in txt_data[data_pointer:data_pointer+sequence_length]]\n",
        "        targets = [char_to_int[ch] for ch in txt_data[data_pointer+1:data_pointer+sequence_length+1]] # t+1        \n",
        "            \n",
        "        if (data_pointer+sequence_length+1 >= len(txt_data) and batch_index == batch_size-1): # processing of the last part of the input data. \n",
        "            targets.append(char_to_int[\" \"])   # When the data doesn't fit, add space(\" \") to the back.\n",
        "\n",
        "\n",
        "        # forward\n",
        "        loss, ps, hs, xs = forwardprop(inputs, targets, h_prev)\n",
        "    \n",
        "        # backward\n",
        "        dWxh, dWhh, dWhy, dbh, dby = backprop(ps, inputs, hs, xs) \n",
        "        \n",
        "        \n",
        "        # perform parameter update with optimiser (adagrad)\n",
        "        for param, dparam, mem in zip([W_xh, W_hh, W_hy, b_h, b_y], \n",
        "                                      [dWxh, dWhh, dWhy, dbh, dby], \n",
        "                                      [mWxh, mWhh, mWhy, mbh, mby]):\n",
        "            mem += dparam * dparam # elementwise\n",
        "            param += -learning_rate * dparam / np.sqrt(mem + 1e-8)      \n",
        "    \n",
        "        data_pointer += sequence_length # move data pointer\n",
        "        \n",
        "    if iteration_index % 100 == 0:\n",
        "        print(f'iteration {iteration_index}, loss: {loss}') # print progress\n"
      ],
      "metadata": {
        "id": "3L8tNWP1tgWn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now let us generate some samples from our model. We will use a predict function for doing this: the function generates new text by sampling from the trained RNN model. The function takes two arguments: test_char, which is the starting character for generating new text, and length, which is the length of the generated text.\n",
        "\n",
        "`predict` initializes an empty array x and sets the value at the index corresponding to test_char to 1, which represents the one-hot encoding of the input character. It also initializes an empty list ixes to store the indices of the predicted characters.\n",
        "\n",
        "The function then uses a for loop to generate length number of characters. In each iteration, it computes the hidden state h and the output y of the RNN using the current input x and the previous hidden state h. It then calculates the probability distribution p over all the possible characters using the softmax function applied to the output y. The function then samples the next character index ix randomly from the probability distribution p, and sets the value at the corresponding index of x to 1. It appends ix to the ixes list, which stores the index of the predicted character.\n",
        "\n",
        "Finally, the function converts the list of indices ixes to the corresponding characters using the int_to_char dictionary, concatenates them into a string txt, and prints the generated text.\n"
      ],
      "metadata": {
        "id": "Fzrn4vASvLvU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def predict(test_char, length):\n",
        "    x = np.zeros((num_chars, 1)) \n",
        "    x[char_to_int[test_char]] = 1\n",
        "    ixes = []\n",
        "    h = np.zeros((hidden_size, 1))\n",
        "\n",
        "    for t in range(length):\n",
        "        h = np.tanh(np.dot(W_xh, x) + np.dot(W_hh, h) + b_h) \n",
        "        y = np.dot(W_hy, h) + b_y\n",
        "        p = np.exp(y) / np.sum(np.exp(y)) \n",
        "        ix = np.random.choice(range(num_chars), p=p.ravel())  # ravel -> rank0\n",
        "        ixes.append(ix)  # list\n",
        "        x = np.zeros((num_chars, 1))  # init\n",
        "        x[ix] = 1 \n",
        "\n",
        "    txt = ''.join(int_to_char[i] for i in ixes)\n",
        "    print('----\\n%s\\n----' % txt)\n"
      ],
      "metadata": {
        "id": "LHNVAeqltwDK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let us sample now: "
      ],
      "metadata": {
        "id": "X4Gnakc1v1DW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "predict('t',10) # (the input characted, length of output)"
      ],
      "metadata": {
        "id": "XmDBm2ZvvTDn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### TODO: \n",
        "- Try a bigger length of output. \n",
        "- Try with a large input file (download the file locally using `!wget <url> -O some_file.txt` (use the exclamation mark too). \n",
        " "
      ],
      "metadata": {
        "id": "FU0VkPVAwFqV"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Neural language model using LSTM with pytorch \n",
        "\n",
        "We have to use GPU for this part of the code - so change the runtime to support the GPU backend. \n",
        "\n",
        "The code below implements a simple LSTM based language model. The RNN class is a subclass of the nn.Module class, which is a base class for all neural network modules in PyTorch. The init method defines the structure of the RNN model by initializing the input embedding layer with nn.Embedding, the LSTM layer with nn.LSTM, and the output layer with nn.Linear.\n",
        "\n",
        "The forward method takes an input sequence and a hidden state as input, applies the embedding layer to convert input sequence to a sequence of embeddings, passes the embeddings sequence through the LSTM layer to get the hidden state and output, and applies the linear layer to the output to get the final output. Finally, it detaches the hidden state to avoid backpropagating through time, and returns both the output and the updated hidden state."
      ],
      "metadata": {
        "id": "r3Bwx7PVxCzF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.distributions import Categorical\n",
        "import numpy as np\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "class RNN(nn.Module):\n",
        "    def __init__(self, input_size, output_size, hidden_size, num_layers):\n",
        "        super(RNN, self).__init__()\n",
        "        self.embedding = nn.Embedding(input_size, input_size)\n",
        "        self.rnn = nn.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers)\n",
        "        self.decoder = nn.Linear(hidden_size, output_size)\n",
        "    \n",
        "    def forward(self, input_seq, hidden_state):\n",
        "        embedding = self.embedding(input_seq)\n",
        "        output, hidden_state = self.rnn(embedding, hidden_state)\n",
        "        output = self.decoder(output)\n",
        "        hidden_state = (hidden_state[0].detach(), hidden_state[1].detach())\n",
        "        return output, hidden_state\n"
      ],
      "metadata": {
        "id": "hRwhJo8SyMvV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will now get shakespearean text to get started. We will download the text from the link below. Click on the text to see text and read it. "
      ],
      "metadata": {
        "id": "vE32AKLn03_6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
      ],
      "metadata": {
        "id": "BZSNNcQXzoAU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will now define the main train loop with adagrad as the optimiser. The training loop iterates over the epochs, and for each epoch, it randomly selects a starting point in the data and trains the RNN model on sequences of length seq_len. The loss is computed, and the optimizer updates the model parameters. After each epoch, the model's state is saved to the specified file, and a new sequence of text is generated by randomly selecting a starting character and iteratively sampling the RNN output to generate the next character until the total number of characters in the generated sequence (op_seq_len) is reached. The generated text is printed to the console after each epoch."
      ],
      "metadata": {
        "id": "xWqUbsCM1FCj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def train():\n",
        "    ########### Hyperparameters ###########\n",
        "    hidden_size = 512   # size of hidden state\n",
        "    seq_len = 100       # length of LSTM sequence\n",
        "    num_layers = 3      # num of layers in LSTM layer stack\n",
        "    lr = 0.002          # learning rate\n",
        "    epochs = 100        # max number of epochs\n",
        "    op_seq_len = 200    # total num of characters in output test sequence\n",
        "    save_path = \"charRNN_shakespeare.pth\"\n",
        "    data_path = \"input.txt\"\n",
        "    #######################################\n",
        "\n",
        "    # load the text file\n",
        "    data = open(data_path, 'r').read()\n",
        "    chars = sorted(list(set(data)))\n",
        "    data_size, vocab_size = len(data), len(chars)\n",
        "    print(\"----------------------------------------\")\n",
        "    print(\"Data has {} characters, {} unique\".format(data_size, vocab_size))\n",
        "    print(\"----------------------------------------\")\n",
        "\n",
        "    # char to index and index to char maps\n",
        "    char_to_ix = { ch:i for i,ch in enumerate(chars) }\n",
        "    ix_to_char = { i:ch for i,ch in enumerate(chars) }\n",
        "\n",
        "    # convert data from chars to indices\n",
        "    data = list(data)\n",
        "    for i, ch in enumerate(data):\n",
        "        data[i] = char_to_ix[ch]\n",
        "\n",
        "    # data tensor on device\n",
        "    data = torch.tensor(data).to(device)\n",
        "    data = torch.unsqueeze(data, dim=1)\n",
        "\n",
        "    # model instance\n",
        "    rnn = RNN(vocab_size, vocab_size, hidden_size, num_layers).to(device)\n",
        "\n",
        "    # loss function and optimizer\n",
        "    loss_fn = nn.CrossEntropyLoss()\n",
        "    optimizer = torch.optim.Adagrad(rnn.parameters(), lr=lr)\n",
        "\n",
        "    # training loop\n",
        "    for i_epoch in range(1, epochs+1):\n",
        "\n",
        "        # random starting point (1st 10000 chars) from data to begin\n",
        "        data_ptr = np.random.randint(10000)\n",
        "        n = 0\n",
        "        running_loss = 0\n",
        "        hidden_state = None\n",
        "\n",
        "        while True:\n",
        "            input_seq = data[data_ptr : data_ptr+seq_len]\n",
        "            target_seq = data[data_ptr+1 : data_ptr+seq_len+1]\n",
        "\n",
        "            # forward pass\n",
        "            output, hidden_state = rnn(input_seq, hidden_state)\n",
        "\n",
        "            # compute loss\n",
        "            loss = loss_fn(torch.squeeze(output), torch.squeeze(target_seq))\n",
        "            running_loss += loss.item()\n",
        "\n",
        "            # compute gradients and take optimizer step\n",
        "            optimizer.zero_grad()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "            # update the data pointer\n",
        "            data_ptr += seq_len\n",
        "            n +=1\n",
        "\n",
        "            # if at end of data : break\n",
        "            if data_ptr + seq_len + 1 > data_size:\n",
        "                break\n",
        "\n",
        "        # print loss and save weights after every epoch\n",
        "        print(\"Epoch: {0} \\t Loss: {1:.8f}\".format(i_epoch, running_loss/n))\n",
        "        torch.save(rnn.state_dict(), save_path)\n",
        "\n",
        "        # sample / generate a text sequence after every epoch\n",
        "        data_ptr = 0\n",
        "        hidden_state = None\n",
        "\n",
        "        # random character from data to begin\n",
        "        rand_index = np.random.randint(data_size-1)\n",
        "        input_seq = data[rand_index : rand_index+1]\n",
        "\n",
        "        print(\"----------------------------------------\")\n",
        "        while True:\n",
        "            # forward pass\n",
        "            output, hidden_state = rnn(input_seq, hidden_state)\n",
        "\n",
        "            # construct categorical distribution and sample a character\n",
        "            output = F.softmax(torch.squeeze(output), dim=0)\n",
        "            dist = Categorical(output)\n",
        "            index = dist.sample()\n",
        "\n",
        "            # print the sampled character\n",
        "            print(ix_to_char[index.item()], end='')\n",
        "\n",
        "            # next input is\n",
        "            input_seq[0][0] = index.item()\n",
        "            data_ptr += 1\n",
        "            \n",
        "            if data_ptr > op_seq_len:\n",
        "                break\n",
        "            \n",
        "        print(\"\\n----------------------------------------\")\n",
        "        \n"
      ],
      "metadata": {
        "id": "lkbGUJdayOhH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let us now begin training. The code is going to sample after every epoch - notice the evolution of the generated text. "
      ],
      "metadata": {
        "id": "GzPhHTPF3g2b"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train() # call the training loop"
      ],
      "metadata": {
        "id": "d29pge1dzhZn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Neural sequence taggers\n",
        "\n",
        "We will now explore an LSTM-based Part-Of-Speech (POS) tagger model using pytorch.\n",
        "\n",
        "Notice that the structure is very similar to the previous example. However, here, the forward method takes a sentence as input and returns the predicted tag scores for each word in the sentence.\n",
        "\n"
      ],
      "metadata": {
        "id": "Q9B_0_Oo2uhi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "\n",
        "# Define the LSTM-based POS tagger model\n",
        "class LSTMTagger(nn.Module):\n",
        "    def __init__(self, vocab_size, embedding_dim, hidden_dim, tagset_size):\n",
        "        super(LSTMTagger, self).__init__()\n",
        "        self.embedding_dim = embedding_dim\n",
        "        self.hidden_dim = hidden_dim\n",
        "        self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)\n",
        "        self.lstm = nn.LSTM(embedding_dim, hidden_dim)\n",
        "        self.hidden2tag = nn.Linear(hidden_dim, tagset_size)\n",
        "\n",
        "    def forward(self, sentence):\n",
        "        embeds = self.word_embeddings(sentence)\n",
        "        lstm_out, _ = self.lstm(embeds.view(len(sentence), 1, -1))\n",
        "        tag_space = self.hidden2tag(lstm_out.view(len(sentence), -1))\n",
        "        tag_scores = nn.functional.log_softmax(tag_space, dim=1)\n",
        "        return tag_scores\n",
        "\n",
        "# Define the training function\n",
        "def train(model, optimizer, loss_function, sentences, tags, num_epochs):\n",
        "    for epoch in range(num_epochs):\n",
        "        for sentence, tag in zip(sentences, tags):\n",
        "            model.zero_grad()\n",
        "            sentence = torch.tensor(sentence, dtype=torch.long).to(device)\n",
        "            tag = torch.tensor(tag, dtype=torch.long).to(device)\n",
        "            tag_scores = model(sentence)\n",
        "            loss = loss_function(tag_scores, tag)\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "# Define the evaluation function\n",
        "def evaluate(model, sentences, tags):\n",
        "    correct = 0\n",
        "    total = 0\n",
        "    with torch.no_grad():\n",
        "        for sentence, tag in zip(sentences, tags):\n",
        "            sentence = torch.tensor(sentence, dtype=torch.long).to(device)\n",
        "            tag = torch.tensor(tag, dtype=torch.long).to(device)\n",
        "            tag_scores = model(sentence)\n",
        "            _, predicted = torch.max(tag_scores.data, 1)\n",
        "            total += tag.size(0)\n",
        "            correct += (predicted == tag).sum().item()\n",
        "    accuracy = 100 * correct / total\n",
        "    return accuracy\n",
        "\n",
        "# Example usage\n",
        "# Define a sample corpus and its corresponding tags\n",
        "corpus = [\n",
        "    \"The cat sat on the mat\",\n",
        "    \"The dog chased the cat\",\n",
        "    \"The mouse ran away from the cat\",\n",
        "    \"The cat purred\",\n",
        "]\n",
        "tags = [\n",
        "    \"DET NOUN VERB ADP DET NOUN\",\n",
        "    \"DET NOUN VERB DET NOUN\",\n",
        "    \"DET NOUN VERB ADV ADP DET NOUN\",\n",
        "    \"DET NOUN VERB\",\n",
        "]\n",
        "# Define the vocabulary and POS tagset\n",
        "word_to_ix = {\"<PAD>\": 0, \"<UNK>\": 1, \"The\": 2, \"cat\": 3, \"sat\": 4, \"on\": 5, \"the\": 6,\n",
        "              \"mat\": 7, \"dog\": 8, \"chased\": 9, \"mouse\": 10, \"ran\": 11, \"away\": 12, \"from\": 13, \"purred\": 14}\n",
        "tag_to_ix = {\"<PAD>\": 0, \"<UNK>\": 1, \"DET\": 2, \"NOUN\": 3, \"VERB\": 4, \"ADP\": 5, \"ADV\": 6}\n",
        "\n",
        "# Convert the corpus and tagset to indices\n",
        "sentences = [[word_to_ix.get(word, word_to_ix[\"<UNK>\"]) for word in sentence.split()] for sentence in corpus]\n",
        "tags = [[tag_to_ix.get(tag, tag_to_ix[\"<UNK>\"]) for tag in sentence.split()] for sentence in tags]\n",
        "\n",
        "# Set hyperparameters and create model, optimizer, and loss function instances\n",
        "vocab_size = len(word_to_ix)\n",
        "embedding_dim = 16\n",
        "hidden_dim = 16\n",
        "tagset_size = len(tag_to_ix)\n",
        "num_epochs = 10\n",
        "learning_rate = 0.1\n",
        "model = LSTMTagger(vocab_size, embedding_dim, hidden_dim, tagset_size).to(device)\n",
        "optimizer = optim.SGD(model.parameters(), lr=learning_rate)\n",
        "loss_function = nn.NLLLoss()\n",
        "#Train the model\n",
        "train(model, optimizer, loss_function, sentences, tags, num_epochs)\n",
        "\n",
        "#Evaluate the model on the same corpus\n",
        "accuracy = evaluate(model, sentences, tags)\n",
        "print(f\"Accuracy: {accuracy:.2f}%\")\n"
      ],
      "metadata": {
        "id": "7gIxryrlv3E0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Todo: \n",
        "- Can you train this on a standard POS dataset from https://github.com/dan-oak/pos/tree/master/data (look for the tagged ones as train and dev set). "
      ],
      "metadata": {
        "id": "ZKiob8_C5Ug7"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "MszfV7EX5SSY"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}