Question Answering

Introduction

With the Question Answering, or Reading Comprehension, task, given a question and a passage of content (context) that may contain an answer for the question, the model will predict the span within the text with a start and end position indicating the answer to the question. For datasets like SQuAD 2.0, this model supports cases when the answer is not contained in the content.

For every word in the context of a given question, the model will be trained to predict:

  • The likelihood this word is the start of the span

  • The likelihood this word is the end of the span

The model chooses the start and end words with maximal probabilities. When the content does not contain the answer, we would like the start and end span to be set for the first token.

A pretrained BERT encoder with two span prediction heads is used for the prediction start and the end position of the answer. The span predictions are token classifiers consisting of a single linear layer.

TAO Toolkit provides a sample notebook to outline the end-to-end workflow on how to train a Question Answering model using TAO Toolkit and deploy it in Riva format on NGC resources.

Downloading Sample Spec files

Before proceeding, let’s download sample spec files that we would need for the rest of the subtasks.

tao question_answering download_specs -r /results/question_answering/default_specs/ \
                                      -o /specs/nlp/questions_answering

Data Format

This model expects the dataset in SQuAD format (i.e., a JSON file for each dataset split). The code snippet below shows an example of the training file. Each title has one or multiple paragraph entries, each consisting of the “context” and question-answer entries. Each question-answer entry has:

  • A question

  • A globally unique id

  • The Boolean flag “is_impossible”, which shows whether a question is answerable or not

  • (if the question is answerable) One answer entry containing the text span and its starting character index in the context.

  • (if the question is not answerable) An empty “answers” list

The evaluation files (for validation and testing) follow the above format, except that it can provide more than one answer to the same question. The inference file also follows the above format, except that it does not require the “answers” and “is_impossible” keywords.

The following is an example of the data format (JSON file):

{
    "data": [
        {
            "title": "Super_Bowl_50",
            "paragraphs": [
                {
                    "context": "Super Bowl 50 was an American football game to determine the champion of the National Football League (NFL) for the 2015 season. The American Football Conference (AFC) champion Denver Broncos defeated the National Football Conference (NFC) champion Carolina Panthers 24\u201310 to earn their third Super Bowl title. The game was played on February 7, 2016, at Levi's Stadium in the San Francisco Bay Area at Santa Clara, California. As this was the 50th Super Bowl, the league emphasized the \"golden anniversary\" with various gold-themed initiatives, as well as temporarily suspending the tradition of naming each Super Bowl game with Roman numerals (under which the game would have been known as \"Super Bowl L\"), so that the logo could prominently feature the Arabic numerals 50.",
                    "qas": [
                        {
                            "question": "Where did Super Bowl 50 take place?",
                            "is_impossible": "false",
                            "id": "56be4db0acb8001400a502ee",
                            "answers": [
                                {
                                    "answer_start": "403",
                                    "text": "Santa Clara, California"
                                }
                            ]
                        },
                        {
                            "question": "What was the winning score of the Super Bowl 50?",
                            "is_impossible": "true",
                            "id": "56be4db0acb8001400a502ez",
                            "answers": [
                            ]
                        }
                    ]
                }
            ]
        }
    ]
}

Dataset Conversion

To perform training of the QA model on the SQuAD dataset, you must first download it from here. You can choose either SQuAD version 1.1, which does not contain questions without the answer and has 100,000+ question-answer pairs on 500+ articles–or the newer SQuAD version 2.0, which combines the 100,000 questions from SQuAD 1.1 with over 50,000 unanswerable questions. To do well with SQuAD2.0, a system must not only answer questions when possible, but also determine when no answer is supported by the paragraph and abstain from answering.

After downloading the files, you should have a squad data folder that contains the following four files for training and evaluation:

|--squad
     |-- v1.1/train-v1.1.json
     |-- v1.1/dev-v1.1.json
     |-- v2.0/train-v2.0.json
     |-- v2.0/dev-v2.0.json

Model Training

The following is an example of the config spec for training (train.yaml) file. You can change any of these parameters and pass them to the training command.

trainer:
  max_epochs: 2

  # Name of the .tlt file where trained model will be saved.
  save_to: trained-model.tlt

model:

  dataset:
      do_lower_case: true
      version_2_with_negative: true

  tokenizer:
      tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece
      vocab_file: null # path to vocab file
      tokenizer_model: null # only used if tokenizer is sentencepiece
      special_tokens: null

  language_model:
    pretrained_model_name: bert-base-uncased
    lm_checkpoint: null
    config_file: null # json file, precedence over config
    config: null

  token_classifier:
    num_layers: 1
    dropout: 0.0
    num_classes: 2
    activation: relu
    log_softmax: false
    use_transformer_init: true


training_ds:
  file: ??? # e.g. squad/v1.1/train-v2.0.json
  batch_size: 12 # per GPU
  shuffle: true
  num_samples: -1

validation_ds:
  file: ??? # e.g. squad/v1.1/dev-v2.0.json
  batch_size: 12 # per GPU
  shuffle: false
  num_samples: -1

optim:
  # optimizer arguments
  name: adamw
  lr: 3e-5
  betas: [0.9, 0.999]
  weight_decay: 0.0
  # scheduler config override
  sched:
    name: SquareRootAnnealing
    warmup_steps: null
    warmup_ratio: 0.0
    last_epoch: -1

    # pytorch lightning args
    monitor: val_loss
    reduce_on_plateau: false

Parameter

Data Type

Default

Description

trainer.max_epochs

integer

2

The number of epochs to train

save_to

string

trained-model.tlt

The filename of the trained model

data_dir

string

The path to the data converted to the specified format

model.tokenizer.tokenizer_name

string

Will be filled automatically based on model.language_model.pretrained_model_name

The tokenizer name

model.tokenizer.vocab_file

string

null

The path to tokenizer vocabulary

model.tokenizer.tokenizer_model

string

null

The path to tokenizer model (for sentencepiece tokenizer only)

model.tokenizer.special_tokens

string

null

Special tokens for the tokenizer (if they exist)

model.language_model.pretrained_model_name

string

bert-base-uncased

The pre-trained language model name (choose from bert-base-cased, bert-base-uncased,
megatron-bert-345m-cased and megatron-bert-345m-uncased, distilbert-base-uncased, biomegatron-bert-345m-uncased)

model.language_model.lm_checkpoint

string

null

The path to the pre-trained language model checkpoint

model.language_model.config_file

string

null

The path to the pre-trained language model config file

model.language_model.config

dictionary

null

The config of the pre-trained language model

model.token_classifier.num_layers

integer

1

The number of fully connected layers of the Classifier on top of the Bert model

model.token_classifier.dropout

float

0.0

The dropout ratio of the fully connected layers

model.token_classifier.num_classes

integer

2

The number of Classifiers (two for QA)

model.token_classifier.activation

string

relu

The activation function to use

model.token_classifier.log_softmax

boolean

false

A flag specifying whether to use log soft max

training_ds.file

string

The training file names

training_ds.batch_size

integer

12

The training data batch size

training_ds.shuffle

bool

true

A flag specifying whether to shuffle the training data

training_ds.num_samples

integer

-1

The number of samples to use from the training dataset (use -1 to specify all samples)

validation_ds.file

string

The validation file names

validation_ds.batch_size

integer

12

The validation data batch size

validation_ds.shuffle

bool

false

A flag specifying whether to shuffle the validation data

validation_ds.num_samples

integer

-1

The number of samples to use from the validation dataset (use -1 to specify all samples)

optim.name

string

adamw

The optimizer to use for training

optim.lr

float

2e-5

The learning rate to use for training

optim.weight_decay

float

0.0

The weight decay to use for training

optim.sched.name

string

SquareRootAnnealing

The warmup schedule

optim.sched.warmup_ratio

float

0.0

The warmup ratio

The following is an example of the command for training the model:

!tao question_answering train -e /specs/nlp/question_answering/train.yaml \
                        data_dir=PATH_TO_DATA \
                        trainer.max_epochs=2 \
                        trainer.amp_level="O1" \
                        trainer.precision=16 \
                        -g 1

Note

The first time you are performing training, it will take an extra 5-10 minutes to process the dataset for training. For future training runs, it will use the processed dataset, which is automatically cached in the files in the same directory as the data.

Required Arguments for Training

  • -e: The experiment specification file to set up training.

  • data_dir: The dataset directory

Optional Arguments

  • trainer.max_epochs: The number of training epochs

  • -g: The number of GPUs to use for training

  • trainer.amp_level and trainer.precision: These fields allow you to use 16-bit mixed precision to accelerate training.

Note

You can use other arguments to override fields in the specification file. To do so, use the name of the config parameter with a desired value and pass it as a parameter in the script call (e.g., trainer.val_check_interval=0.25).

