## 预训练
本文档将介绍如何使用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)
