RetroMAE

레트롬

레트로마 및 그 너머의 코드베이스.

새로운 것

• 2022 년 10 월, Retromae : 마스크 된 자동 인코더를 통한 사전 훈련 검색 지향 언어 모델은 EMNLP 2022 에 허용됩니다. Bert-Base Scale 조밀 한 리트리버의 Ms Marco와 Beir에 대한 SOTA 공연!
• 2022 년 11 월, Retromae v2 : 사전 훈련 검색 지향 언어 모델을위한 이중 마스크 자동 인코더가 이제 Arxiv에 있습니다. V1에서 또 다른 큰 진전과 Ms Marco와 Beir의 주요 개선! 모델과 코드가 곧 출시 될 예정입니다!

릴리스 모델

우리는 Huggingface Hub에 몇 가지 검문소를 업로드했습니다.

모델	설명	링크
레트롬	Wikipedia와 BookCorpus에서 사전 트라이얼	Shitao/Retromae
retromae_msmarco	MSMarco 구절에 사전 트라이얼	shitao/retromae_msmarco
retromae_msmarco_finetune	MSMarco 통과 데이터에 대한 Retromae_msmarco를 Finetune	shitao/retromae_msmarco_finetune
retromae_msmarco_distill	Cross-Encoder로 KL-Divergence를 최소화하여 MSMarco 통과 데이터의 Retromae_msmarco를 Finetune	shitao/retromae_msmarco_distill
retromae_beir	BEIR의 MSMARCO 통과 데이터에 대한 레트로 마인 (BEIR이 제공 한 공식 네거티브 사용).	shitao/retromae_beir

식별자 문자열을 사용하여 쉽게로드 할 수 있습니다. 예를 들어:

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

최첨단 공연

Retromae는 고밀도 리트리버의 강력한 초기화를 제공 할 수 있습니다. 도메인 인 도메인 데이터로 미세 조정 한 후 해당 시나리오에서 고품질 감독 검색 성능이 발생합니다. 또한 미리 훈련 된 모델의 전송 가능성을 크게 향상시켜 도메인 외 데이터 세트에서 우수한 제로 샷 성능을 제공합니다.

msmarco 구절

• Wikipedia 및 BookCorpus에서 미리 훈련 된 모델 :

모델	MRR@10	리콜@1000
버트	0.346	0.964
레트롬	0.382	0.981

• MSMarco에서 미리 훈련 된 모델 :

모델	MRR@10	리콜@1000
Cocondenser	0.382	0.984
레트롬	0.393	0.985
레트롬 (증류)	0.416	0.988

베어 벤치 마크

모델	AVG NDCG@10 (18 개의 데이터 세트)
버트	0.371
콘덴서	0.407
레트롬	0.452
Retromae v2	0.491

설치

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

개발을 위해 편집 가능한 것으로 설치하십시오.

 pip install -e .

워크 플로

이 repo에는 사전 훈련과 Finetune의 두 가지 기능이 포함되어 있습니다. 먼저, 마스크 언어 모델링 손실로 일반 데이터 세트 (또는 다운 스트림 데이터 세트)에서 레트로를 훈련하십시오. 그런 다음 대조적 인 손실로 다운 스트림 데이터 세트의 레트로 메인을 미세화하십시오. 더 나은 성능을 달성하기 위해 크로스 인코더가 제공 한 점수를 증류하여 레트로 마를 미팅 할 수도 있습니다. 자세한 워크 플로우 예제를 참조하십시오.

프리 트레인

 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}

FINETUNE

 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 

예

• 사전 훈련
  • Wikipedia에서 사전 훈련
  • MSMARCO 통로에서 사전 훈련
• 바이오 코더
  • MSMARCO 통로에 대한 FINETUNE
  • 베어 벤치 마크
• 크로스 코더
  • MSMARCO 통로에 재선거

소환

우리의 일이 도움이된다면 우리를 인용하는 것을 고려하십시오.

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