Training Procedure

At the start of evaluation, TAO Toolkit will print out a log of the experiment specification, then load and preprocess the trained data. For the SQuAD dataset, it can initially take several minutes to Tokenize the content. For subsequent faster runs, the preprocessed dataset will be cached in the files at the same directory as the original dataset. Then it will also display the detailed model architecture.

As the model starts training, you should see a progress bar per epoch. Since QA datasets like SQuAD are big, it is usually enough to train for two epochs. If you want to better see the training progress, you can add the trainer.val_check_interval parameter to the script with a value less than one (e.g., trainer.val_check_interval=0.25, which specifies four evaluations on the validation dataset per 1 epoch of training).

At the end of training, TAO Toolkit will save the best checkpoint on the validation dataset at the path specified by the experiment spec file before finishing.

TPU available: None, using: 0 TPU cores
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2]
[NeMo W 2021-01-28 14:52:19 exp_manager:299] There was no checkpoint folder at checkpoint_dir :results/checkpoints. Training from scratch.
[NeMo I 2021-01-28 14:52:19 exp_manager:186] Experiments will be logged at results
...
Validating: 100%|███████████████████████████| 1020/1020 [01:00<00:00, 21.60it/s][NeMo I 2021-01-29 10:17:18 qa_model:175] val exact match 50.10528088941295
[NeMo I 2021-01-29 10:17:18 qa_model:176] val f1 50.10528088941295
Epoch 0:  25%|██▎      | 3770/15076 [09:18<27:54,  6.75it/s, loss=1.34, lr=3e-5]
      Epoch 0, global step 2748: val_loss reached 1.19158 (best 1.19158), saving model to...

Model Fine-tuning

The following is an example spec for fine-tuning of the model:

trainer:
  max_epochs: 1

# Name of the .tlt file where finetuned model will be saved.
save_to: finetuned-model.tlt

# Fine-tuning settings: training dataset.
finetuning_ds:
  file: ??? # e.g. squad/v1.1/train-v1.1.json
  num_samples: 500 # DEMO purposes # -1 # number of samples to be considered, -1 means all the dataset

# Fine-tuning settings: validation dataset.
validation_ds:
  file: ??? # e.g. squad/v1.1/dev-v1.1.json
  num_samples: 500 # DEMO purposes # -1 # number of samples to be considered, -1 means all the dataset

# Fine-tuning settings: different optimizer.
optim:
  name: adamw
  lr: 5e-6

Parameter

Data Type

Default

Description

trainer.max_epochs

integer

2

The number of epochs to train

restore_from

string

trained-model.tlt

The path to the pre-trained model

save_to

string

finetuned-model.tlt

The path to save trained model to

finetuning_ds.file

string

The data file for fine tuning

finetuning_ds.num_samples

integer

500

The number of samples to use from the fine-tuning dataset (use -1 to specify all samples)

validation_ds.file

string

The validation data file

validation_ds.num_samples

integer

500

The number of samples to use from the validation dataset (use -1 to specify all samples)

optim.name

string

adam

The optimizer to use for training

optim.lr

float

1e-5

The learning rate to use for training

Use the following command to fine-tune the model:

!tao question_answering finetune \
                        -e /specs/nlp/question_answering/finetune.yaml \
                        -g 1 \
                        data_dir=PATH_TO_DATA

Required Arguments for Fine-tuning

  • -e: The experiment specification file to set up fine-tuning

  • data_dir: The path to the data

Optional Arguments

  • -g: The number of GPUs to be use for evaluation in a multi-GPU scenario (default: 1)

Note

You can use other arguments to override fields in the specification file. To do so, use the name of the config parameter with a desired value and pass it as a parameter in the script call (e.g. trainer.val_check_interval=0.25).

Fine-tuning Procedure

Fine-tuning procedure and logs will look similar to described in the Model Training section, with the addition of the model that is initially loaded from a previously trained checkpoint.

Model Evaluation

The following is an example spec to evaluate the pre-trained model:

# Test settings: dataset.
test_ds:
  file: ??? # e.g. squad/v1.1/dev-v1.1.json
  batch_size: 32
  shuffle: false
  num_samples: 500 # DEMO purposes -1 # number of samples to be considered, -1 means the whole the dataset

Parameter

Data Type

Default

Description

test_ds.file

string

The evaluation data file

test_ds.batch_size

