Question_Answering_with_ALBERT.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Question Answering with ALBERT.ipynb",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyPtt7t+plApQ1pDExTO4vM/",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/spark-ming/albert-qa-demo/blob/master/Question_Answering_with_ALBERT.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1qfQAtRsMVl7",
        "colab_type": "text"
      },
      "source": [
        "# Reading Comprehension with ALBERT (and similar)\n",
        "\n",
        "Author: [@techno246](https://twitter.com/techno246)\n",
        "\n",
        "Github Repo: https://github.com/spark-ming/albert-qa-demo/\n",
        "\n",
        "Blog Post: https://www.spark64.com/post/machine-comprehension\n",
        "\n",
        "\n",
        "## Introduction\n",
        "\n",
        "Reading comprehension, otherwise known as question answering systems, are one of the tasks that NLP tries to solve. The goal of this task is to be able to answer an arbitary question given a context. For instance, given the following context:\n",
        "\n",
        "> New Zealand (Māori: Aotearoa) is a sovereign island country in the southwestern Pacific Ocean. It has a total land area of 268,000 square kilometres (103,500 sq mi), and a population of 4.9 million. New Zealand's capital city is Wellington, and its most populous city is Auckland.\n",
        "\n",
        "We ask the question\n",
        "\n",
        "> How many people live in New Zealand?\n",
        "\n",
        "We expect the QA system is to respond with something like this:\n",
        "\n",
        "> 4.9 million\n",
        "\n",
        "Since 2017, transformer models have shown to outperform existing approaches for this task. Many pretrained transformer models exist, including BERT, GPT-2, XLNET. One of the newcomers to the group is ALBERT (A Lite BERT) which was published in September 2019. The research group claims that it outperforms BERT, with much less parameters (shorter training and inference time). \n",
        "\n",
        "This tutorial demonstrates how you can fine-tune ALBERT for the task of QnA and use it for inference. For this tutorial, we will use the transformer library built by [Hugging Face](https://huggingface.co/), which is an extremely nice implementation of the transformer models (including ALBERT) in both TensorFlow and PyTorch. You can  just use a fine-tuned model from their [model repository](https://huggingface.co/models) (which I encourage in general to save money and reduce emissions). However for educational purposes I will also show you how to finetune it yourself so you can adapt it for your own data. \n",
        "\n",
        "Note that the goal of this is not to build an optimised, production ready system, but to demonstrate the concept with as little code as possible. Therefore a lot of code will be retrofitted for this purpose. \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBBHbGvQN5vX",
        "colab_type": "text"
      },
      "source": [
        "## 1.0 Setup\n",
        "\n",
        "Let's check out what kind of GPU our friends at Google gave us. This notebook should be configured to give you a P100 😃 (saved in metadata)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "frTeTcy4WdbY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D5RImM3oWbrZ",
        "colab_type": "text"
      },
      "source": [
        "First, we clone the Hugging Face transformer library from Github.\n",
        "\n",
        "\n",
        "Note it's checking out a specific commit only because I've tested this"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QOAoUwBFMQCg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!git clone https://github.com/huggingface/transformers \\\n",
        "&& cd transformers \\\n",
        "&& git checkout a3085020ed0d81d4903c50967687192e3101e770 "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TRZned-8WJrj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install ./transformers\n",
        "!pip install tensorboardX"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UHCuzhPptH0M",
        "colab_type": "text"
      },
      "source": [
        "## 2.0 Train Model\n",
        "\n",
        "This is where we can train our own model. Note you can skip this step if you don't want to wait 1.5 hours!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OaQGsAiWXcnd",
        "colab_type": "text"
      },
      "source": [
        "### 2.1 Get Training and Evaluation Data\n",
        "\n",
        "The SQuAD dataset contains question/answer pairs to for training the ALBERT model for the QA task. \n",
        "\n",
        "Now get the SQuAD V2.0 dataset. `train-v2.0.json` is for training and `dev-v2.0.json` is for evaluation to see how well your model trained.\n",
        "\n",
        "Read more about this dataset here: https://rajpurkar.github.io/SQuAD-explorer/"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dI6e-PfOXSnO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir dataset \\\n",
        "&& cd dataset \\\n",
        "&& wget https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v2.0.json \\\n",
        "&& wget https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v2.0.json"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dZ87q93GDeeL",
        "colab_type": "text"
      },
      "source": [
        "### 2.2 Run training \n",
        "\n",
        "We can now train the model with the training set. \n",
        "\n",
        "### Notes about parameters:\n",
        "`per_gpu_train_batch_size` specifies the number of training examples per iteration per GPU. *In general*, higher means more accuracy and faster training. However, the biggest limitation is the size of the GPU. 12 is what I use for a GPU with 16GB memory. \n",
        "\n",
        "`save_steps` specifies number of steps before it outputs a checkpoint file. I've increased it to save disk space.\n",
        "\n",
        "`num_train_epochs` I recommend two epochs here. It's currently set to one for the purpose of time\n",
        "\n",
        "`version_2_with_negative` is required for SQuAD V2.0. If training with V1.1, take out this flag\n",
        "\n",
        "Warning: it takes about 1.5 hours to train an epoch! If you don't want to wait this long, feel free to skip this step and note the comment in the code to use a pretrained model!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-Eg53t3QXZAb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!export SQUAD_DIR=/content/dataset \\\n",
        "&& python transformers/examples/run_squad.py \\\n",
        "  --model_type albert \\\n",
        "  --model_name_or_path albert-base-v2 \\\n",
        "  --do_train \\\n",
        "  --do_eval \\\n",
        "  --do_lower_case \\\n",
        "  --train_file $SQUAD_DIR/train-v2.0.json \\\n",
        "  --predict_file $SQUAD_DIR/dev-v2.0.json \\\n",
        "  --per_gpu_train_batch_size 12 \\\n",
        "  --learning_rate 3e-5 \\\n",
        "  --num_train_epochs 1.0 \\\n",
        "  --max_seq_length 384 \\\n",
        "  --doc_stride 128 \\\n",
        "  --output_dir /content/model_output \\\n",
        "  --save_steps 1000 \\\n",
        "  --threads 4 \\\n",
        "  --version_2_with_negative "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-JCNRkQwUD56",
        "colab_type": "text"
      },
      "source": [
        "## 3.0 Setup prediction code\n",
        "\n",
        "Now we can use the Hugging Face library to make predictions using our newly trained model. Note that a lot of the code is pulled from `run_squad.py` in the Hugging Face repository, with all the training parts removed. This modified code allows to run predictions we pass in directly as strings, rather .json format like the training/test set.\n",
        "\n",
        "NOTE if you decided train your own mode, change the flag `use_own_model` to `True`\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qp0Pq9z9Y4S0",
        "colab_type": "code",
        "cellView": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import torch\n",
        "import time\n",
        "from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n",
        "\n",
        "from transformers import (\n",
        "    AlbertConfig,\n",
        "    AlbertForQuestionAnswering,\n",
        "    AlbertTokenizer,\n",
        "    squad_convert_examples_to_features\n",
        ")\n",
        "\n",
        "from transformers.data.processors.squad import SquadResult, SquadV2Processor, SquadExample\n",
        "\n",
        "from transformers.data.metrics.squad_metrics import compute_predictions_logits\n",
        "\n",
        "# READER NOTE: Set this flag to use own model, or use pretrained model in the Hugging Face repository\n",
        "use_own_model = False\n",
        "\n",
        "if use_own_model:\n",
        "  model_name_or_path = \"/content/model_output\"\n",
        "else:\n",
        "  model_name_or_path = \"ktrapeznikov/albert-xlarge-v2-squad-v2\"\n",
        "\n",
        "output_dir = \"\"\n",
        "\n",
        "# Config\n",
        "n_best_size = 1\n",
        "max_answer_length = 30\n",
        "do_lower_case = True\n",
        "null_score_diff_threshold = 0.0\n",
        "\n",
        "def to_list(tensor):\n",
        "    return tensor.detach().cpu().tolist()\n",
        "\n",
        "# Setup model\n",
        "config_class, model_class, tokenizer_class = (\n",
        "    AlbertConfig, AlbertForQuestionAnswering, AlbertTokenizer)\n",
        "config = config_class.from_pretrained(model_name_or_path)\n",
        "tokenizer = tokenizer_class.from_pretrained(\n",
        "    model_name_or_path, do_lower_case=True)\n",
        "model = model_class.from_pretrained(model_name_or_path, config=config)\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "model.to(device)\n",
        "\n",
        "processor = SquadV2Processor()\n",
        "\n",
        "def run_prediction(question_texts, context_text):\n",
        "    \"\"\"Setup function to compute predictions\"\"\"\n",
        "    examples = []\n",
        "\n",
        "    for i, question_text in enumerate(question_texts):\n",
        "        example = SquadExample(\n",
        "            qas_id=str(i),\n",
        "            question_text=question_text,\n",
        "            context_text=context_text,\n",
        "            answer_text=None,\n",
        "            start_position_character=None,\n",
        "            title=\"Predict\",\n",
        "            is_impossible=False,\n",
        "            answers=None,\n",
        "        )\n",
        "\n",
        "        examples.append(example)\n",
        "\n",
        "    features, dataset = squad_convert_examples_to_features(\n",
        "        examples=examples,\n",
        "        tokenizer=tokenizer,\n",
        "        max_seq_length=384,\n",
        "        doc_stride=128,\n",
        "        max_query_length=64,\n",
        "        is_training=False,\n",
        "        return_dataset=\"pt\",\n",
        "        threads=1,\n",
        "    )\n",
        "\n",
        "    eval_sampler = SequentialSampler(dataset)\n",
        "    eval_dataloader = DataLoader(dataset, sampler=eval_sampler, batch_size=10)\n",
        "\n",
        "    all_results = []\n",
        "\n",
        "    for batch in eval_dataloader:\n",
        "        model.eval()\n",
        "        batch = tuple(t.to(device) for t in batch)\n",
        "\n",
        "        with torch.no_grad():\n",
        "            inputs = {\n",
        "                \"input_ids\": batch[0],\n",
        "                \"attention_mask\": batch[1],\n",
        "                \"token_type_ids\": batch[2],\n",
        "            }\n",
        "\n",
        "            example_indices = batch[3]\n",
        "\n",
        "            outputs = model(**inputs)\n",
        "\n",
        "            for i, example_index in enumerate(example_indices):\n",
        "                eval_feature = features[example_index.item()]\n",
        "                unique_id = int(eval_feature.unique_id)\n",
        "\n",
        "                output = [to_list(output[i]) for output in outputs]\n",
        "\n",
        "                start_logits, end_logits = output\n",
        "                result = SquadResult(unique_id, start_logits, end_logits)\n",
        "                all_results.append(result)\n",
        "\n",
        "    output_prediction_file = \"predictions.json\"\n",
        "    output_nbest_file = \"nbest_predictions.json\"\n",
        "    output_null_log_odds_file = \"null_predictions.json\"\n",
        "\n",
        "    predictions = compute_predictions_logits(\n",
        "        examples,\n",
        "        features,\n",
        "        all_results,\n",
        "        n_best_size,\n",
        "        max_answer_length,\n",
        "        do_lower_case,\n",
        "        output_prediction_file,\n",
        "        output_nbest_file,\n",
        "        output_null_log_odds_file,\n",
        "        False,  # verbose_logging\n",
        "        True,  # version_2_with_negative\n",
        "        null_score_diff_threshold,\n",
        "        tokenizer,\n",
        "    )\n",
        "\n",
        "    return predictions"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nIQOB8vhpcKs",
        "colab_type": "text"
      },
      "source": [
        "## 4.0 Run predictions\n",
        "\n",
        "Now for the fun part... testing out your model on different inputs. Pretty rudimentary example here. But the possibilities are endless with this function."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F-sUrcA5nXTH",
        "colab_type": "code",
        "cellView": "code",
        "colab": {}
      },
      "source": [
        "context = \"New Zealand (Māori: Aotearoa) is a sovereign island country in the southwestern Pacific Ocean. It has a total land area of 268,000 square kilometres (103,500 sq mi), and a population of 4.9 million. New Zealand's capital city is Wellington, and its most populous city is Auckland.\"\n",
        "questions = [\"How many people live in New Zealand?\", \n",
        "             \"What's the largest city?\"]\n",
        "\n",
        "# Run method\n",
        "predictions = run_prediction(questions, context)\n",
        "\n",
        "# Print results\n",
        "for key in predictions.keys():\n",
        "  print(predictions[key])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rkivu8FOqp_8",
        "colab_type": "text"
      },
      "source": [
        "## 5.0 Next Steps\n",
        "\n",
        "In this tutorial, you learnt how to fine-tune an ALBERT model for the task of question answering, using the SQuAD dataset. Then, you learnt how you can make predictions using the model. \n",
        "\n",
        "We retrofitted `compute_predictions_logits` to make the prediction for the purpose of simplicity and minimising dependencies in the tutorial. Take a peak inside that module to see how it works. If you want to serve this as an API, you will want to strip out a lot of the stuff it's doing (such as writing the predictions to a JSON, etc)\n",
        "\n",
        "You can now turn this into an API by serving it using a web framework. I recommend checking out FastAPI, which is what [Albert Learns to Read](https://littlealbert.now.sh) is built on. \n",
        "\n",
        "Feel free to open an issue in the [Github respository](https://github.com/spark-ming/albert-qa-demo/) for this notebook, or tweet me @techno246 if you have any questions! \n",
        "\n"
      ]
    }
  ]
}