FlagEmbedding项目中文LLM嵌入模型微调实践

背景介绍

FlagEmbedding项目中的LLM-Embedder是一个强大的嵌入模型框架，但其官方版本BAAI/llm-embedder默认不支持中文。本文将详细介绍如何通过微调方法，使该框架能够有效处理中文文本。

中文支持方案

要让LLM-Embedder支持中文，核心思路是使用一个已经支持中文的基础模型作为起点，然后通过特定数据集进行微调。具体实现步骤如下：

1. 选择中文基础模型：推荐使用BAAI/bge-large-zh-v1.5作为query_encoder，这是一个已经预训练好的中文嵌入模型。

2. 准备训练数据：训练数据应采用JSONL格式，每条记录包含查询文本、正例和负例样本，并标注任务类型。例如：

{"query":"今天星期几","pos":[xxxx],"neg":[xxxx],"task":"icl"}

3. 配置中文指令：需要为中文版本配置特定的指令模板，例如：

"zh_llm_embedder": {
    "instruction": {
        "icl": {
            "query": "为这个句子生成表示以用于检索相关文章: ",
            "key": ""
        }
    }
}

微调实施步骤

执行微调的命令如下：

CUDA_VISIBLE_DEVICES="0,2" nohup torchrun --nproc_per_node=2 run_dense.py \
--output_dir data/outputs/icl \
--train_data llm-embedder:train.jsonl \
--save_steps 6000 \
--query_encoder BAAI/bge-large-zh-v1.5 \
--version zh_llm_embedder \
--inbatch_same_dataset epoch \
--learning_rate 2e-5 \
--num_train_epochs 2 \
--train_group_size 15 \
--per_device_train_batch_size 16 \
--cos_temperature 0.02 \
--dtype fp32 \
--report_to tensorboard \
--data_root data >log.log 2>&1 &

关键参数说明：

• query_encoder：指定中文基础模型
• version：使用配置的中文版本标识
• train_group_size：控制训练时正负样本比例
• cos_temperature：调节相似度计算的温度参数

模型转换与使用

微调完成后，模型会保存在指定目录中，包含以下文件结构：

├── args.json
├── encoder
│   ├── config.json
│   ├── model.safetensors
│   ├── special_tokens_map.json
│   ├── tokenizer_config.json
│   ├── tokenizer.json
│   └── vocab.txt
...

为了使用sentence_transformers库加载模型，需要进行格式转换。转换过程中可能会出现关于pooler层未初始化的警告，这是正常现象，可以安全忽略，因为实际使用中不会用到这些层。

技术要点

1. 跨语言迁移：通过使用中文预训练模型作为基础，可以有效实现知识迁移，使模型具备中文处理能力。

2. 微调策略：采用in-batch负采样技术，通过train_group_size参数控制样本组合，提高模型区分相关和不相关文本的能力。

3. 温度参数：cos_temperature参数影响相似度得分的分布，适当调整可以提高模型区分度。

应用建议

1. 对于不同的中文任务场景，可以调整指令模板以获得更好的效果。

2. 微调后的模型可以广泛应用于中文信息检索、问答系统等场景。

3. 建议在实际应用前，使用领域相关数据对模型进行进一步微调，以获得最佳性能。

通过这种方法，开发者可以有效地将LLM-Embedder框架应用于中文场景，为中文自然语言处理任务提供高质量的文本嵌入能力。