integer

32

The training data batch size

test_ds.shuffle

bool

false

A flag specifying whether to shuffle the training data

test_ds.num_samples

integer

500

The number of samples to use from the training dataset (use -1 to specify all samples)

Use the following command to evaluate the model:

!tao question_answering evaluate \
                        -e /specs/nlp/question_answering/evaluate.yaml \
                        data_dir=PATH_TO_DATA

Required Arguments for Evaluation

  • -e: The experiment specification file to set up evaluation

  • data_dir: The path to the pre-processed data to run evaluation on

Evaluation Procedure

After the previously trained model is initialized, it will run evaluation against the provided test set. With Extractive QA models, when the answer span is returned by the model, accuracy evaluation uses two metrics: The exact match (EM) and F1 score of the returned answer spans compared to the right answers. The overall EM and F1 scores are computed for a model by averaging the individual example scores.

  • Exact match: If the answer span is exactly equal to the correct one, it returns 1; otherwise, it returns 0. When assessing against a negative example (SQuAD 2.0), if the model predicts any text at all, it automatically receives a 0 for that example.

  • F1: The F1 score is a common metric for classification problems and widely used in QA. It is appropriate when we care equally about precision and recall. In this case, it is computed over the individual words in the prediction against those in the True Answer. The number of shared words between the prediction and the truth is the basis of the F1 score: Precision is the ratio of the number of shared words to the total number of words in the prediction, and recall is the ratio of the number of shared words to the total number of words in the ground truth. F1 = 2 *  (precision * recall) / (precision + recall)

Testing: 100%|███████████████████████████████| 383/383 [01:36<00:00,  3.77it/s][NeMo I 2021-01-29 10:26:40 qa_model:175] test exact match 50.11370336056599
[NeMo I 2021-01-29 10:26:40 qa_model:176] test f1 50.11370336056599
Testing: 100%|████████████████████████████████| 383/383 [02:07<00:00,  3.01it/s]
--------------------------------------------------------------------------------
DATALOADER:0 TEST RESULTS
{'test_exact_match': 50.11370336056599,
 'test_f1': 50.11370336056599,
 'test_loss': tensor(1.1229, device='cuda:0')}

Model Inference

# Name of  file containing data used as inputs during the inference.
input_file: ??? # e.g. squad/v1.1/dev-v1.1.json

# Name of output nbest list file to store predictions to
output_nbest_file: nbest.txt

# Name of output file to store predictions to
output_prediction_file: prediction.txt

Parameter

Data Type

Default

Description

input_file

string

The file containing the data used as inputs during the inference

output_nbest_file

string

nbest.txt

The name of the output nbest list file to store predictions in

output_prediction_file

string

prediction.txt

The name of the output file to store predictions in

The following example shows how to run inference:

!tao question_answering infer \
                        -e /specs/nlp/question_answering/infer.yaml \
                        -m trained-model.tlt \

Required Arguments for Inference

  • -e: The experiment specification file to set up inference. This requires the input_batch with the list of examples to run inference on.

  • -m: The path to the pre-trained model checkpoint from which to infer. The value should be a .tlt file.

Inference Procedure

After the trained model is loaded, it will run on the input file, which is in the same format as the file used for training and evaluation. It will create a prediction.txt prediction output file with prediction spans for each question in the input file.

Model Export

The following is an example of the spec file for model export:

# Name of the .tlt EFF archive to be loaded/model to be exported.
restore_from: trained-model.tlt

# Set export format: ONNX | RIVA
export_format: ONNX

# Output EFF archive containing ONNX.
export_to: exported-model.eonnx

Parameter

Data Type

Default

Description

restore_from

string

trained-model.tlt

The path to the pre-trained model

export_format

string

ONNX

The export format (either “ONNX” or “RIVA”)

export_to

string

exported-model.eonnx

The path to the exported model

To export a pre-trained model, run the following:

 ### For export to ONNX
!tao question_answering export \
                        -e /specs/nlp/question_answering/export.yaml \
                        -m finetuned-model.tlt \
                        -k $KEY

Required Arguments for Export

  • -e: The experiment specification file to set up inference. This requires the input_batch with a list of examples to run inference on.

  • -m: The path to the pre-trained model checkpoint from which to infer. The file should have a .tlt extension.

  • -k: The encryption key

Model Deployment

You can use the Riva framework for the deployment of the trained model in the runtime. For more details, refer to the Riva documentation