Jerry Liu • 2023-08-25

使用合成数据微调 RAG 嵌入模型

更新 2023/9/10： 我们已将嵌入微调抽象包含到 LlamaIndex 仓库中，因此此仓库技术上已过时！请查阅我们的核心文档中的嵌入微调指南。

我们创建了一个全面、端到端的指南，向您展示如何微调嵌入模型，以改进在任何非结构化文本语料库上检索增强生成 (RAG) 系统的性能（无需标签！）。

结果是检索评估指标性能提高了 5-10%——我们微调后的 bge 模型在命中率方面几乎达到了 text-embedding-ada-002 的检索性能水平。这使得检索更加准确，从而整体改进了 RAG 系统。

本教程对构建 RAG 系统的任何人都有帮助

如果您是微调新手，没问题！我们提供了分步 Notebook，详细讲解了关键步骤。只需将文件链接替换为您自己的数据，然后运行每个单元格即可。
微调嵌入模型是轻量级的，不需要 GPU。这些 Notebook 是在 M2 Macbook Pro 上测试的。

资源

仓库：https://github.com/run-llama/finetune-embedding
Notebook：数据集生成，微调，评估

背景/上下文

当前的 RAG 技术栈

RAG 是一种流行范式，用于将大型语言模型 (LLM) 与其训练语料库中不存在的外部数据源连接起来。它通过输入提示空间将知识库上的检索模型与 LLM 配对。RAG 技术栈通常如下所示

索引：准备一个非结构化文本语料库，对其进行解析/分块。然后嵌入每个块并放入向量数据库。
查询时：使用 top-k 嵌入相似性查找从向量数据库中检索上下文，并将上下文填充到 LLM 输入空间中。

（当然 RAG 可以比这复杂得多，LlamaIndex 为简单和高级 RAG 都提供了工具）

不幸的是，RAG 很容易通过拼凑不同组件来制作原型，但难以投入生产。简单的技术栈有许多故障模式，问题往往出在糟糕的检索上——如果返回的上下文与查询无关，那么 LLM 的能力就无关紧要了；答案总是会很糟糕。

如何改进检索？

我们可以尝试更复杂的检索算法（例如，混合搜索、重排序）。

然而，我们最近的生产级 RAG 网络研讨会中的一个见解是，嵌入本身可能并未处于对您的数据而言最优的潜在空间中。预训练模型生成的嵌入可能基于预训练目标彼此接近/远离，但可能无法完全与您自己的检索目标对齐。例如，如果您正在构建机器学习 ArXiv 论文的搜索，您可能希望嵌入在语义上与特定的机器学习概念（例如“LLMs”、“NLP”）对齐，而不是与填充词“本文是…”对齐。

微调是解决这个问题的一种方法。随着技术进步以及易于使用的服务的出现，微调的概念在 LLM 领域变得越来越流行。

在本教程中，我们重点介绍微调嵌入模型。 我们将展示如何通过微调嵌入模型来获得更好的检索性能。

挑战/注意事项

当您微调嵌入模型时，您需要训练示例。对于嵌入模型，这通常意味着您需要“正例”和“负例”——即应该彼此接近和彼此远离的文本对。

一个问题是我们事先没有这些正例或负例。给定一个非结构化文本数据集，是否可能自动生成这些示例对？

使用 LlamaIndex 您可以做到！我们使用 LlamaIndex 模块从非结构化文本块中自动生成一组问题。然后将这些（问题，块）对用作训练模型的正例信号（负例则从其他块中随机采样）。

下一节将全面演示我们的所有模块。

演示

总体而言，我们执行以下步骤

生成用于训练和评估的合成数据集（Notebook）
微调一个开源嵌入模型（Notebook）
评估嵌入模型（Notebook）

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

这里的核心思想是，我们可以利用 LLM 生成可以通过给定上下文片段最好地回答的假设问题。这使我们能够以可扩展的方式生成合成的（查询，相关文档）正例对，而无需人工标注。

更具体地说，我们首先将给定文档处理成文本块语料库。我们使用 LlamaIndex 中的 SimpleNodeParser 模块完成此操作

parser = SimpleNodeParser()
nodes = parser.get_nodes_from_documents(docs, show_progress=verbose)
corpus = {
  node.node_id: node.get_content(metadata_mode=MetadataMode.NONE) 
  for node in nodes
}

然后对于每个文本块，我们使用 LLM 生成一些可以通过该文本块中的信息回答的假设问题。示例如下的提示。

prompt_template = prompt_template or """\
  Context information is below.
  
  ---------------------
  {context_str}
  ---------------------
  
  Given the context information and not prior knowledge.
  generate only questions based on the below query.
  
  You are a Teacher/ Professor. Your task is to setup \
  {num_questions_per_chunk} questions for an upcoming \
  quiz/examination. The questions should be diverse in nature \
  across the document. Restrict the questions to the \
  context information provided."
  """

# for a given node, extract questions (do this over all nodes in outer loop)
query = prompt_template.format(context_str=text, num_questions_per_chunk=num_questions_per_chunk)
response = llm.complete(query)

result = str(response).strip().split("\n")
questions = [
    re.sub(r"^\d+[\).\s]", "", question).strip() for question in result
]
questions = [question for question in questions if len(question) &gt; 0]

最后，我们将所有问题和文本块的配对收集起来作为数据集。示例如下的查询、块和映射。


# example query
f331640a-b407-4028-8db8-4b8db691dd34: "What is the market value of Lyft's common stock held by non-affiliates as of June 30, 2021, based on the closing sales price of the Class A common stock on that date?"

# example corpus
d5554f3e-cdaf-41d7-ac49-8f0ffe3f5759:"UNITED STATESSECURITIES AND..."

# example mapping
f331640a-b407-4028-8db8-4b8db691dd34: d5554f3e-cdaf-41d7-ac49-8f0ffe3f5759

微调开源嵌入模型

我们利用 sentencetransformers 的高级模型拟合 API 来非常轻松地设置训练过程。

我们使用 MultipleNegativesRankingLoss 作为训练目标，使用 InformationRetrievalEvaluator 作为训练过程中的评估器。此外，我们使用 Hugging Face 上的 BAAI/bge-small-en 作为基础模型，并训练少量 epoch。

# define model
model_id = "BAAI/bge-small-en"
model = SentenceTransformer(model_id)

...

# define loss
from sentence_transformers import losses
loss = losses.MultipleNegativesRankingLoss(model)

# define evaluator
from sentence_transformers.evaluation import InformationRetrievalEvaluator
# define over validation dataset
...
evaluator = InformationRetrievalEvaluator(queries, corpus, relevant_docs)

# run training
...
model.fit(
    train_objectives=[(loader, loss)],
    epochs=EPOCHS,
    warmup_steps=warmup_steps,
    output_path='exp_finetune',
    show_progress_bar=True,
    evaluator=evaluator, 
    evaluation_steps=50,
)

评估嵌入模型

我们将微调后的模型与基础模型以及 OpenAI 嵌入模型 text-embedding-ada-002 进行比较。

我们使用两个主要指标进行评估

命中率指标： 对于每个（查询，相关文档）对，我们使用该查询检索 top-k 文档。如果结果包含 relevant_doc，则视为一次命中。
来自 sentence_transformers 的 InformationRetrievalEvaluator。它提供了一整套全面的指标，例如余弦相似性准确率、不同 top-k 值下的精确率和召回率。

结果

在命中率指标方面，基础模型在验证数据集上的命中率为 78%，微调后的模型为 84%。text-embedding-ada-002 达到 87%，这意味着我们微调后的模型只差 3%！