EMDR2 is an end-to-end training algorithm developed for the task of open-domain question answering.
Comparison with other approaches
End-to-End Training
Results
To use this repo, we recommend using one of NGC's more recent PyTorch containers.
The image version used in this paper can be pulled with the command docker pull nvcr.io/nvidia/pytorch:20.03-py3.
An installation of the Nvidia container toolkit may also be required.
There are additional dependencies that are needed to be installed. We have provided a Dockerfile for the same under the directory docker.
For instance, to build a new docker image (nvcr.io/nvidia/pytorch:20.10-py3-faiss-compiled) over the base container, please use this command
cd docker sudo docker build -t nvcr.io/nvidia/pytorch:20.10-py3-faiss-compiled .
To run the above image in an interactive mode, please use this command
sudo docker run --ipc=host --gpus all -it --rm -v /mnt/disks:/mnt/disks nvcr.io/nvidia/pytorch:20.10-py3-faiss-compiled bash
, where /mnt/disks is the directory to be mounted.
We've provided pretrained checkpoints and datasets on Dropbox for use to train models for dense retrieval and open-domain QA tasks. This data can be downloaded here:
Required training files
- Wikipedia Evidence Documents from DPR paper
- Indexed evidence documents and titles
- Dataset-specific question-answer pairs
- BERT-large vocabulary file
Required checkpoints and embeddings
- Masked Salient Span (MSS) pre-trained retriver
- Masked Salient Span (MSS) pre-trained reader
- Precomputed Evidence Embedding using MSS retriever: This is a big file with 32 GB size.
We've provided several scripts for training models for both dense retriever and open-domain QA tasks in examples directory.
Please ensure to change the data and checkpoint paths in these scripts.
To replicate the answer generation results on the Natural Questions (NQ) dataset, run the script as
bash examples/openqa/emdr2_nq.sh
Similar scripts are provided for TriviaQA, WebQuestions and also for training dense retriever.
For end-to-end training, we used a single node of 16 A100 GPUs with 40GB GPU memory.
In the codebase, the first set of 8 GPUs are used for model training, the second set of 8 GPUs are used for asynchronous evidence embedding, and all the 16 GPUs are used for online retrieval at every step.
The code can also be run on a node with 8 GPUs by disabling asynchronous evidence embedding computation. However, this can lead to some loss in performance.
| Dataset | Dev EM | Test EM | Checkpoint | Precomputed Evidence Embedding |
|---|---|---|---|---|
| Natural Questions | 50.42 | 52.49 | link | link |
| TriviaQA | 71.13 | 71.43 | link | link |
| WebQuestions | 49.86 | 48.67 | link | link |
To use these checkpoints, please set the variables of CHECKPOINT_PATH and EMBEDDING_PATH to point to the above checkpoint and embedding index, respectively.
Also, add the option of --no-load-optim and remove the options of --emdr2-training --async-indexer --index-reload-interval 500 from the example script, so that it works in inference mode.
As the memory requirement for inference is lower, evaluation can also be performed on 4-8 GPUs.
- Sometimes, we need to save the retriever model for tasks such as top-K recall evaluation. To just save the retriever model from the checkpoints, please use this cmd
python tools/save_emdr2_models.py --submodel-name retriever --load e2eqa/trivia --save e2eqa/trivia/retriever/- To create evidence embeddings from a retriever checkpoint and perform top-K recall evaluation, please use this script. Make sure to correctly set the paths of datasets and checkpoints.
bash examples/helper-scripts/create_wiki_indexes_and_evaluate.shFor any errors or bugs in the codebase, please either open a new issue or send an email to Devendra Singh Sachan ([email protected]) .
If you find this code useful, please consider citing our paper as:
@inproceedings{sachan2021endtoend,
title={End-to-End Training of Multi-Document Reader and Retriever for Open-Domain Question Answering},
author={Devendra Singh Sachan and Siva Reddy and William L. Hamilton and Chris Dyer and Dani Yogatama},
booktitle={Advances in Neural Information Processing Systems},
editor={A. Beygelzimer and Y. Dauphin and P. Liang and J. Wortman Vaughan},
year={2021},
url={https://openreview.net/forum?id=5KWmB6JePx}
}
Some of the ideas and implementations in this work were based on a previous paper. Please also consider citing the following paper, if the code is helpful.
@inproceedings{sachan-etal-2021-end,
title = "End-to-End Training of Neural Retrievers for Open-Domain Question Answering",
author = "Sachan, Devendra and Patwary, Mostofa and Shoeybi, Mohammad and Kant, Neel and Ping, Wei and Hamilton, William L. and Catanzaro, Bryan",
booktitle = "Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)",
month = aug,
year = "2021",
address = "Online",
publisher = "Association for Computational Linguistics",
url = "https://aclanthology.org/2021.acl-long.519",
doi = "10.18653/v1/2021.acl-long.519",
pages = "6648--6662"
}