AugTriever: Unsupervised Dense Retrieval by Scalable Data Augmentation

This repository contains the code and models of the paper "AugTriever: Unsupervised Dense Retrieval by Scalable Data Augmentation"

Our code is based on the following repositories:

• SimCSE released with the SimCSE paper.
• Contriever released with the Contriever paper.
• MoCo released with the MoCo paper.
• DPR released with the DPR paper.
• SPIDER released with the SPIDER paper.
• MTEB released with the MTEB paper.
• BEIR released with the BEIR paper.
• SentEval released with the SentEval paper.
• E5 released with the E5 paper.
• rank_bm25

Data

AugQ-Wiki and AugQ-CC can be downloaded from Huggingface Hub.

Checkpoints

Naming corresponds to Table 1 in the paper.

Aug method	Model	MM	BEIR (14 tasks)	Download Link
Hybrid-TQGen+	MoCo	24.6	41.1	[download]
Hybrid-All	MoCo	23.5	39.4	[download]
Hybrid-TQGen	MoCo	23.3	39.4	[download]
Doc-Title	MoCo	21.8	38.7	[download]
QExt-PLM	MoCo	20.6	38.2	[download]
TQGen-Topic	MoCo	21.2	38.9	[download]
TQGen-Title	MoCo	21.8	39.3	[download]
TQGen-AbSum	MoCo	23.2	39.6	[download]
TQGen-ExSum	MoCo	23.0	39.4	[download]
TQGen-Topic	InBatch	20.7	39.0	[download]

Run Training

A few scripts for starting training are placed in the folder examples/traning. For example:

cd $PATH_TO_REPO
sh examples/training/cc.moco.topic50.bs2048.gpu8.sh

Run Evaluation

BEIR

Please refer to BEIR for data download.

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python torch.distributed.launch --nproc_per_node=8 --master_addr=127.0.0.1 --master_port=2255 eval_beir.py --model_name_or_path output_dir/augtriever-release/cc.T03b_title50.moco-2e14.contriever256-special50.bert-base-uncased.avg.dot.q128d256.step100k.bs1024.lr5e5/ --dataset fiqa --metric dot --pooling average --per_gpu_batch_size 128 --beir_data_path data/beir/ --output_dir eval_dir/beir

ODQA

Please refer to Spider for details about QA data download and processing.

export EXP_DIR="output_dir/cc-hybrid.RC20+T0gen80.seed477.moco-2e14.contriever256-special50.bert-base-uncased.avg.dot.q128d256.step200k.bs2048.lr5e5/"
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python torch.distributed.launch --nproc_per_node=8 --master_port=31133 --max_restarts=0 generate_passage_embeddings.py --model_name_or_path $EXP_DIR --output_dir $EXP_DIR/embeddings --passages data/nq/psgs_w100.tsv --per_gpu_batch_size 512
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python eval_qa.py --model_name_or_path facebook/contriever --passages data/nq/psgs_w100.tsv --passages_embeddings "$EXP_DIR/embeddings/*" --qa_file data/nq/qas/*-test.csv,data/nq/qas/entityqs/test/P*.test.json --output_dir $EXP_DIR/qa_output --save_or_load_index

Convert models

Convert to Huggingface BERT model

python convert_checkpoint_to_hf_bert.py --ckpt_path output_dir/cc.T03b_topic.inbatch.contriever256-special.bert-base-uncased.avg.dot.q128d256.step100k.bs1024.lr5e5 --output_dir output_dir/cc.T03b_topic.inbatch.contriever256-special.bert-base-uncased.avg.dot.q128d256.step100k.bs1024.lr5e5/hf_ckpt_bert --model_type shared

Convert to Huggingface DPR model

python convert_checkpoint_to_hf_dpr.py --ckpt_path output_dir/cc-hybrid.RC20+T0gen80.seed477.moco-2e14.contriever256-special50.bert-base-uncased.avg.dot.q128d256.step200k.bs2048.lr5e5 --output_dir output_dir/cc-hybrid.RC20+T0gen80.seed477.moco-2e14.contriever256-special50.bert-base-uncased.avg.dot.q128d256.step200k.bs2048.lr5e5/hf_ckpt_dpr --model_type shared

Export scores

Replace the exp path in gather_score_beir.py/gather_score_qa.py/gather_score_senteval.py and run it. For example

python gather_score_beir.py

License

AugTriever is licensed under the BSD 3-Clause License.

Evaluation codes that are forked from external repositories are placed in subfolders (e.g. src/beir, src/beireval, src/mteb, src/mtebeval, src/qa, src/senteval). Please refer to LICENSE in each subfolder for their Copyright information.

Citation

If you find the AugTriever code or models useful, please cite it by using the following BibTeX entry.

@article{meng2022augtriever,
  title={AugTriever: Unsupervised Dense Retrieval by Scalable Data
Augmentation},
  author={Meng, Rui and Liu, Ye and Yavuz, Semih and Agarwal, Divyansh and Tu, Lifu and Yu, Ning and Zhang, Jianguo and Bhat, Meghana and Zhou, Yingbo},
  journal={arXiv preprint arXiv:2212.08841},
  year={2022}
}