RetroMAE

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

最新动态

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

已发布的模型

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

模型	描述	链接
RetroMAE	在维基百科和图书语料库上预训练	Shitao/RetroMAE
RetroMAE_MSMARCO	在MSMARCO段落上预训练	Shitao/RetroMAE_MSMARCO
RetroMAE_MSMARCO_finetune	在MSMARCO段落数据上微调RetroMAE_MSMARCO	Shitao/RetroMAE_MSMARCO_finetune
RetroMAE_MSMARCO_distill	通过最小化与交叉编码器的KL散度在MSMARCO段落数据上微调RetroMAE_MSMARCO	Shitao/RetroMAE_MSMARCO_distill
RetroMAE_BEIR	在MSMARCO段落数据上为BEIR微调RetroMAE（使用BEIR提供的官方负样本）	Shitao/RetroMAE_BEIR

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

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

最先进的性能

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

MSMARCO段落

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

模型	MRR@10	Recall@1000
Bert	0.346	0.964
RetroMAE	0.382	0.981

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

BEIR 基准测试

模型	平均NDCG@10 (18个数据集)
Bert	0.371
Condenser	0.407
RetroMAE	0.452
RetroMAE v2	0.491

安装

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

示例

引用

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

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