使用合成数据对 RAG 的嵌入模型进行微调

更新日期：2023年9月10日：我们已将嵌入模型微调的相关抽象封装到 LlamaIndex 仓库中，因此本仓库在技术上已过时！请参阅核心文档中的嵌入微调指南。

本仓库展示了如何在缺乏标注数据（即查询与相关文档的正样本对）的情况下，对嵌入模型进行微调，从而提升 RAG 的性能。

我们将逐步演示如何利用大语言模型生成合成数据集、微调开源嵌入模型，并最终评估微调后的模型效果。我们以一小规模的金融 PDF 文档数据集为例，证明通过微调嵌入模型可以显著提升检索性能。

环境搭建

首先，请克隆本仓库并安装依赖项。此外，您还需要克隆 llama_index 仓库以获取示例 PDF 文件。

git clone git@github.com:jerryjliu/llama_index
git clone git@github.com:run-llama/finetune-embedding.git
cd finetune-embedding
pip install -r requirements.txt

随后即可运行笔记本文件（例如通过 jupyter lab）。

这些笔记本文件较为轻量，几乎可在任何设备上顺利运行。

运行步骤

1. 执行 generate_dataset.ipynb，生成用于训练和评估的合成数据集。
2. 执行 finetune.ipynb，对预训练的开源嵌入模型进行微调。
3. 执行 evaluate.ipynb，将微调后的模型与预训练的基础嵌入模型以及专有的 OpenAI 嵌入模型进行对比评估。

工作原理

1. 生成用于训练和评估的合成数据集

其核心思想是利用大语言模型生成最能由特定上下文回答的假设性问题。这样，我们便可以在无需人工标注的情况下，以可扩展的方式生成 (query, relevant documents) 的合成正样本对。

具体而言，我们首先将给定的文档处理为一系列文本块组成的语料库。接着，针对每个文本块，使用大语言模型生成若干个能够基于该文本块内容回答的假设性问题。最后，我们将所有问题与文本块的配对收集起来，作为最终的数据集。

2. 微调开源嵌入模型

我们借助 sentence-transformers 提供的高级模型训练 API，轻松搭建训练流程。训练目标采用 MultipleNegativesRankingLoss，并在训练过程中使用 InformationRetrievalEvaluator 作为评估指标。我们以开源模型 “BAAI/bge-small-en” 作为基础模型，仅进行少量轮次的训练。

3. 评估嵌入模型

我们将微调后的模型与基础模型以及 OpenAI 的嵌入模型进行对比。评估指标包括 InformationRetrievalEvaluator 和简单的命中率指标。

finetune-embedding 快速上手指南

本指南介绍如何使用 finetune-embedding 项目，通过合成数据微调开源嵌入模型，以提升检索增强生成（RAG）系统的性能。即使没有人工标注的“查询 - 文档”配对数据，也能利用大语言模型（LLM）自动生成训练集。

重要提示：截至 2023 年 9 月 10 日，该仓库的功能已集成至 LlamaIndex 核心库。生产环境建议优先参考 LlamaIndex 官方微调文档。本指南仅用于理解原理或在本地复现该特定流程。

环境准备

• 操作系统：Linux / macOS / Windows (推荐 Linux)
• Python 版本：3.8+
• 硬件要求：
  • 笔记本轻量级运行：任意现代 CPU 即可。
  • 微调训练：建议使用带有 CUDA 支持的 NVIDIA GPU 以加速训练过程。
• 前置依赖：
  • git
  • pip
  • jupyter lab (用于运行示例 Notebook)
  • LLM API Key (用于生成合成数据，如 OpenAI Key)

安装步骤

1. 克隆仓库 需要同时克隆 llama_index (获取示例 PDF 数据) 和 finetune-embedding 项目。

   git clone git@github.com:jerryjliu/llama_index
   git clone git@github.com:run-llama/finetune-embedding.git
   cd finetune-embedding
   

2. 安装依赖 建议使用国内镜像源加速 Python 包下载。

   # 使用清华源安装依赖
   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 启动开发环境 安装完成后，启动 Jupyter Lab 以运行后续步骤。

   jupyter lab
   

基本使用

项目包含三个核心 Notebook，按顺序执行即可完成从数据生成到评估的全流程。请在 Jupyter Lab 中依次打开并运行以下文件：

第一步：生成合成数据集

运行 generate_dataset.ipynb。

• 功能：利用 LLM 阅读文档片段，自动生成假设性问题，构建“问题 - 相关文档”的正样本对。
• 操作：确保在 Notebook 中配置好你的 LLM API Key，运行所有单元格以生成训练和评估所需的数据集。

第二步：微调嵌入模型

运行 finetune.ipynb。

• 功能：基于 sentence-transformers 库，使用生成的合成数据微调开源模型（默认基座模型为 BAAI/bge-small-en）。
• 原理：采用 MultipleNegativesRankingLoss 作为损失函数进行训练。
• 操作：直接运行单元格，等待训练完成并保存微调后的模型权重。

第三步：评估模型效果

运行 evaluate.ipynb。

• 功能：对比微调后的模型、原始基座模型以及 OpenAI 嵌入模型的性能。
• 指标：使用 InformationRetrievalEvaluator 和命中率（Hit Rate）进行评估。
• 操作：运行单元格查看评估报告，验证微调是否提升了检索效果。