RetroMAE

RetroMAE及其后续发展的代码库。

最新动态

• 🔥 2022年10月，RetroMAE: 通过掩码自编码器预训练面向检索的语言模型被EMNLP 2022接收；在MS MARCO和BEIR上取得了BERT-base规模密集检索器的最优性能！
• 🔥 2022年11月，RetroMAE v2: 用于预训练面向检索语言模型的双工掩码自编码器已发布在ArXiv上。相比v1又有重大进步，在MS MARCO和BEIR上取得了显著改进！模型和代码即将发布！

已发布的模型

我们已将一些检查点上传到Huggingface Hub。

你可以使用标识符字符串轻松加载它们。例如：

from transformers import AutoModel
model = AutoModel.from_pretrained('Shitao/RetroMAE')

最先进的性能

RetroMAE可以为密集检索器提供强大的初始化；在领域内数据上微调后，它可以在相应场景中产生高质量的有监督检索性能。此外，它显著提高了预训练模型的迁移能力，有助于在领域外数据集上获得卓越的零样本性能。

MSMARCO段落

• 在维基百科和图书语料库上预训练的模型：

• 在MSMARCO上预训练的模型： | 模型 | MRR@10 | 召回率@1000 | |-------------------|---|---| | coCondenser | 0.382 | 0.984 | | RetroMAE | 0.393 | 0.985 | | RetroMAE(蒸馏) | 0.416 | 0.988 |

BEIR 基准测试

安装

git clone https://github.com/staoxiao/RetroMAE.git
cd RetroMAE
pip install .

对于开发用途，可以安装为可编辑模式：

pip install -e .

工作流程

本仓库包含两个功能：预训练和微调。首先，在通用数据集(或下游数据集)上使用掩码语言建模损失训练RetroMAE。然后在下游数据集上使用对比损失对RetroMAE进行微调。为了获得更好的性能，您还可以通过蒸馏交叉编码器提供的分数来微调RetroMAE。详细工作流程请参考我们的示例。

预训练

torchrun --nproc_per_node 8 \
  -m pretrain.run \
  --output_dir {保存检查点的路径} \
  --data_dir {您的数据} \
  --do_train True \
  --model_name_or_path bert-base-uncased \
  --pretrain_method {retromae 或 dupmae}

微调

torchrun --nproc_per_node 8 \
-m bi_encoder.run \
--output_dir {保存检查点的路径} \
--model_name_or_path Shitao/RetroMAE \
--do_train  \
--corpus_file ./data/BertTokenizer_data/corpus \
--train_query_file ./data/BertTokenizer_data/train_query \
--train_qrels ./data/BertTokenizer_data/train_qrels.txt \
--neg_file ./data/train_negs.tsv 

示例

• 预训练
  • 在维基百科上预训练
  • 在MSMARCO Passage上预训练
• 双编码器
  • 在MSMARCO Passage上微调
  • BEIR基准测试
• 交叉编码器
  • 在MSMARCO Passage上重排序

引用

如果您觉得我们的工作有帮助，请考虑引用我们：

@inproceedings{RetroMAE,
  title={RetroMAE: 通过掩码自编码器预训练面向检索的语言模型},
  author={肖世涛, 刘政, 邵颖侠, 曹昭},
  url={https://arxiv.org/abs/2205.12035},
  booktitle ={EMNLP},
  year={2022},
}