RetroMAE

Retromae

Base de código para retromae y más allá.

Qué hay de nuevo

• Octubre de 2022, RetRomae: los modelos de lenguaje orientados a la recuperación previa al entrenamiento a través de un codificador automático enmascarado se aceptan para EMNLP 2022 ; ¡Actuaciones de SOTA en la Sra. Marco y Beir de un Retriever denso a escala Bert-base!
• Noviembre de 2022, RetRomae V2: el codificador automático enmascarado dúplex para los modelos de lenguaje orientados a la recuperación previa al entrenamiento ahora está en ARXIV. ¡Otro gran paso adelante de V1 y mejoras importantes en la Sra. Marco y Beir! ¡Los modelos y el código llegarán pronto!

Modelos lanzados

Hemos subido algunos puntos de control al Hub Huggingface.

Modelo	Descripción	Enlace
Retromae	Prerediado en Wikipedia y Bookcorpus	Shitao/retromae
Retromae_msmarco	Previamente tripulado en el pasaje de Msmarco	Shitao/retromae_msmarco
Retromae_msmarco_finetune	Finetune el retromae_msmarco en los datos de pasaje de msmarco	Shitao/retromae_msmarco_finetune
Retromae_msmarco_distill	Finetune el retromae_msmarco en los datos de pasaje de msmarco minimizando la divergencia KL con el codificador cruzado	Shitao/retromae_msmarco_distill
Retromae_beir	Finetune las retromae en los datos de pasaje de Msmarco para Beir (use los negativos oficiales proporcionados por Beir)	Shitao/retromae_beir

Puede cargarlos fácilmente usando las cadenas de identificadores. Por ejemplo:

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

Actuación de última generación

RETROMAE puede proporcionar una sólida inicialización del denso retriever; Después de ajustar con datos internos, da lugar a un rendimiento de recuperación supervisado de alta calidad en el escenario correspondiente. Además, mejora sustancialmente la transferibilidad del modelo previamente capacitado, lo que ayuda a dar como resultado un rendimiento superior de disparo cero en conjuntos de datos fuera de dominio.

Pasaje de msmarco

• Modelo previamente entrenado en Wikipedia y BookCorpus:

Modelo	Mrr@10	Recuerde@1000
Bert	0.346	0.964
Retromae	0.382	0.981

• Modelo previamente entrenado en Msmarco:

Modelo	Mrr@10	Recuerde@1000
cocondensador	0.382	0.984
Retromae	0.393	0.985
Retromae (destilación)	0.416	0.988

Marca benjera

Modelo	AVG NDCG@10 (18 conjuntos de datos)
Bert	0.371
Condensador	0.407
Retromae	0.452
Retromae v2	0.491

Instalación

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

Para el desarrollo, instale como editable:

 pip install -e .

Flujo de trabajo

Este repositorio incluye dos funciones: pre-tren y Finetune. En primer lugar, capacite a las retromae en el conjunto de datos generales (o un conjunto de datos aguas abajo) con pérdida de modelado de lenguaje de máscaras. Luego fineta las retromasas en el conjunto de datos aguas abajo con pérdida de contraste. Para lograr un mejor rendimiento, también puede finentar las retromasas por destilación los puntajes proporcionados por el entrelazador. Flujo de trabajo detallado, consulte nuestros ejemplos.

Predicha

 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}

Afinar

 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 

Ejemplos

• Pre-entrante
  • Precinente en Wikipedia
  • Presina en el pasaje de Msmarco
• Bi-codificador
  • Finetune en el paso de Msmarco
  • Marca benjera
• Entrante
  • Reranker en Pasaje de Msmarco

Citación

Si encuentra útil nuestro trabajo, considere citarnos:

 @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},
}