RetroMAE

Rétromae

Base de code pour les rétromae et au-delà.

Quoi de neuf

• Oct. 2022, Retomae: Les modèles de langage orientés vers la récupération avant la formation via un encodeur auto-encodeur masqué sont acceptés à EMNLP 2022 ; Performances SOTA sur Mme Marco et Beir à partir d'un retriever dense à l'échelle de la base de Bert!
• Nov. 2022, Retomae V2: le cocodeur automatique masqué en duplex pour les modèles de langage orientés vers la récupération pré-formation est maintenant sur arXIV. Un autre grand pas en avant de V1 et des améliorations majeures sur Mme Marco et Beir! Les modèles et le code arrivent bientôt!

Modèles publiés

Nous avons téléchargé quelques points de contrôle sur HuggingFace Hub.

Modèle	Description	Lien
Rétromae	Pré-triané sur le Wikipedia et BookCorpus	Shitao / rétromae
Rétromae_msmarco	Pré-trien sur le passage Msmarco	Shitao / rétromae_msmarco
Rétromae_msmarco_finetune	Finetune le rétromae_msmarco sur les données de passage MSMARCO	Shitao / rétromae_msmarco_finetune
Rétromae_msmarco_distill	Finetune le rétromae_msmarco sur les données de passage MSMARCO en minimisant la divergence KL avec le coder	Shitao / rétromae_msmarco_distill
Rétromae_beir	Finetune les rétromae sur les données de passage MSMARCO pour Beir (utilisez les négatifs officiels fournis par Beir)	Shitao / rétromae_beir

Vous pouvez les charger facilement en utilisant les chaînes d'identifiant. Par exemple:

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

Performance de pointe

Les rétromae peuvent fournir une forte initialisation du retriever dense; Après avoir réglé fin avec des données dans le domaine, il donne lieu à une performance de récupération supervisée de haute qualité dans le scénario correspondant. En outre, il améliore considérablement la transférabilité du modèle pré-formé, ce qui contribue à entraîner des performances supérieures à zéro sur des ensembles de données hors du domaine.

Passage MSMARC

• Modèle pré-formé sur Wikipedia et BookCorpus:

Modèle	Mrr @ 10	Rappel @ 1000
Bert	0,346	0,964
Rétromae	0,382	0,981

• Modèle pré-formé sur MSMARCO:

Modèle	Mrr @ 10	Rappel @ 1000
coco-condeseur	0,382	0,984
Rétromae	0,393	0,985
Rétromae (distillation)	0,416	0,988

Benchers de Beir

Modèle	AVG NDCG @ 10 (18 ensembles de données)
Bert	0,371
Condenseur	0,407
Rétromae	0,452
Rétromae v2	0,491

Installation

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

Pour le développement, installer comme modifiable:

 pip install -e .

Flux de travail

Ce repo comprend deux fonctions: pré-train et finetune. Tout d'abord, entraînez les rétromae sur l'ensemble de données général (ou l'ensemble de données en aval) avec une perte de modélisation du langage du masque. Ensuite, Finetune les rétromae sur l'ensemble de données en aval avec perte contrastive. Pour obtenir une meilleure performance, vous pouvez également Finetune les rétromae par distillation les scores fournis par l'encodeur croisé. Flux de travail détaillé, veuillez vous référer à nos exemples.

Présager

 torchrun --nproc_per_node 8 
  -m pretrain.run 
  --output_dir {path to save ckpt} 
  --data_dir {your data} 
  --do_train True 
  --model_name_or_path bert-base-uncased 
  --pretrain_method {retromae or dupmae}

Affiner

 torchrun --nproc_per_node 8 
-m bi_encoder.run 
--output_dir {path to save ckpt} 
--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 

Exemples

• Prétraitement
  • Pré-train sur wikipedia
  • Pré-train sur le passage MSMARCO
• Encodeur
  • Finetune sur le passage Msmarco
  • Benchers de Beir
• Encodeur
  • Reranker sur le passage MSMARCO

Citation

Si vous trouvez notre travail utile, envisagez de nous citer:

 @inproceedings{RetroMAE,
  title={RetroMAE: Pre-Training Retrieval-oriented Language Models Via Masked Auto-Encoder},
  author={Shitao Xiao, Zheng Liu, Yingxia Shao, Zhao Cao},
  url={https://arxiv.org/abs/2205.12035},
  booktitle ={EMNLP},
  year={2022},
}