QA_LSTM.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "QA_LSTM.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m2ye9_3R1IDV",
        "colab_type": "text"
      },
      "source": [
        "# **QA-LSTM**\n",
        "As an alternative to running the scripts on the command line, this notebook shows how to train and evaluate a baseline LSTM model for non-factoid question answering."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CBPpjAOmp51l",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!git clone https://github.com/yuanbit/FinBERT-QA\n",
        "\n",
        "%cd FinBERT-QA/\n",
        "from src.utils import *\n",
        "from src.evaluate import *"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z9NOakGS8WnB",
        "colab_type": "code",
        "outputId": "2b52e1b9-8318-4f1e-a18f-0e2e3f4bcb0a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 146
        }
      },
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import random\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n",
        "import torchtext\n",
        "from tqdm import tqdm\n",
        "\n",
        "# Setting device on GPU if available\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print('Using device:', device)\n",
        "print()\n",
        "\n",
        "if device.type == 'cuda':\n",
        "    print(torch.cuda.get_device_name(0))\n",
        "    print('Memory Usage:')\n",
        "    print('Allocated:', round(torch.cuda.memory_allocated(0)/1024**3,1), 'GB')\n",
        "    print('Cached:   ', round(torch.cuda.memory_cached(0)/1024**3,1), 'GB')\n",
        "\n",
        "torch.backends.cudnn.deterministic = True\n",
        "\n",
        "# Set the random seed manually for reproducibility.\n",
        "torch.manual_seed(1234)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Using device: cuda\n",
            "\n",
            "Tesla P100-PCIE-16GB\n",
            "Memory Usage:\n",
            "Allocated: 0.0 GB\n",
            "Cached:    0.0 GB\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x7f8e95b3ee50>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5InBSV7w1egQ",
        "colab_type": "text"
      },
      "source": [
        "## **Configure hyperparameters and load data**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2c-0GAEUVYJx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "config = {'max_seq_len': 128,\n",
        "          'batch_size': 64,\n",
        "          'n_epochs': 3,\n",
        "          'lr': 1e-3,\n",
        "          'emb_dim': 100,\n",
        "          'hidden_size': 256,\n",
        "          'dropout': 0.2,\n",
        "          'margin': 0.2}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K9oexhEqqJd0",
        "colab_type": "code",
        "outputId": "318dc578-a3f8-42aa-e92c-68d8076b9f42",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 72
        }
      },
      "source": [
        "# Dictonary with token to id mapping\n",
        "vocab = load_pickle(\"data/qa_lstm_tokenizer/word2index.pickle\")\n",
        "# Dictonary with qid to tokenized text mapping\n",
        "qid_to_tokenized_text = load_pickle('data/qa_lstm_tokenizer/qid_to_tokenized_text.pickle')\n",
        "# Dictionary with docid to tokenized text mapping\n",
        "docid_to_tokenized_text = load_pickle('data/qa_lstm_tokenizer/docid_to_tokenized_text.pickle')\n",
        "\n",
        "# List of lists:\n",
        "# Each element is a list contraining [qid, list of pos docid, list of candidate docid]\n",
        "train_set = load_pickle('data/data_pickle/train_set_50.pickle')\n",
        "valid_set = load_pickle('data/data_pickle/valid_set_50.pickle')\n",
        "test_set = load_pickle('data/data_pickle/test_set_50.pickle')\n",
        "\n",
        "# Labels\n",
        "labels = load_pickle('data/data_pickle/labels.pickle')\n",
        "\n",
        "print(\"Number of questions in the training set: {}\".format(len(train_set)))\n",
        "print(\"Number of questions in the validation set: {}\".format(len(valid_set)))\n",
        "print(\"Number of questions in the test set: {}\".format(len(test_set)))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Number of questions in the training set: 5676\n",
            "Number of questions in the validation set: 631\n",
            "Number of questions in the test set: 333\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wcX7mbfDWKhP",
        "colab_type": "text"
      },
      "source": [
        "## **Prepare data**\n",
        "Create the required inputs and convert them to DataLoader objects to be trained in batches.\n",
        "\n",
        "**Required training inputs:**\n",
        "1. question input ids\n",
        "2. positive answer input ids\n",
        "3. negative answer input ids"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8fwqmca0M1PD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def pad_seq(seq_idx, max_seq_len):\n",
        "    \"\"\"Creates padded or truncated sequence.\n",
        "\n",
        "    Returns:\n",
        "        seq: list of padded vectorized sequence\n",
        "    ----------\n",
        "    Arguements:\n",
        "        seq_idx: tensor with similarity of a question and a positive answer\n",
        "        max_seq_len: int\n",
        "    \"\"\"\n",
        "    # Pad each sequence to be the same length to process in batches\n",
        "    # pad_token = 0\n",
        "    if len(seq_idx) >= max_seq_len:\n",
        "        seq_idx = seq_idx[:max_seq_len]\n",
        "    else:\n",
        "        seq_idx += [0]*(max_seq_len - len(seq_idx))\n",
        "    seq = seq_idx\n",
        "\n",
        "    return seq\n",
        "\n",
        "def vectorize(seq, max_seq_len):\n",
        "    \"\"\"Creates vectorized sequence.\n",
        "\n",
        "    Returns:\n",
        "        vectorized_seq: List of padded vectorized sequence\n",
        "    ----------\n",
        "    Arguements:\n",
        "        seq: List of tokens in a sequence\n",
        "        max_seq_len: int\n",
        "    \"\"\"\n",
        "    # Map tokens in seq to idx\n",
        "    seq_idx = [vocab[token] for token in seq]\n",
        "    # Pad seq idx\n",
        "    vectorized_seq = pad_seq(seq_idx, max_seq_len)\n",
        "\n",
        "    return vectorized_seq"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-VJ15qnBW9jS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_lstm_input_data(dataset, max_seq_len):\n",
        "    \"\"\"Creates input data for model.\n",
        "\n",
        "    Returns:\n",
        "        q_input_ids: List of lists of vectorized question sequence\n",
        "        pos_input_ids: List of lists of vectorized positve ans sequence\n",
        "        neg_input_ids: List of lists of vectorized negative ans sequence\n",
        "    ----------\n",
        "    Arguements:\n",
        "        dataset: List of lists in the form of [qid, [pos ans], [ans cands]]\n",
        "        max_seq_len: int\n",
        "    \"\"\"\n",
        "    q_input_ids = []\n",
        "    pos_input_ids = []\n",
        "    neg_input_ids = []\n",
        "\n",
        "    for i, seq in enumerate(tqdm(dataset)):\n",
        "        qid, ans_labels, cands = seq[0], seq[1], seq[2]\n",
        "\n",
        "        # Remove the positive answers for the candidates\n",
        "        filtered_cands = list(set(cands)-set(ans_labels))\n",
        "        # Select a positive answer from the list of positive answers\n",
        "        pos_docid = random.choice(ans_labels)\n",
        "        # Map question id to text\n",
        "        q_text = qid_to_tokenized_text[qid]\n",
        "        # Pad and vectorize text\n",
        "        q_input_id = vectorize(q_text, max_seq_len)\n",
        "\n",
        "        # For all the negative answers\n",
        "        for neg_docid in filtered_cands:\n",
        "            # Map the docid to text\n",
        "            pos_ans_text = docid_to_tokenized_text[pos_docid]\n",
        "            neg_ans_text = docid_to_tokenized_text[neg_docid]\n",
        "            # Pad and vectorize sequences\n",
        "            pos_input_id = vectorize(pos_ans_text, max_seq_len)\n",
        "            neg_input_id = vectorize(neg_ans_text, max_seq_len)\n",
        "\n",
        "            q_input_ids.append(q_input_id)\n",
        "            pos_input_ids.append(pos_input_id)\n",
        "            neg_input_ids.append(neg_input_id)\n",
        "\n",
        "    return q_input_ids, pos_input_ids, neg_input_ids"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bm11X-X7XAxL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_dataloader(dataset, type, max_seq_len, batch_size):\n",
        "    \"\"\"Creates train and validation DataLoaders with question, positive\n",
        "    answer, and negative answer vectorized inputs.\n",
        "\n",
        "    Returns:\n",
        "        train_dataloader: DataLoader object\n",
        "        validation_dataloader: DataLoader object\n",
        "    ----------\n",
        "    Arguements:\n",
        "        dataset: List of lists in the form of [qid, [pos ans], [ans cands]]\n",
        "        type: str - 'train' or 'validation'\n",
        "        max_seq_len: int\n",
        "        batch_size: int\n",
        "    \"\"\"\n",
        "    question_input, pos_ans_input, neg_ans_input = get_lstm_input_data(dataset, max_seq_len)\n",
        "\n",
        "    question_inputs = torch.tensor(question_input)\n",
        "    pos_ans_inputs = torch.tensor(pos_ans_input)\n",
        "    neg_ans_inputs = torch.tensor(neg_ans_input)\n",
        "\n",
        "    # Create the DataLoader\n",
        "    data = TensorDataset(question_inputs, pos_ans_inputs, neg_ans_inputs)\n",
        "    if type == \"train\":\n",
        "        sampler = RandomSampler(data)\n",
        "    else:\n",
        "        sampler = SequentialSampler(data)\n",
        "    dataloader = DataLoader(data, sampler=sampler, batch_size=batch_size)\n",
        "\n",
        "    return dataloader"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OHn1sFajaL6Y",
        "colab_type": "code",
        "outputId": "8505ff2d-5990-4a0d-a6fe-07b2a2861310",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 127
        }
      },
      "source": [
        "# Get dataloaders\n",
        "train_dataloader = get_dataloader(train_set, 'train', \n",
        "                                  config['max_seq_len'], \n",
        "                                  config['batch_size'])\n",
        "validation_dataloader = get_dataloader(valid_set, 'validation', \n",
        "                                       config['max_seq_len'], \n",
        "                                       config['batch_size'])\n",
        "\n",
        "print(\"\\n\\nSize of the training DataLoader: {}\".format(len(train_dataloader)))\n",
        "print(\"Size of the validation DataLoader: {}\".format(len(validation_dataloader)))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 5676/5676 [00:13<00:00, 414.70it/s]\n",
            "100%|██████████| 631/631 [00:01<00:00, 472.74it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "Size of the training DataLoader: 4342\n",
            "Size of the validation DataLoader: 483\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tkbw3PcRazLT",
        "colab_type": "text"
      },
      "source": [
        "## **Model**\n",
        "Uses pre-trained GloVe embeddings and a biLSTM network with the cosine similarity measure."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E3-VHSXrLSfF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class QA_LSTM(nn.Module):\n",
        "    \"\"\"\n",
        "    QA-LSTM model\n",
        "    \"\"\"\n",
        "    def __init__(self, config):\n",
        "        super(QA_LSTM, self).__init__()\n",
        "        # Embedding dimension\n",
        "        self.emb_dim = config['emb_dim']\n",
        "        # Hidden size\n",
        "        self.hidden_size = config['hidden_size']\n",
        "        # Dropout rate\n",
        "        self.dropout = config['dropout']\n",
        "        # Vocabulary size\n",
        "        self.vocab_size = len(vocab)\n",
        "        # Create embedding layer\n",
        "        self.embedding = self.create_emb_layer()\n",
        "        # The question and answer representations share the same biLSTM network\n",
        "        self.lstm = nn.LSTM(self.emb_dim, \\\n",
        "                            self.hidden_size, \\\n",
        "                            num_layers=1, \\\n",
        "                            batch_first=True, \\\n",
        "                            bidirectional=True)\n",
        "        # Cosine similiarty metric\n",
        "        self.sim = nn.CosineSimilarity(dim=1)\n",
        "        # Apply dropout\n",
        "        self.dropout = nn.Dropout(self.dropout)\n",
        "\n",
        "    def create_emb_layer(self):\n",
        "        \"\"\"Creates embedding layerself using pre-trained\n",
        "        GloVe embeddings (6B tokens)\n",
        "\n",
        "        Returns:\n",
        "            emb_layer: Torch embedding layer\n",
        "        \"\"\"\n",
        "        print(\"\\n\\nInitializing model...\")\n",
        "        print(\"\\nDownloading pre-trained GloVe embeddings...\\n\")\n",
        "        # Use GloVe embeddings from torchtext\n",
        "        emb = torchtext.vocab.GloVe(\"6B\", dim=self.emb_dim)\n",
        "        # Dictionary mapping of word idx to GloVe vectors\n",
        "        emb_weights = np.zeros((self.vocab_size, self.emb_dim))\n",
        "        # Count\n",
        "        words_found = 0\n",
        "\n",
        "        for token, idx in vocab.items():\n",
        "            # emb.stoi is a dict of token to idx mapping\n",
        "            # If token from the vocabulary exist in GloVe\n",
        "            if token in emb.stoi:\n",
        "                # Add the embedding to the index\n",
        "                emb_weights[idx] = emb[token]\n",
        "                words_found += 1\n",
        "\n",
        "        print(\"\\n\")\n",
        "        print(words_found, \"words are found in GloVe\")\n",
        "\n",
        "        # Convert matrix to tensor\n",
        "        emb_weights = torch.from_numpy(emb_weights).float()\n",
        "\n",
        "        vocab_size, emb_dim = emb_weights.shape\n",
        "        # Create embedding layer\n",
        "        emb_layer = nn.Embedding(vocab_size, emb_dim)\n",
        "        # Load the embeddings\n",
        "        emb_layer.load_state_dict({'weight': emb_weights})\n",
        "\n",
        "        return emb_layer\n",
        "\n",
        "    def forward(self, question, answer):\n",
        "        \"\"\"Forward pass to generate biLSTM representations for the question and\n",
        "        answer independently, and then utilize cosine similarity to measure\n",
        "        their distance.\n",
        "\n",
        "        Returns:\n",
        "            similarity: Torch tensor with cosine similarity score.\n",
        "        ----------\n",
        "        Arguements:\n",
        "            question: Torch tensor of vectorized question\n",
        "            answer: Torch tensor of vectorized answer\n",
        "        \"\"\"\n",
        "        # Embedding layers - (batch_size, max_seq_len, emb_dim)\n",
        "        question_embedding = self.embedding(question)\n",
        "        answer_embedding = self.embedding(answer)\n",
        "\n",
        "        # biLSTM - (batch_size, max_seq_len, 2*hidden_size)\n",
        "        question_lstm, (hidden, cell) = self.lstm(question_embedding)\n",
        "        answer_lstm, (hidden, cell) = self.lstm(answer_embedding)\n",
        "\n",
        "        # Max-pooling - (batch_size, 2*hidden_size)\n",
        "        # There are n word level biLSTM representations where n is the max_seq_len\n",
        "        # Use max pooling to generate the best representation\n",
        "        question_maxpool = torch.max(question_lstm, 1)[0]\n",
        "        answer_maxpool = torch.max(answer_lstm, 1)[0]\n",
        "\n",
        "        # Apply dropout\n",
        "        question_output = self.dropout(question_maxpool)\n",
        "        answer_output = self.dropout(answer_maxpool)\n",
        "\n",
        "        # Similarity -(batch_size,)\n",
        "        similarity = self.sim(question_output, answer_output)\n",
        "\n",
        "        return similarity"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sj-tbs6Wb6dw",
        "colab_type": "code",
        "outputId": "82179fc1-7e01-40d4-a0e0-050d1aa7c244",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 182
        }
      },
      "source": [
        "model = QA_LSTM(config).to(device)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "Initializing model...\n",
            "\n",
            "Downloading pre-trained GloVe embeddings...\n",
            "\n",
            "\n",
            "\n",
            "50456 words are found in GloVe\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rJYLALN4cLuk",
        "colab_type": "text"
      },
      "source": [
        "## **Training/Validation methods**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C5As531eJCLX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def hinge_loss(pos_sim, neg_sim):\n",
        "    \"\"\"\n",
        "    Returns:\n",
        "        loss: Tensor with hinge loss value\n",
        "    ----------\n",
        "    Arguements:\n",
        "        pos_sim: Tensor with similarity of a question and a positive answer\n",
        "        neg_sim: Tensor with similarity of a question and a negative answer\n",
        "    \"\"\"\n",
        "    margin = config['margin']\n",
        "\n",
        "    loss = torch.max(torch.tensor(0, dtype=torch.float).to(device), \n",
        "                     margin - pos_sim + neg_sim)\n",
        "    return loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z-GFmyfiAXXd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(model, train_dataloader, optimizer):\n",
        "    \"\"\"Trains the model and returns the average loss\n",
        "\n",
        "    Returns:\n",
        "        avg_loss: Float\n",
        "    ----------\n",
        "    Arguements:\n",
        "        model: Torch model\n",
        "        train_dataloader: DataLoader object\n",
        "        optimizer: Optimizer object\n",
        "    \"\"\"\n",
        "    # Cumulated Training loss\n",
        "    train_loss = 0.0\n",
        "    # Set model to training mode\n",
        "    model.train()\n",
        "    # For each batch of training data\n",
        "    for step, batch in enumerate(tqdm(train_dataloader)):\n",
        "        # batch contains 3 PyTorch tensors\n",
        "        # Move tensors to gpu\n",
        "        question = batch[0].to(device)\n",
        "        pos_ans = batch[1].to(device)\n",
        "        neg_ans = batch[2].to(device)\n",
        "\n",
        "        # 1. Zero gradients\n",
        "        model.zero_grad()\n",
        "        # 2. Compute similarity scores of pos and neg QA pairs\n",
        "        pos_sim = model(question, pos_ans)\n",
        "        neg_sim = model(question, neg_ans)\n",
        "        # 3. Compute loss\n",
        "        loss = hinge_loss(pos_sim, neg_sim).mean()\n",
        "        # 4. Use loss to compute gradients\n",
        "        loss.backward()\n",
        "        # 5. Use optimizer to take gradient step\n",
        "        optimizer.step()\n",
        "        # Cumulate loss\n",
        "        train_loss += loss.item()\n",
        "    # Compute average loss\n",
        "    avg_loss = train_loss/len(train_dataloader)\n",
        "\n",
        "    return avg_loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zhQfdIU_GEjs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def validate(model, validation_dataloader):\n",
        "    \"\"\"Validates the model and returns the average loss\n",
        "\n",
        "    Returns:\n",
        "        avg_loss: Float\n",
        "    ----------\n",
        "    Arguements:\n",
        "        model: Torch model\n",
        "        validation_dataloader: DataLoader object\n",
        "    \"\"\"\n",
        "    # Cumulated validation loss\n",
        "    valid_loss = 0.0\n",
        "    # Set model to evaluation mode\n",
        "    model.eval()\n",
        "    # Evaluate data\n",
        "    for batch in tqdm(validation_dataloader):\n",
        "        # Add batch to GPU\n",
        "        batch = tuple(t.to(device) for t in batch)\n",
        "        # Unpack the inputs from Dataloader\n",
        "        question, pos_ans, neg_ans = batch\n",
        "        # Don't calculate the gradients\n",
        "        with torch.no_grad():\n",
        "            # Compute similarity score of pos and neg QA pairs\n",
        "            pos_sim = model(question, pos_ans)\n",
        "            neg_sim = model(question, neg_ans)\n",
        "            # Compute loss\n",
        "            loss = hinge_loss(pos_sim, neg_sim).mean()\n",
        "            # Coumulate loss\n",
        "            valid_loss += loss.item()\n",
        "    # Compute average loss\n",
        "    avg_loss = valid_loss/len(validation_dataloader)\n",
        "\n",
        "    return avg_loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "of0-uXlwc7LV",
        "colab_type": "text"
      },
      "source": [
        "## **Train model**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7uJv3UkoqvLh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TotmjsQTc6qs",
        "colab_type": "code",
        "outputId": "f1fc208e-b29a-4616-c012-f8b1270e3f61",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 458
        }
      },
      "source": [
        "optimizer = optim.Adam(model.parameters(), lr=config['lr'])\n",
        "\n",
        "# Lowest validation lost\n",
        "best_valid_loss = float('inf')\n",
        "\n",
        "for epoch in range(config['n_epochs']):\n",
        "    # Evaluate training loss\n",
        "    train_loss = train(model, train_dataloader, optimizer)\n",
        "    # Evaluate validation loss\n",
        "    valid_loss = validate(model, validation_dataloader)\n",
        "    # At each epoch, if the validation loss is the best\n",
        "    if valid_loss < best_valid_loss:\n",
        "        best_valid_loss = valid_loss\n",
        "        # Save the parameters of the model\n",
        "        torch.save(model.state_dict(), 'model/'+str(epoch+1)+'_qa_lstm.pt')\n",
        "\n",
        "    print(\"\\n\\n Epoch {}:\".format(epoch+1))\n",
        "    print(\"\\t Train Loss: {0:.3f}\".format(train_loss))\n",
        "    print(\"\\t Validation Loss: {0:.3f}\\n\".format(valid_loss))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 4342/4342 [05:48<00:00, 12.44it/s]\n",
            "100%|██████████| 483/483 [00:15<00:00, 31.28it/s]\n",
            "  0%|          | 2/4342 [00:00<06:31, 11.09it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            " Epoch 1:\n",
            "\t Train Loss: 0.073\n",
            "\t Validation Loss: 0.269\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 4342/4342 [05:47<00:00, 12.48it/s]\n",
            "100%|██████████| 483/483 [00:15<00:00, 31.27it/s]\n",
            "  0%|          | 2/4342 [00:00<06:27, 11.19it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            " Epoch 2:\n",
            "\t Train Loss: 0.022\n",
            "\t Validation Loss: 0.267\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 4342/4342 [05:48<00:00, 12.46it/s]\n",
            "100%|██████████| 483/483 [00:15<00:00, 31.41it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            " Epoch 3:\n",
            "\t Train Loss: 0.012\n",
            "\t Validation Loss: 0.258\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LQNoBE_VeS34",
        "colab_type": "text"
      },
      "source": [
        "## **Evaluate**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s9YIIOZt3waV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_rank(model, test_set, max_seq_len):\n",
        "    \"\"\"Re-ranks the answer candidates per question using trained model.\n",
        "\n",
        "    Returns:\n",
        "        qid_pred_rank: Dictionary\n",
        "                key - qid\n",
        "                value - List of re-ranked candidate answers\n",
        "    -------------------\n",
        "    Arguments:\n",
        "        model - Trained PyTorch model\n",
        "        test_set - List of lists\n",
        "                Each element is a list contraining\n",
        "                [qid, list of pos docid, list of candidate docid]\n",
        "        max_seq_len - int   \n",
        "    \"\"\"\n",
        "    # Dictionary - key: qid, value: ranked list of docids\n",
        "    qid_pred_rank = {}\n",
        "    # Set model to evaluation mode\n",
        "    model.eval()\n",
        "    # For each sample in the test set\n",
        "    for i, seq in enumerate(tqdm(test_set)):\n",
        "        # Extract input data\n",
        "        ques, pos_ans, cands = seq[0], seq[1], seq[2]\n",
        "        # Tokenize and vectorize question\n",
        "        q_text = qid_to_tokenized_text[ques]\n",
        "        q_vec = torch.tensor([vectorize(q_text, max_seq_len)]).to(device)\n",
        "        # Tokenize candidate answers\n",
        "        cands_text = [docid_to_tokenized_text[c] for c in cands]\n",
        "        cands_id = np.array(cands)\n",
        "        # List to store similarity score of QA pair\n",
        "        scores = []\n",
        "        # For each candidate answer\n",
        "        for cand in cands_text:\n",
        "            # Vectorize the answers\n",
        "            a_vec = torch.tensor([vectorize(cand, max_seq_len)]).to(device)\n",
        "            # Compute similarity score of QA pair and add to scores\n",
        "            scores.append(model(q_vec, a_vec).item())\n",
        "\n",
        "        # Get the indices of the sorted (descending) similarity scores\n",
        "        sorted_index = np.argsort(scores)[::-1]\n",
        "        # Get the cand ans docid from the sorted indices\n",
        "        ranked_ans = cands_id[sorted_index]\n",
        "        # Set the qid keys to the list of re-ranked docids\n",
        "        qid_pred_rank[ques] = ranked_ans\n",
        "\n",
        "    return qid_pred_rank"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dej0NZ8xrBjn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Download trained model\n",
        "checkpoint = get_trained_model('qa-lstm')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X8TCRr4YMrJ7",
        "colab_type": "code",
        "outputId": "a3778631-c6fb-4537-aad5-9210af5a3891",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 127
        }
      },
      "source": [
        "k = 10\n",
        "\n",
        "trained_model_path = \"model/trained/qa-lstm/\" + checkpoint\n",
        "# Load model\n",
        "model.load_state_dict(torch.load(trained_model_path), strict=False)\n",
        "\n",
        "# Get rank\n",
        "qid_pred_rank = get_rank(model, test_set, config['max_seq_len'])\n",
        "\n",
        "k = 10\n",
        "num_q = len(test_set)\n",
        "\n",
        "# Evaluate\n",
        "MRR, average_ndcg, precision, rank_pos = evaluate(qid_pred_rank, labels, k)\n",
        "\n",
        "print(\"\\n\\nAverage nDCG@{0} for {1} queries: {2:.3f}\".format(k, num_q, average_ndcg))\n",
        "print(\"MRR@{0} for {1} queries: {2:.3f}\".format(k, num_q, MRR))\n",
        "print(\"Average Precision@1 for {0} queries: {1:.3f}\".format(num_q, precision))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 333/333 [01:49<00:00,  3.04it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "Average nDCG@10 for 333 queries: 0.115\n",
            "MRR@10 for 333 queries: 0.073\n",
            "Average Precision@1 for 333 queries: 0.033\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stderr"
        }
      ]
    }
  ]
}