README.md 1.43 KB
Newer Older
Rayyyyy's avatar
Rayyyyy committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
## 预训练
本文档将介绍如何使用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)