Update 02/10/2020: Please check this link for an improved version of doc2query
Doc2query: Document Expansion by Query Prediction
This repository contains the code to reproduce our entry to the MSMARCO passage ranking task, which was placed first on April 8th, 2019. The paper describing our implementation is here.
| MSMARCO Passage Re-Ranking Leaderboard (Apr 8th 2019) | Eval MRR@10 | Dev MRR@10 |
|---|---|---|
| 1st Place - BERTter Indexing (this code) | 36.8 | 37.5 |
| 2nd Place - SAN + BERT base | 35.9 | 37.0 |
| 3rd Place - BERT + Small Training | 35.9 | 36.5 |
Installation
We first need to install OpenNMT so we can train a model to predict queries from documents. We clone from 0.8.2 because that was the version we trained our models. However, feel free to use a newer version, but we cannot guarantee that the commands below will work.
git clone --branch 0.8.2 https://github.com/OpenNMT/OpenNMT-py.git
cd OpenNMT-py
pip install -r requirements.txt
cd ..
We also need to install Anserini, so we can index and retrieve the expanded documents.
sudo apt-get install maven
git clone https://github.com/castorini/Anserini.git
cd Anserini
mvn clean package appassembler:assemble
tar xvfz eval/trec_eval.9.0.4.tar.gz -C eval/ && cd eval/trec_eval.9.0.4 && make
cd ../ndeval && make
cd ../../../
MS MARCO
Data Preprocessing
First, we need to download and extract the MS MARCO dataset:
DATA_DIR=./msmarco_data
mkdir ${DATA_DIR}
wget https://msmarco.blob.core.windows.net/msmarcoranking/collectionandqueries.tar.gz -P ${DATA_DIR}
tar -xvf ${DATA_DIR}/collectionandqueries.tar.gz -C ${DATA_DIR}
To confirm, collectionandqueries.tar.gz should have MD5 checksum of fed5aa512935c7b62787cb68ac9597d6.
The scripts below convert the data to a format that can be consumed by OpenNMT training and inference scripts:
python ./convert_msmarco_to_opennmt.py \
--collection_path=${DATA_DIR}/collection.tsv \
--train_queries=${DATA_DIR}/queries.train.tsv \
--train_qrels=${DATA_DIR}/qrels.train.tsv \
--dev_queries=${DATA_DIR}/queries.dev.tsv \
--dev_qrels=${DATA_DIR}/qrels.dev.small.tsv \
--output_folder=${DATA_DIR}/opennmt_format
The output files and their number of lines should be:
$ wc -l ./msmarco_data/opennmt_format/*
8841823 ./msmarco_data/opennmt_format/src-collection.txt
7437 ./msmarco_data/opennmt_format/src-dev.txt
532751 ./msmarco_data/opennmt_format/src-train.txt
7437 ./msmarco_data/opennmt_format/tgt-dev.txt
532751 ./msmarco_data/opennmt_format/tgt-train.txt
The last step is to preprocess train and dev files with the following command:
python ./OpenNMT-py/preprocess.py \
-train_src ${DATA_DIR}/opennmt_format/src-train.txt \
-train_tgt ${DATA_DIR}/opennmt_format/tgt-train.txt \
-valid_src ${DATA_DIR}/opennmt_format/src-dev.txt \
-valid_tgt ${DATA_DIR}/opennmt_format/tgt-dev.txt \
-save_data ${DATA_DIR}/opennmt_format/preprocessed \
-src_seq_length 10000 \
-tgt_seq_length 10000 \
-src_seq_length_trunc 400 \
-tgt_seq_length_trunc 100 \
-dynamic_dict \
-share_vocab \
-src_vocab_size 32000 \
-tgt_vocab_size 32000 \
-shard_size 100000
Training doc2query (i.e. a transformer model)
python -u ./OpenNMT-py/train.py \
-data ${DATA_DIR}/opennmt_format/preprocessed \
-save_model ${DATA_DIR}/doc2query \
-layers 6 \
-rnn_size 512 \
-word_vec_size 512 \
-transformer_ff 2048 \
-heads 8 \
-encoder_type transformer \
-decoder_type transformer \
-position_encoding \
-train_steps 10000 \
-max_generator_batches 2 \
-dropout 0.1 \
-batch_size 4096 \
-batch_type tokens \
-normalization tokens \
-accum_count 2 \
-optim adam \
-adam_beta2 0.998 \
-decay_method noam \
-warmup_steps 8000 \
-learning_rate 2.0 \
-max_grad_norm 0.0 \
-param_init 0.0 \
-param_init_glorot \
-label_smoothing 0.1 \
-valid_steps 5000 \
-save_checkpoint_steps 5000 \
-world_size 4 \
-share_embeddings \
-gpu_ranks 0 1 2 3
The command above starts training a transformer model using four GPUs (you can
train with one GPU by setting gpu_ranks 0 and world_size 1). It should
take approximately 3-6 hours to reach iteration 10,000, which contains the
lowest perplexity on the dev set (~15.2).
We can now evaluate in BLEU points the performance of our model:
python ./OpenNMT-py/translate.py \
-gpu 0 \
-model ${DATA_DIR}/doc2query_step_10000.pt \
-src ${DATA_DIR}/opennmt_format/src-dev.txt \
-tgt ${DATA_DIR}/opennmt_format/tgt-dev.txt \
-output ${DATA_DIR}/opennmt_format/pred-dev.txt \
-replace_unk \
-verbose \
-report_time \
-beam_size 1
perl ./OpenNMT-py/tools/multi-bleu.perl \
${DATA_DIR}/opennmt_format/tgt-dev.txt < ${DATA_DIR}/opennmt_format/pred-dev.txt
The output should be similar to this:
BLEU = 8.82, 35.0/14.3/5.7/2.5 (BP=0.957, ratio=0.958, hyp_len=34050, ref_len=35553)
In case you don't want to train a doc2query model yourself, you can download our trained model here.
Predicting Queries
We use our best model checkpoint (iteration 10,000) to predict 5 queries for each document in the collection:
python ./OpenNMT-py/translate.py \
-gpu 0 \
-model ${DATA_DIR}/doc2query_step_10000.pt \
-src ${DATA_DIR}/opennmt_format/src-collection.txt \
-output ${DATA_DIR}/opennmt_format/pred-collection_beam5.txt \
-batch_size 32 \
-beam_size 5 \
--n_best 5 \
-replace_unk \
-report_time
The step above takes many hours. Alternatively, you can split
src-collection.txt into multiple files, process them in parallel and merge the
output. For example:
# Split in 9 files, each with a 1M docs.
split -l 1000000 --numeric-suffixes ${DATA_DIR}/opennmt_format/src-collection.txt ${DATA_DIR}/opennmt_format/src-collection.txt
...
# Execute 9 translate.py in parallel, one for each split file.
...
# Merge the predictions into a single file.
cat ${DATA_DIR}/opennmt_format/pred-collection_beam5.txt?? > ${DATA_DIR}/opennmt_format/pred-collection_beam5.txt
In any case, you can download the predicted queries here.
Expanding docs
Next, we need to merge the original documents with the predicted queries into Anserini's jsonl files (which have one json object per line):
python ./convert_collection_to_jsonl.py \
--collection_path=${DATA_DIR}/collection.tsv \
--predictions=${DATA_DIR}/opennmt_format/pred-collection_beam5.txt \
--beam_size=5 \
--output_folder=${DATA_DIR}/collection_jsonl
The above script should generate 9 jsonl files in ${DATA_DIR}/collection_jsonl, each with 1M lines/docs (except for the last one, which should have 841,823 lines).
We can now index these docs as a JsonCollection using Anserini:
sh ./Anserini/target/appassembler/bin/IndexCollection -collection JsonCollection \
-generator DefaultLuceneDocumentGenerator -threads 9 -input ${DATA_DIR}/collection_jsonl \
-index ${DATA_DIR}/lucene-index-msmarco -optimize
The output message should be something like this:
2019-04-26 07:49:14,549 INFO [main] index.IndexCollection (IndexCollection.java:647) - Total 8,841,823 documents indexed in 00:06:02
Your speed may vary... with a modern desktop machine with an SSD, indexing takes around a minute.
Retrieving and Evaluating the Dev set
Since queries of the dev set are too many (+100k), it would take a long time to retrieve all of them. To speed this up, we use only the queries that are in the qrels file:
python ./Anserini/src/main/python/msmarco/filter_queries.py --qrels=${DATA_DIR}/qrels.dev.small.tsv \
--queries=${DATA_DIR}/queries.dev.tsv --output_queries=${DATA_DIR}/queries.dev.small.tsv
The output queries file should contain 6980 lines.
$ wc -l ${DATA_DIR}/queries.dev.small.tsv
6980 /scratch/rfn216/msmarco_data//queries.dev.small.tsv
We can now retrieve this smaller set of queries.
cd Anserini
python ./src/main/python/msmarco/retrieve.py --index ${DATA_DIR}/lucene-index-msmarco \
--qid_queries ${DATA_DIR}/queries.dev.small.tsv --output ${DATA_DIR}/run.dev.small.tsv --hits 1000
cd ..
Retrieval speed will vary by machine:
On a modern desktop with an SSD, we can get ~0.04 seconds per query (taking about five minutes).
On a slower machine with mechanical disks, the entire process might take as long as a couple of hours.
The option -hits specifies the of documents per query to be retrieved.
Thus, the output file should have approximately 6980 * 1000 = 6.9M lines.
Finally, we can evaluate the retrieved documents using this the official MS MARCO evaluation script:
python ./src/main/python/msmarco/msmarco_eval.py ${DATA_DIR}/qrels.dev.small.tsv ${DATA_DIR}/run.dev.small.tsv
And the output should be like this:
#####################
MRR @10: 0.22155540774093688
QueriesRanked: 6980
#####################
Note that these results are 0.6 higher than the ones in the paper. This is due to better BM25 tuning (b1=0.8, k=0.6).
In case you want to compare your retrieved docs against ours, you can download our retrieved docs here.
Reranking with BERT
Most of the gains come from re-ranking with BERT the passages retrieved with BM25 + Doc2query. To implement BERT re-ranker, we follow the same procedure described in the BERT for Passage Re-ranking repository.
We first need to convert dev queries and retrieved docs into the TFRecord format that will be consumed by BERT:
python convert_msmarco_to_tfrecord.py \
--output_folder=${DATA_DIR}/bert_tfrecord \
--collection_path=${DATA_DIR}/collection.tsv \
--vocab=vocab.txt \
--queries=${DATA_DIR}/queries.dev.small.tsv \
--run=${DATA_DIR}/run.dev.small.tsv \
--qrels=${DATA_DIR}/qrels.dev.small.tsv
This script above produces the files dataset.tf and query_doc_ids.txt, and
they should be moved to a folder in the Google Cloud Storage. For you convenience, you can download these files here.
We are now ready to use our Google's Colab to re-rank with BERT.
Because we did not see any difference from training BERT with the expanded vs original docs, we simple re-rank dev queries using the same checkpoint from the BERT for Passage Re-ranking repository, that is, no training is required in this step.
The Colab is configured to use TPUs and it should take 5-10 hours to re-rank all 6980 dev set queries. If you use a GPU, expect this step to be 10x longer.
After it finishes, we can download the run file msmarco_predictions_dev.tsv
(which is in the Google Storage folder you specified in OUTPUT_DIR) and evaluate it:
python ./src/main/python/msmarco/msmarco_eval.py ${DATA_DIR}/qrels.dev.small.tsv ${DATA_DIR}/msmarco_predictions_dev.tsv
The output should be like this:
#####################
MRR @10: 0.3763750170555333
QueriesRanked: 6980
#####################
Note that this MRR@10 is slightly higher than our leadearboard entry, probably because the better tuned BM25.
You can download our run file here.
TREC-CAR
Download our doc2query model trained on TREC-CAR here.
How do I cite this work?
@article{nogueira2019document,
title={Document Expansion by Query Prediction},
author={Nogueira, Rodrigo and Yang, Wei and Lin, Jimmy and Cho, Kyunghyun},
journal={arXiv preprint arXiv:1904.08375},
year={2019}
}