## 预训练 本文档将介绍如何使用retromae进行预训练,提高检索效果。 ## 环境配置 参考[环境配置](../../README.md#环境配置) ## 数据集 ``` ├── pretrain │ └── toy_pretrain_data.jsonl ``` - `预训练`的训练数据是一个json文件,其中每一行都是这样的字典: ``` {"text": str} ``` 如[toy_pretrain_data.jsonl](./examples/pretrain/toy_pretrain_data.jsonl) ## 训练 使用[retromae](https://github.com/staoxiao/RetroMAE)对模型进行预训练,再用对比学习在大规模成对数据上训练模型。 **预训练的目标是重构文本,预训练后的模型无法直接用于相似度计算,需要进行微调之后才可以用于相似度计算。** 训练脚本如下,encoder模型默认保存地址:`{output_dir}/encoder_model` ```bash torchrun --nproc_per_node {number of gpus} \ -m FlagEmbedding.baai_general_embedding.retromae_pretrain.run \ --output_dir {path to save model} \ --model_name_or_path BAAI/bge-large-en \ --train_data toy_pretrain_data.jsonl \ --learning_rate 2e-5 \ --num_train_epochs 2 \ --per_device_train_batch_size {batch size; set 1 for toy data} \ --dataloader_drop_last True \ --max_seq_length 512 \ --logging_steps 10 \ --dataloader_num_workers 12 ``` 更多参数信息请参考[transformers.TrainingArguments](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments)