FinBERT_QA.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "FinBERT_QA.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": "ZYTJmWOgraFa",
        "colab_type": "text"
      },
      "source": [
        "# **FinBERT-QA**\n",
        "As an alternative to running the scripts on the command line, this notebook shows the process of fine-tuning a pre-trained BERT model to the Opionated Financial Question and Answering (FiQA) dataset.\n",
        "\n",
        "The evaluation, prediction, and analysis of the FinBERT-QA model can be found after the fine-tuning process."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "usfS4Iw1psfC"
      },
      "source": [
        "## **Setup**\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zBMhV22cZpDI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install transformers\n",
        "!pip install pyserini==0.8.1.0"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ySXHdHsEUxNd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!git clone https://github.com/yuanbit/FinBERT-QA"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WdAFUpPyk-4v",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%cd FinBERT-QA/\n",
        "from src.utils import *\n",
        "from src.evaluate import *"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qyJO_z_skUvv",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 126
        },
        "outputId": "d3d39beb-9837-44e3-bf8b-4447e5a436b0"
      },
      "source": [
        "from tqdm import tqdm\n",
        "import numpy as np\n",
        "import random\n",
        "import torch\n",
        "import json\n",
        "import os\n",
        "from torch.nn import CrossEntropyLoss\n",
        "from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler\n",
        "from torch.nn.functional import softmax\n",
        "from transformers import BertTokenizer, BertForSequenceClassification, AdamW, get_linear_schedule_with_warmup, BertConfig\n",
        "\n",
        "os.environ[\"JAVA_HOME\"] = \"/usr/lib/jvm/java-11-openjdk-amd64\"\n",
        "from pyserini.search import pysearch\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')"
      ],
      "execution_count": 4,
      "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"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xz-pwt46vIcX",
        "colab_type": "text"
      },
      "source": [
        "## **Configure**\n",
        "\n",
        "**bert_model_name configuration:**\n",
        "\n",
        "1. 'bert-qa': uses a pre-trained BERT model fine-tuned on the MS Macro passage dataset of [Nogueira and Cho](https://github.com/nyu-dl/dl4marco-bert).\n",
        "2. 'finbert-domain': uses a further pre-trained BERT model of [Araci](https://github.com/ProsusAI/finBERT) on a large financial corpus\n",
        "3. 'finbert-task': uses a further pre-trained BERT model on the FiQA dataset\n",
        "4. 'bert-base': uses the 'bert-base-uncased' model from huggingface\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L4KGOdhviDpR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "config = {'bert_model_name': 'bert-qa',\n",
        "          'max_seq_len': 512,\n",
        "          'batch_size': 16,\n",
        "          'learning_rate': 3e-6,\n",
        "          'weight_decay': 0.01,\n",
        "          'n_epochs': 2,\n",
        "          'num_warmup_steps': 10000}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u2YEBYUoZvfG",
        "colab_type": "code",
        "outputId": "142b059c-7e84-425f-f1bb-b4200c3597d2",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 108
        }
      },
      "source": [
        "# Dictionary mapping docid and qid to raw text\n",
        "docid_to_text = load_pickle('data/id_to_text/docid_to_text.pickle')\n",
        "qid_to_text = load_pickle('data/id_to_text/qid_to_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)))\n",
        "\n",
        "# Load the BERT tokenizer.\n",
        "print('\\nLoading BERT tokenizer...')\n",
        "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)"
      ],
      "execution_count": 11,
      "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",
            "\n",
            "Loading BERT tokenizer...\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x9_cmDMdmnGj",
        "colab_type": "text"
      },
      "source": [
        "## **Prepare data**\n",
        "Create the required inputs for fine-tuning and convert them to DataLoader objects to be trained in batches.\n",
        "\n",
        "**Required input data:**\n",
        "1. input_ids of concatenated question and answer sequences\n",
        "2. token_type_ids are segment ids incidiating if the index is part of a question or answer\n",
        "3. att_masks are ids indicating if the index is part of the sequence or padding\n",
        "4. labels indicating if a QA pair is positive or negative"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1-hWS0v0sx7J",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_input_data(dataset, max_seq_len):\n",
        "    \"\"\"Creates input parameters for training and validation.\n",
        "\n",
        "    Returns:\n",
        "        input_ids: List of lists\n",
        "                Each element contains a list of padded/truncated numericalized\n",
        "                tokens of the sequences including [CLS] and [SEP] tokens\n",
        "                e.g. [[101, 2054, 2003, 102, 2449, 1029, 102], ...]\n",
        "        token_type_ids: List of lists\n",
        "                Each element contains a list of segment token indices to\n",
        "                indicate first (question) and second (answer) parts of the inputs.\n",
        "                0 corresponds to a question token, 1 corresponds an answer token\n",
        "                e.g. [[0, 0, 0, 0, 1, 1, 1], ...]\n",
        "        att_masks: List of lists\n",
        "                Each element contains a list of mask values to avoid\n",
        "                performing attention on padding token indices.\n",
        "                1 for tokens that are NOT MASKED, 0 for MASKED tokens.\n",
        "                e.g. [[1, 1, 1, 1, 1, 1, 1], ...]\n",
        "        labels: List of 1's and 0's incidating relevacy of answer\n",
        "    -----------------\n",
        "    Arguements:\n",
        "        dataset: List of lists in the form of [qid, [pos ans], [ans cands]]\n",
        "    \"\"\"\n",
        "    input_ids = []\n",
        "    token_type_ids = []\n",
        "    att_masks = []\n",
        "    labels = []\n",
        "\n",
        "    for i, seq in enumerate(tqdm(dataset)):\n",
        "        qid, ans_labels, cands = seq[0], seq[1], seq[2]\n",
        "        # Map question id to text\n",
        "        q_text = qid_to_text[qid]\n",
        "        # For each answer in the candidates\n",
        "        for docid in cands:\n",
        "            # Map the docid to text\n",
        "            ans_text = docid_to_text[docid]\n",
        "            # Encode the sequence using BERT tokenizer\n",
        "            encoded_seq = tokenizer.encode_plus(q_text, ans_text,\n",
        "                                                max_length=max_seq_len,\n",
        "                                                pad_to_max_length=True,\n",
        "                                                return_token_type_ids=True,\n",
        "                                                return_attention_mask = True)\n",
        "            # Get parameters\n",
        "            input_id = encoded_seq['input_ids']\n",
        "            token_type_id = encoded_seq['token_type_ids']\n",
        "            att_mask = encoded_seq['attention_mask']\n",
        "\n",
        "            # If an answer is in the list of relevant answers assign\n",
        "            # positive label\n",
        "            if docid in ans_labels:\n",
        "                label = 1\n",
        "            else:\n",
        "                label = 0\n",
        "\n",
        "            # Each parameter list has the length of the max_seq_len\n",
        "            assert len(input_id) == max_seq_len, \"Input id dimension incorrect!\"\n",
        "            assert len(token_type_id) == max_seq_len, \"Token type id dimension incorrect!\"\n",
        "            assert len(att_mask) == max_seq_len, \"Attention mask dimension incorrect!\"\n",
        "\n",
        "            input_ids.append(input_id)\n",
        "            token_type_ids.append(token_type_id)\n",
        "            att_masks.append(att_mask)\n",
        "            labels.append(label)\n",
        "\n",
        "    return input_ids, token_type_ids, att_masks, labels"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mX3dEiY-MzxU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_dataloader(dataset, type, max_seq_len, batch_size):\n",
        "    \"\"\"Creates train and validation DataLoaders with input_ids,\n",
        "    token_type_ids, att_masks, and labels\n",
        "\n",
        "    Returns:\n",
        "        train_dataloader: DataLoader object\n",
        "        validation_dataloader: DataLoader object\n",
        "\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",
        "    input_id, token_type_id, \\\n",
        "    att_mask, label = get_input_data(dataset, max_seq_len)\n",
        "\n",
        "    # Convert all inputs to torch tensors\n",
        "    input_ids = torch.tensor(input_id)\n",
        "    token_type_ids = torch.tensor(token_type_id)\n",
        "    att_masks = torch.tensor(att_mask)\n",
        "    labels = torch.tensor(label)\n",
        "\n",
        "    # Create the DataLoader for our training set.\n",
        "    data = TensorDataset(input_ids, token_type_ids, att_masks, labels)\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": "RhPBExB2bjG-",
        "colab_type": "code",
        "outputId": "fbb183e8-3791-4d4f-9eef-788ee311d034",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 126
        }
      },
      "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 [20:03<00:00,  4.72it/s]\n",
            "100%|██████████| 631/631 [02:13<00:00,  4.74it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "Size of the training DataLoader: 17738\n",
            "Size of the validation DataLoader: 1972\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2zaovfEtm7Mf",
        "colab_type": "text"
      },
      "source": [
        "## **Model**\n",
        "Initiate a pre-trained BERT model with a single linear classification layer on top."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MBGukeI_cCPy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Download pre-trained BERT model\n",
        "# The model was converted from TensorFlow to PyTorch format\n",
        "get_model(config['bert_model_name'])\n",
        "\n",
        "if config['bert_model_name'] == 'bert-base-uncased':\n",
        "    model_path = config['bert_model_name']\n",
        "else:\n",
        "    model_path = \"model/\" + config['bert_model_name']\n",
        "\n",
        "# Load BertForSequenceClassification - pretrained BERT model \n",
        "# with a single linear classification layer on top\n",
        "model = BertForSequenceClassification.from_pretrained(model_path, cache_dir=None, num_labels=2)\n",
        "\n",
        "model.to(device)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sPQfVrBHnaIO",
        "colab_type": "text"
      },
      "source": [
        "## **Training and Validation Methods**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "diWEgRaLGoyb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_accuracy(preds, labels):\n",
        "    \"\"\"Compute the accuracy of binary predictions.\n",
        "\n",
        "    Returns:\n",
        "        accuracy: float\n",
        "    -----------------\n",
        "    Arguments:\n",
        "        preds: Numpy list with two columns of probabilities for each label\n",
        "        labels: List of labels\n",
        "    \"\"\"\n",
        "    # Get the label (column) with the higher probability\n",
        "    predictions = np.argmax(preds, axis=1).flatten()\n",
        "    labels = labels.flatten()\n",
        "    # Compute accuracy\n",
        "    accuracy = np.sum(predictions == labels) / len(labels)\n",
        "\n",
        "    return accuracy"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A9kxwJskHG1Y",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(model, train_dataloader, optimizer, scheduler):\n",
        "    \"\"\"Trains the model and returns the average loss and accuracy.\n",
        "\n",
        "    Returns:\n",
        "        avg_loss: Float\n",
        "        avg_acc: Float\n",
        "    ----------\n",
        "    Arguements:\n",
        "        model: Torch model\n",
        "        train_dataloader: DataLoader object\n",
        "        optimizer: Optimizer object\n",
        "        scheduler: Scheduler object\n",
        "    \"\"\"\n",
        "    # Cumulated Training loss and accuracy\n",
        "    total_loss = 0\n",
        "    train_accuracy = 0\n",
        "    # Track the number of steps\n",
        "    num_steps = 0\n",
        "    # Set model in train mode\n",
        "    model.train()\n",
        "    # For each batch of training data\n",
        "    for step, batch in enumerate(tqdm(train_dataloader)):\n",
        "        # Get tensors and move to gpu\n",
        "        # batch contains four PyTorch tensors:\n",
        "        #   [0]: input ids\n",
        "        #   [1]: token_type_ids\n",
        "        #   [2]: attention masks\n",
        "        #   [3]: labels\n",
        "        b_input_ids = batch[0].to(device)\n",
        "        b_token_type_ids = batch[1].to(device)\n",
        "        b_input_mask = batch[2].to(device)\n",
        "        b_labels = batch[3].to(device)\n",
        "\n",
        "        # Zero the gradients\n",
        "        model.zero_grad()\n",
        "        # Forward pass: the model will return the loss and the logits\n",
        "        outputs = model(b_input_ids,\n",
        "                        token_type_ids = b_token_type_ids,\n",
        "                        attention_mask = b_input_mask,\n",
        "                        labels = b_labels)\n",
        "\n",
        "        # Get loss and predictions\n",
        "        loss = outputs[0]\n",
        "        logits = outputs[1]\n",
        "        # Move logits and labels to CPU\n",
        "        logits = logits.detach().cpu().numpy()\n",
        "        label_ids = b_labels.to('cpu').numpy()\n",
        "\n",
        "        # Calculate the accuracy for a batch\n",
        "        tmp_accuracy = get_accuracy(logits, label_ids)\n",
        "\n",
        "        # Accumulate the total accuracy.\n",
        "        train_accuracy += tmp_accuracy\n",
        "\n",
        "        # Track the number of batches\n",
        "        num_steps += 1\n",
        "\n",
        "        # Accumulate the training loss over all of the batches\n",
        "        total_loss += loss.item()\n",
        "\n",
        "        # Perform a backward pass to calculate the gradients\n",
        "        loss.backward()\n",
        "\n",
        "        # Clip the norm of the gradients to 1.0.\n",
        "        # This is to help prevent the \"exploding gradients\" problem\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
        "\n",
        "        # Update parameters and take a step using the computed gradient\n",
        "        optimizer.step()\n",
        "\n",
        "        # Update scheduler\n",
        "        scheduler.step()\n",
        "\n",
        "    # Calculate the average loss over the training data.\n",
        "    avg_loss = total_loss / len(train_dataloader)\n",
        "    avg_acc = train_accuracy/num_steps\n",
        "\n",
        "    return avg_loss, avg_acc"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7iCcUYUvb6-r",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def validate(model, validation_dataloader):\n",
        "    \"\"\"Validates the model and returns the average loss and accuracy.\n",
        "\n",
        "    Returns:\n",
        "        avg_loss: Float\n",
        "        avg_acc: Float\n",
        "    ----------\n",
        "    Arguements:\n",
        "        model: Torch model\n",
        "        validation_dataloader: DataLoader object\n",
        "    \"\"\"\n",
        "    # Set model to evaluation mode\n",
        "    model.eval()\n",
        "    # Cumulated Training loss and accuracy\n",
        "    total_loss = 0\n",
        "    eval_accuracy = 0\n",
        "    # Track the number of steps\n",
        "    num_steps = 0\n",
        "\n",
        "    # For each batch of the validation data\n",
        "    for batch in tqdm(validation_dataloader):\n",
        "        # Move tensors from batch to GPU\n",
        "        batch = tuple(t.to(device) for t in batch)\n",
        "        # Unpack the inputs from the dataloader\n",
        "        b_input_ids, b_token_type_ids, b_input_masks, b_labels = batch\n",
        "        # Don't to compute or store gradients\n",
        "        with torch.no_grad():\n",
        "            outputs = model(b_input_ids,\n",
        "                            token_type_ids = b_token_type_ids,\n",
        "                            attention_mask = b_input_masks,\n",
        "                            labels= b_labels)\n",
        "        # Get loss and logits\n",
        "        loss = outputs[0]\n",
        "        logits = outputs[1]\n",
        "        # Move logits and labels to CPU\n",
        "        logits = logits.detach().cpu().numpy()\n",
        "        label_ids = b_labels.to('cpu').numpy()\n",
        "\n",
        "        # Calculate the accuracy for this batch of test sentences.\n",
        "        tmp_eval_accuracy = get_accuracy(logits, label_ids)\n",
        "\n",
        "        # Accumulate the total accuracy.\n",
        "        eval_accuracy += tmp_eval_accuracy\n",
        "\n",
        "        # Track the number of steps\n",
        "        num_steps += 1\n",
        "\n",
        "        total_loss += loss.item()\n",
        "\n",
        "    # Calculate loss and accuracy\n",
        "    avg_loss = total_loss / len(validation_dataloader)\n",
        "    avg_acc = eval_accuracy/num_steps\n",
        "\n",
        "    return avg_loss, avg_acc"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Z2f40LmnmJy",
        "colab_type": "text"
      },
      "source": [
        "## **Fine-tune FinBERT-QA**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s6lI80kbTo-h",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "optimizer = AdamW(model.parameters(), \n",
        "                  lr = config['learning_rate'], \n",
        "                  weight_decay = config['weight_decay'])\n",
        "\n",
        "n_epochs = config['n_epochs']\n",
        "\n",
        "\n",
        "# Total number of training steps is number of batches * number of epochs.\n",
        "total_steps = len(train_dataloader) * n_epochs\n",
        "\n",
        "# Create the learning rate scheduler.\n",
        "scheduler = get_linear_schedule_with_warmup(optimizer, \n",
        "                                            num_warmup_steps = config['num_warmup_steps'],\n",
        "                                            num_training_steps = total_steps)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rdWsesbTVMt0",
        "colab_type": "code",
        "outputId": "aadbd01d-f953-410c-ed10-8c1cc0af9f6c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 308
        }
      },
      "source": [
        "# Train and validate the model and print the average loss and accuracy.\n",
        "\n",
        "# Lowest validation lost\n",
        "best_valid_loss = float('inf')\n",
        "\n",
        "for epoch in range(n_epochs):\n",
        "    # Evaluate training loss\n",
        "    train_loss, train_acc = train(model, train_dataloader, optimizer, scheduler)\n",
        "    # Evaluate validation loss\n",
        "    valid_loss, valid_acc = 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",
        "        torch.save(model.state_dict(), 'model/' + str(epoch+1)+ '_finbert-qa.pt')\n",
        "\n",
        "    print(\"\\n\\n Epoch {}:\".format(epoch+1))\n",
        "    print(\"\\t Train Loss: {} | Train Accuracy: {}%\".format(round(train_loss, 3), round(train_acc*100, 2)))\n",
        "    print(\"\\t Validation Loss: {} | Validation Accuracy: {}%\\n\".format(round(valid_loss, 3), round(valid_acc*100, 2)))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 17738/17738 [4:09:59<00:00,  1.18it/s]\n",
            "100%|██████████| 1972/1972 [08:56<00:00,  3.68it/s]\n",
            "  0%|          | 0/17738 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            " Epoch 1:\n",
            "\t Train Loss: 0.091 | Train Accuracy: 97.79%\n",
            "\t Validation Loss: 0.085 | Validation Accuracy: 98.11%\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 17738/17738 [4:09:52<00:00,  1.18it/s]\n",
            "100%|██████████| 1972/1972 [08:56<00:00,  3.68it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            " Epoch 2:\n",
            "\t Train Loss: 0.069 | Train Accuracy: 98.4%\n",
            "\t Validation Loss: 0.089 | Validation Accuracy: 98.05%\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J3Qppb1-nvRg",
        "colab_type": "text"
      },
      "source": [
        "## **Evaluate**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6npVgMDNewIW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def predict(model, q_text, cands, max_seq_len):\n",
        "    \"\"\"Re-ranks the candidates answers for each question.\n",
        "\n",
        "    Returns:\n",
        "        ranked_ans: list of re-ranked candidate docids\n",
        "        sorted_scores: list of relevancy scores of the answers\n",
        "    -------------------\n",
        "    Arguments:\n",
        "        model - PyTorch model\n",
        "        q_text - str - query\n",
        "        cands -List of retrieved candidate docids\n",
        "        max_seq_len - int\n",
        "    \"\"\"\n",
        "    # Convert list to numpy array\n",
        "    cands_id = np.array(cands)\n",
        "    # Empty list for the probability scores of relevancy\n",
        "    scores = []\n",
        "    # For each answer in the candidates\n",
        "    for docid in cands:\n",
        "        # Map the docid to text\n",
        "        ans_text = docid_to_text[docid]\n",
        "        # Create inputs for the model\n",
        "        encoded_seq = tokenizer.encode_plus(q_text, ans_text,\n",
        "                                            max_length=max_seq_len,\n",
        "                                            pad_to_max_length=True,\n",
        "                                            return_token_type_ids=True,\n",
        "                                            return_attention_mask = True)\n",
        "\n",
        "        # Numericalized, padded, clipped seq with special tokens\n",
        "        input_ids = torch.tensor([encoded_seq['input_ids']]).to(device)\n",
        "        # Specify question seq and answer seq\n",
        "        token_type_ids = torch.tensor([encoded_seq['token_type_ids']]).to(device)\n",
        "        # Sepecify which position is part of the seq which is padded\n",
        "        att_mask = torch.tensor([encoded_seq['attention_mask']]).to(device)\n",
        "        # Don't calculate gradients\n",
        "        with torch.no_grad():\n",
        "            # Forward pass, calculate logit predictions for each QA pair\n",
        "            outputs = model(input_ids, token_type_ids=token_type_ids, attention_mask=att_mask)\n",
        "        # Get the predictions\n",
        "        logits = outputs[0]\n",
        "        # Apply activation function\n",
        "        pred = softmax(logits, dim=1)\n",
        "        # Move logits and labels to CPU\n",
        "        pred = pred.detach().cpu().numpy()\n",
        "        # Append relevant scores to list (where label = 1)\n",
        "        scores.append(pred[:,1][0])\n",
        "        # Get the indices of the sorted similarity scores\n",
        "        sorted_index = np.argsort(scores)[::-1]\n",
        "        # Get the list of docid from the sorted indices\n",
        "        ranked_ans = list(cands_id[sorted_index])\n",
        "        sorted_scores = list(np.around(sorted(scores, reverse=True),decimals=3))\n",
        "\n",
        "    return ranked_ans, sorted_scores"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xRSvaM8BEvH8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_rank(model, test_set, max_seq_len):\n",
        "    \"\"\"Re-ranks the candidates answers for each question.\n",
        "\n",
        "    Returns:\n",
        "        qid_pred_rank: Dictionary\n",
        "            key - qid\n",
        "            value - list of re-ranked candidates\n",
        "    -------------------\n",
        "    Arguments:\n",
        "        model - PyTorch model\n",
        "        test_set  List of lists\n",
        "        max_seq_len - int\n",
        "    \"\"\"\n",
        "    # Initiate empty dictionary\n",
        "    qid_pred_rank = {}\n",
        "    # Set model to evaluation mode\n",
        "    model.eval()\n",
        "    # For each element in the test set\n",
        "    for i, seq in enumerate(tqdm(test_set)):\n",
        "        # question id, list of rel answers, list of candidates\n",
        "        qid, label, cands = seq[0], seq[1], seq[2]\n",
        "        # Map question id to text\n",
        "        q_text = qid_to_text[qid]\n",
        "\n",
        "        # List of re-ranked docids and the corresponding probabilities\n",
        "        ranked_ans, sorted_scores = predict(model, q_text, cands, max_seq_len)\n",
        "\n",
        "        # Dict - key: qid, value: ranked list of docids\n",
        "        qid_pred_rank[qid] = ranked_ans\n",
        "\n",
        "    return qid_pred_rank"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZCSL8fW1y6Xn",
        "colab_type": "text"
      },
      "source": [
        "**get_trained_model()** downloads a fine-tuned model - the available model_name are:\n",
        "1. finbert-qa: 'bert-qa' fine-tuned on FiQA\n",
        "2. finbert-domain: 'finbert-domain' fine-tuned on FiQA\n",
        "3. finbert-task: 'finberr-task' fine-tuned on FiQA\n",
        "4. bert-pointwise: 'bert-base-uncase' fine-tuned on FiQA using the cross-entropy loss\n",
        "5. bert-pairwise: 'bert-base-uncase' fine-tuned on FiQA using a pairwise loss"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YZ5dr2hKo1X3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model_name = 'finbert-qa'\n",
        "checkpoint = get_trained_model(model_name)\n",
        "\n",
        "trained_model_path = \"model/trained/\" + model_name + \"/\" + checkpoint"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F2dgKDgyow98",
        "colab_type": "code",
        "outputId": "0d5ed1fd-6df8-4c12-ce4a-ac863dc09940",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 163
        }
      },
      "source": [
        "print(\"Evaluating:\\n\")\n",
        "# Load model\n",
        "model.load_state_dict(torch.load(trained_model_path))\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": 15,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Evaluating:\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 333/333 [07:23<00:00,  1.33s/it]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "Average nDCG@10 for 333 queries: 0.482\n",
            "MRR@10 for 333 queries: 0.436\n",
            "Average Precision@1 for 333 queries: 0.366\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nQEhAYG2pBDl",
        "colab_type": "text"
      },
      "source": [
        "## **Answer Re-ranking**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r3OeoTAEl0Ur",
        "colab_type": "text"
      },
      "source": [
        "### **Retriever**\n",
        "Use BM25 from [Pyserini](https://github.com/castorini/pyserini) to retrieve the top-50 answer candidates"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xdaBYwZ0pI0I",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "59572db2-3d1c-42ae-eebd-d2e95cec9ea6"
      },
      "source": [
        "# Lucene index\n",
        "FIQA_INDEX = \"retriever/lucene-index-fiqa\"\n",
        "\n",
        "# Retriever using Pyserini\n",
        "searcher = pysearch.SimpleSearcher(FIQA_INDEX)\n",
        "\n",
        "# Get a sample from the test set\n",
        "seq = test_set[91]\n",
        "qid, label, cands = seq[0], seq[1], seq[2]\n",
        "q_text = qid_to_text[qid]\n",
        "query = q_text\n",
        "print(query)\n",
        "# query = \"Which investments are the best?\"\n",
        "\n",
        "# Retrieve top-50 answer candidates\n",
        "hits = searcher.search(query, k=50)\n",
        "cands = []\n",
        "\n",
        "for i in range(0, len(hits)):\n",
        "    cands.append(int(hits[i].docid))"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Are credit histories/scores international?\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "YcFxKCoNmUbM"
      },
      "source": [
        "### **Re-ranker**\n",
        "Use FinBERT-QA to re-rank the answer candidates"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oS0KmGvwpW15",
        "colab_type": "code",
        "outputId": "45d2c667-f0fe-4952-bca9-404ddc0420f0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 310
        }
      },
      "source": [
        "# Load model\n",
        "model.load_state_dict(torch.load(trained_model_path))\n",
        "model.eval()\n",
        "\n",
        "# Re-rank candidates\n",
        "rank, scores = predict(model, query, cands, config['max_seq_len'])\n",
        "# Print the Top-k answers\n",
        "k = 5\n",
        "\n",
        "print(\"Query:\\n\\t{}\\n\".format(query))\n",
        "print(\"Top-{} Answers: \\n\".format(k))\n",
        "for i in range(0, k):\n",
        "    print(\"{}.\\t{}\\n\".format(i+1, docid_to_text[rank[i]]))"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Query:\n",
            "\tAre credit histories/scores international?\n",
            "\n",
            "Top-5 Answers: \n",
            "\n",
            "1.\tCurrently the credit history are not International but are local. Many countries don't have a concept of credit history yet.   Having said that, if you are moving to US, depending on your history in your country, you can ask the same bank to provide you with a card and then start building history. For example in India I had a card with Citi Bank and when I moved to US for a short period, I was given a card based on my India Card, with equivalent credit in USD. If you are moving often internationally, it would make sense to Bank with a leading bank that provide services in geographies of your interest [Citi, HSBC, etc] and then in a new country approach these institutions to get you some starting credit for you to build a history.\n",
            "\n",
            "2.\tIt's not just that credit history is local; it's that it's a private business run for profit. The \"big three\" credit bureaus in the US are Experian, Equifax and Transunion.  They collect information on debt usage and abuse from various companies in the US, and charge a fee to provide that information (and their judgement of you) to companies interested in offering you further credit.  But there's nothing stopping a company from collecting international credit histories, or specialized credit histories either (for instance, there's a company called ChexSystems which focuses on retail purchase financing (mostly auto) and checking account abuse, while ignoring other types of lending). That being said, I don't know of any companies which currently collect international credit histories.  Perhaps in Europe, with more nations in close geographic proximity, there would be, but not in North America.\n",
            "\n",
            "3.\tFor instance and to give a comparison to the US - in Austria, almost everybody gets a credit card (without a credit history (e.g. a young person) / with a bad credit history & with a good credit history).  The credit history is in the USA much more important than in Austria. In future, the way to assess a credit history will change due to analysis of social networks for instance. This can be considered in addition to traditional scoring procedures. Is your credit history/score like a criminal record? Nope. I mean is it always with you?  Not really cause a criminal record will be retained on a central storage (to state it abstract) and a credit history can be calculated by private companies. Also, are there other ways to get credit cards besides with a bank? That depends on the country. In Austria, yes.\n",
            "\n",
            "4.\tI came to US as an international student several years ago, and I have also experienced the same situation like most of the international students in finding ways to build credit history. Below I list out some possible approaches you may want to consider: I. Get a student job at campus (recommended) I think the best way is to get a student job in university, say a teaching assistant or student helper. In this case, you can be provided with a social security number and start to build your own credit history. II. Get credit card You can also consider to apply for a credit card. There are indeed some financial institutions that can provide credit cards for international students with no or limited credit scores requirement, say Discover and Bank of America. However, it is relatively hard to get approved, simply because hey may put more restriction in other aspects. For example, you may be required to keep sufficient bank balance above several thousand dollars during a period of time, or you should prove that you have relatives with citizenship in US who can provide your financial aid if needed. III. Apply for a loan (recommended) Getting a loan product is another alternative to get out of this difficult situation, but most of people don’t realize that. There are some FinTech start-ups in United States that specifically focus on international students’ loan financing. One representative example is Westbon (Westbon ), an online lending company that specializes in providing car loan for international students with no SSN or credit history. I once used their loan product to finance a Honda Accord, and Westbon reported my loan transaction records to US credit bureau during my repayment process. Later when I officially got my SSN number, I found my credit history has been automatically synchronized and I don’t have to start from all over again. It never be an easy journey for international students to build credit history in United States. What approach you should make really depends on you own situation. I hope the information above can be useful and good luck for your credit journey!\n",
            "\n",
            "5.\tWhen you start living in US, it doesn't actually matter what was your Credit history in another country. Your Credit History in US is tied to your SSN (Social Security Number), which will be awarded once you are in the country legally and apply for it. Getting an SSN also doesn't guarantee you nothing and you have to build your credit history slowly. Opening a Checking or Savings account will not help you in building a credit history. You need to have some type of Credit Account (credit card, car loan, mortgage etc.) linked to your SSN to start building your credit history. When you are new to US, you probably won't find any bank that will give you a Credit Card as you have no Credit history. One alternative is to apply for a secured credit card. A secured credit card is one you get by putting money or paying money to a bank and open a Credit Card against that money, thereby the bank can be secure that they won't lose any money. Once you have that, you can use that to build up your credit history slowly and once you have a good credit history and score, apply for regular Credit Card or apply for a car loan, mortgage etc. When I came to US 8 years ago, my Credit History was nothing, even though I had pretty good balance and credit history back in my country. I applied for secured credit card by paying $500 to a bank ( which got acquired by CapitalOne ), got it approved and used it for everything, for three years. I applied for other cards in the mean time but got rejected every time. Finally got approved for a regular credit card after three years and in one year added a mortgage and car loan, which helped me to get a decent score now. And Yes, a good Credit Score is important and essential for renting an apartment, leasing a car, getting a Credit Card etc.  but normally your employer can always arrange for an apartment given your situation or you need to share apartment with someone else. You can rent a car without and credit score, but need a valid US / International Drivers license and a Credit Card :-) Best option will be to open a secured credit card and start building your credit. When your wife and family arrives, they also will be assigned individual SSN and can start building their credit history themselves. Please keep in mind that Credit Score and Credit History is always individual here...\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gTjIzeaFmfzW"
      },
      "source": [
        "### **Analysis**\n",
        "Analysis of the ranking of the top-10 answer candidate ids by the retriever and the re-ranker. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YdNPwLAvpZGw",
        "colab_type": "code",
        "outputId": "3dcd4be0-1ed5-4396-ad56-7abe2a788685",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 272
        }
      },
      "source": [
        "# Get the relevancy of the answers\n",
        "cand_rel = get_rel(label, cands)\n",
        "pred_rel = get_rel(label, rank)\n",
        "\n",
        "# Print the ranking of the answer ids and their corresponding relevancy\n",
        "print(\"Retriever: \\n\\t Ranking: {}\\n\\n\\t Relevancy: {}\\n\".format(cands[:10], cand_rel[:10]))\n",
        "print(\"Re-ranker: \\n\\t Ranking: {}\\n\\n\\t Probability: {}\\n\\n\\t Relevancy: {}\".format(rank[:10], scores[:10], pred_rel[:10]))\n",
        "# Actual labels for the question\n",
        "print(\"\\nLabel: \\n\\t{}\".format(label))"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Retriever: \n",
            "\t Ranking: [111466, 509739, 166875, 206267, 304578, 267422, 293363, 192641, 336468, 82472]\n",
            "\n",
            "\t Relevancy: [0, 0, 0, 0, 0, 1, 0, 0, 0, 0]\n",
            "\n",
            "Re-ranker: \n",
            "\t Ranking: [346166, 184365, 304578, 509739, 111466, 267422, 245729, 206267, 23016, 124699]\n",
            "\n",
            "\t Probability: [0.959, 0.94, 0.032, 0.003, 0.002, 0.001, 0.001, 0.001, 0.001, 0.001]\n",
            "\n",
            "\t Relevancy: [1, 1, 0, 0, 0, 1, 0, 0, 0, 0]\n",
            "\n",
            "Label: \n",
            "\t[346166, 184365, 267422, 50080]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VhdBEnhgpa_U",
        "colab_type": "code",
        "outputId": "d38ff9f7-a0cb-4bf2-fa19-84aa72e61220",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 310
        }
      },
      "source": [
        "print(\"Question: \\n\\t{}\\n\".format(query))\n",
        "print(\"Answer Re-ranker\\n\")\n",
        "print(\"Answer: \\n\\t{}\\n\".format(docid_to_text[rank[0]]))\n",
        "print(\"Answer Retriever\\n\")\n",
        "print(\"Answer: \\n\\t{}\\n\".format(docid_to_text[cands[0]]))\n",
        "print(\"Label: \\n\\t{}\".format(docid_to_text[label[0]]))"
      ],
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Question: \n",
            "\tAre credit histories/scores international?\n",
            "\n",
            "Answer Re-ranker\n",
            "\n",
            "Answer: \n",
            "\tCurrently the credit history are not International but are local. Many countries don't have a concept of credit history yet.   Having said that, if you are moving to US, depending on your history in your country, you can ask the same bank to provide you with a card and then start building history. For example in India I had a card with Citi Bank and when I moved to US for a short period, I was given a card based on my India Card, with equivalent credit in USD. If you are moving often internationally, it would make sense to Bank with a leading bank that provide services in geographies of your interest [Citi, HSBC, etc] and then in a new country approach these institutions to get you some starting credit for you to build a history.\n",
            "\n",
            "Answer Retriever\n",
            "\n",
            "Answer: \n",
            "\tWhen you start living in US, it doesn't actually matter what was your Credit history in another country. Your Credit History in US is tied to your SSN (Social Security Number), which will be awarded once you are in the country legally and apply for it. Getting an SSN also doesn't guarantee you nothing and you have to build your credit history slowly. Opening a Checking or Savings account will not help you in building a credit history. You need to have some type of Credit Account (credit card, car loan, mortgage etc.) linked to your SSN to start building your credit history. When you are new to US, you probably won't find any bank that will give you a Credit Card as you have no Credit history. One alternative is to apply for a secured credit card. A secured credit card is one you get by putting money or paying money to a bank and open a Credit Card against that money, thereby the bank can be secure that they won't lose any money. Once you have that, you can use that to build up your credit history slowly and once you have a good credit history and score, apply for regular Credit Card or apply for a car loan, mortgage etc. When I came to US 8 years ago, my Credit History was nothing, even though I had pretty good balance and credit history back in my country. I applied for secured credit card by paying $500 to a bank ( which got acquired by CapitalOne ), got it approved and used it for everything, for three years. I applied for other cards in the mean time but got rejected every time. Finally got approved for a regular credit card after three years and in one year added a mortgage and car loan, which helped me to get a decent score now. And Yes, a good Credit Score is important and essential for renting an apartment, leasing a car, getting a Credit Card etc.  but normally your employer can always arrange for an apartment given your situation or you need to share apartment with someone else. You can rent a car without and credit score, but need a valid US / International Drivers license and a Credit Card :-) Best option will be to open a secured credit card and start building your credit. When your wife and family arrives, they also will be assigned individual SSN and can start building their credit history themselves. Please keep in mind that Credit Score and Credit History is always individual here...\n",
            "\n",
            "Label: \n",
            "\tCurrently the credit history are not International but are local. Many countries don't have a concept of credit history yet.   Having said that, if you are moving to US, depending on your history in your country, you can ask the same bank to provide you with a card and then start building history. For example in India I had a card with Citi Bank and when I moved to US for a short period, I was given a card based on my India Card, with equivalent credit in USD. If you are moving often internationally, it would make sense to Bank with a leading bank that provide services in geographies of your interest [Citi, HSBC, etc] and then in a new country approach these institutions to get you some starting credit for you to build a history.\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}