通过适配嵌入模型提高您的 RAG 上下文召回率 95%。

原文：towardsdatascience.com/improve-your-rag-context-recall-by-40-with-an-adapted-embedding-model-5d4a8f583f32

检索增强生成（RAG）是将 LLM 集成到业务用例中的一种突出技术，允许将专有知识注入到 LLM 中。本文假设您已经了解 RAG，您来这里是为了提高您的 RAG 准确性。

让我们简要回顾一下这个过程。RAG 模型由两个主要步骤组成：检索和生成。在检索步骤中，涉及多个子步骤，包括将上下文文本转换为向量、索引上下文向量、检索用户查询的上下文向量以及重新排序上下文向量。一旦检索到查询的上下文，我们就进入生成阶段。在生成阶段，上下文与提示结合，并发送到 LLM 以生成响应。在发送到 LLM 之前，可能对充满上下文的提示进行缓存和路由步骤以优化效率。

对于每个管道步骤，我们将进行多次实验，以共同提高 RAG 的准确性。您可以参考以下图片，该图片列出了（但不限于）每个步骤中进行的实验。

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/4dc2be8668d5a5a5d64037859acedb29.png

开发者面临的一个主要问题是，在生产环境中部署应用程序时准确性大幅下降。

RAG 在 POC 阶段表现最佳，在生产阶段表现最差。这种挫败感在构建通用人工智能应用的开发者中很常见。

生成阶段已经基本解决，部分得益于提示工程。然而，主要挑战是检索用户查询的正确和完整上下文。这通过一个称为上下文召回率的指标来衡量，该指标考虑了针对给定查询检索的相关上下文数量。检索阶段实验的目标是提高上下文召回率。

嵌入模型适配——一个圣杯

在检索阶段的实验中，通过适配嵌入模型，可以将您的上下文召回率分数显著提高+95%。

在适应嵌入模型之前，让我们先理解其背后的概念。这个想法始于词向量，我们将训练模型理解单词在其周围上下文中的含义（了解更多关于 CBOW 和 Skipgram 的信息）。在词向量之后，嵌入模型是专门设计来捕捉文本之间关系的神经网络。它们超越了词级理解，以把握句子级语义。它使用掩码语言建模目标进行训练，其中特定百分比的输入文本将被掩码以训练嵌入模型，以预测掩码的单词。这种方法使模型能够在使用数十亿个标记进行训练时理解更深层次的语言结构和细微差别，并且生成的嵌入模型能够产生相似句子的相似向量，然后可以使用距离度量（如余弦相似度）进行测量，基于此检索上下文将被优先考虑。

因此，现在我们知道了这些模型训练的目标。它将为以下句子产生相似的嵌入：

句子 1：玫瑰是红色的

句子 2：紫罗兰是蓝色的

它们密切相关，因为这两句话都谈论颜色。

对于 RAG，查询和上下文之间的相似度得分应该更高，以便检索到相关上下文。让我们看一下下面的查询和来自 PubmedQA 数据集的上下文。

查询：肿瘤浸润性免疫细胞特征及其在新辅助化疗后的变化能否预测乳腺癌的响应和预后？

上下文：肿瘤微环境免疫与乳腺癌预后相关。高淋巴细胞浸润已被与对新辅助化疗的响应相关联，但免疫细胞亚群特征在预处理和治疗后残留肿瘤中的贡献仍不清楚。我们通过免疫组化分析了 121 名接受新辅助化疗的同质性治疗的乳腺癌患者的预处理和治疗后肿瘤浸润性免疫细胞（CD3、CD4、CD8、CD20、CD68、Foxp3）。分析了免疫细胞特征，并与响应和生存相关联。我们确定了三种肿瘤浸润性免疫细胞特征，这些特征能够预测对新辅助化疗的病理完全缓解（pCR）（簇 B：58%，与簇 A 和 C：7%相比）。CD4 淋巴细胞的高浸润是 pCR 发生的主要因素，这种关联在六个公共基因组数据集中得到了验证。化疗对淋巴细胞浸润的影响更大，包括 CD4/CD8 比率的逆转，与 pCR 和更好的预后相关。对化疗后残留肿瘤中的免疫浸润分析确定了一个特征（簇 Y），主要由高 CD3 和 CD68 浸润组成，与更差的疾病无进展生存率相关。

查询和上下文看起来相似吗？我们是否按照设计的方式使用嵌入模型？显然不是！

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/95a814d47b04312409e32bbf585e9f94.png

左图由作者提供；右图来源：github.com/UKPLab/sentence-transformers/blob/master/docs/img/SemanticSearch.png，Apache-2.0 许可

我们需要微调嵌入模型的原因是确保查询和相关上下文的表示更接近。为什么不从头开始训练呢？这是因为嵌入模型已经从数十亿个标记的训练中理解了语言结构，这些结构仍然可以加以利用。

微调嵌入模型

为了精炼嵌入模型，我们需要包含与预期用户查询相似的问题和相关的公司文档的数据。我们可以利用语言模型（LLM）根据知识库文档生成查询。使用公司的知识库训练 LLM 相当于提供了一个捷径，因为它允许嵌入模型在训练阶段本身访问上下文。

准备数据集 – 训练和测试：

这里是数据准备步骤：

对于训练集：

1. 使用 LLM 从公司的知识库中挖掘所有可能的问题。

2. 如果你正在分割知识库，确保从所有分割中挖掘问题。

对于测试集：

1. 从知识库中挖掘较少的问题。

2. 如果可用，使用真实用户的问题。

3. 改写训练集中的问题。

4. 结合并改写训练集和测试集中的问题。

我们中的大多数人不从事开发领域范围内的嵌入模型。我们创建的嵌入模型旨在在公司的知识库上表现更好。因此，使用公司的内部数据集来训练嵌入模型并无害处。

对于本文，我们将使用来自 Hugging Face 的"*qiaojin/PubMedQ"*数据集，该数据集包含 pubid、问题和上下文等列。pubid 将被用作问题 ID。

fromdatasetsimportload_dataset med_data=load_dataset("qiaojin/PubMedQA","pqa_artificial",split="train")med_data

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/7a6d6e379a8b1896674b5ac82b0d9a6d.png

pubid 是一个唯一的 ID，它指向行。我们将使用 pubid 作为问题 ID。

为了训练嵌入模型，我们将使用 sentence-transformer 训练器进行训练，但你也可以使用 huggingface 训练器。此外，我们正在使用MultipleNegativeRankingLoss来微调我们的模型，但也可以使用各种损失函数如TripletLoss、ContrastiveLoss等达到相同的效果。但是，对于每种损失函数，所需的数据将不同。例如，对于 tripletloss，你需要（Query，Positive Context，Negative Context）对，而在 MultipleNegativeRankingLoss 中，你只需要（Query，Positive Context）对。除了给定 Query 的正向上下文之外的所有上下文都将被视为负向。

在我们的 PubMedQA 数据集中，每一行的"question"列包含一个问题，"context"列包含该问题的合适上下文列表。因此，我们需要扩展上下文列表列，并在新列中创建具有相应上下文的单独行。

dataset=med_data.remove_columns(['long_answer','final_decision'])df=pd.DataFrame(dataset)df['contexts']=df['context'].apply(lambdax:x['contexts'])# Flatten the context list and repeat the questionexpanded_df=df.explode('contexts')# Optionally: Reset index if neededexpanded_df.reset_index(drop=True,inplace=True)expanded_df=expanded_df[['question','contexts']]splitted_dataset=Dataset.from_pandas(expanded_df).train_test_split(test_size=0.05)expanded_df.head()

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/5ccc66bfd389a5a6aadbf3bf9a72f83b.png

准备数据集 – 评估：

现在，我们有了用于训练和测试的数据集。让我们形成用于评估的数据集。对于评估，我们将使用 LLM 从上下文中挖掘问题，从而可以更真实地了解我们的嵌入模型改进得如何。从 PubMedDataset 中，我们将取前 250 行，将每个上下文列表合并成每行的字符串，并将其发送到 LLM 进行问题挖掘。因此，对于每一行，LLM 可能会输出大约 10 个问题。因此，我们将有大约 2500 个问题上下文对用于评估。

fromopenaiimportOpenAIfromtqdm.autoimporttqdm eval_med_data_seed=med_data.shuffle().take(251)client=OpenAI(api_key="<YOUR_API_KEY>")prompt="""Your task is to mine questions from the given context. Example question is also given to you. You have to create questions and return as pipe separated values(|) <Context> {context} </Context> <Example> {example_question} </Example> """questions=[]forrowintqdm(eval_med_data_seed):question=row["question"]context="nn".join(row["context"]["contexts"])question_count=len(row["context"]["contexts"])message=prompt.format(question_count=question_count,context=context,example_question=question)completion=client.chat.completions.create(model="gpt-4o",messages=[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":message}])questions.append(completion.choices[0].message.content.split("|"))eval_med_data_seed=eval_med_data_seed.add_column("test_questions",questions)df=eval_med_data_seed.to_pandas()eval_data=Dataset.from_pandas(df.explode("test_questions"))eval_data.to_parquet("test_med_data2.parquet")

在我们开始训练之前，我们需要使用上面创建的评估数据集准备评估器。

准备评估器：

sentence transformer 库提供了各种评估器，如EmbeddingSimilarityEvaluator、BinaryClassificationEvaluator和InformationRetrievalEvaluator。对于我们为 RAG 训练嵌入模型的特定用例，InformationRetrievalEvaluator 是最合适的选择。此外，可以添加多个评估器并用于评分。

给定一组查询和大型语料库集，信息检索评估器将为每个查询检索最相似的 top-k 个文档。信息检索评估器将使用各种指标评估模型，如 Recall@k、Precision@k、MRR 和 Accuracy@K，其中 k 将为 1、3、5 和 10。对于 RAG，Recall@K 指标是最重要的，因为它表明检索器可以成功检索多少相关上下文。这是至关重要的，因为如果检索器可以检索正确的上下文，即使我们有额外的无关上下文，生成的内容也可能会更准确。

eval_context_id_map={}forrowineval_data:contexts=row["context"]["contexts"]forcontext,context_idinzip(contexts,row["context_ids"]):eval_context_id_map[context_id]=context eval_corpus={}# Our corpus (cid => document)eval_queries={}# Our queries (qid => question)eval_relevant_docs={}# Query ID to relevant documents (qid => set([relevant_cids])forrowineval_data:pubid=row.get("pubid")eval_queries[pubid]=row.get("test_questions")eval_relevant_docs[pubid]=row.get("context_ids")forcontext_idinrow.get("context_ids"):eval_corpus[context_id]=eval_context_id_map[context_id]

_查询：将每个出版物 ID 映射到其对应的问题。

*语料库：将每个上下文 ID 映射到上下文映射中的内容。

*相关文档：将每个出版物 ID 与一组相关上下文 ID 关联。

在形成所有字典后，我们可以从 sentence_transformer 包中创建一个 InformationRetrievalEvaluator 实例。

ir_evaluator=InformationRetrievalEvaluator(queries=eval_queries,corpus=eval_corpus,relevant_docs=eval_relevant_docs,name="med-eval-test",)

模型训练：

最后，让我们训练我们的模型，使用 sentence-transformer 训练器非常简单。只需设置训练配置参数，例如

1. eval_steps – 指定模型需要评估的频率。

2. save_steps – 指定模型需要保存的频率。

3. num_train_epochs – 训练的 epoch 数量

4. per_device_train_batch_size – 在单个 GPU 的情况下，这是一个批处理大小。

5. save_total_limit – 指定允许保存的最大模型。

6. run_name – 日志将发布在 wandb.ai，因此运行名称是必要的。

然后，我们将我们的参数、训练数据集、测试数据集、损失函数、评估器和模型名称传递给训练器。现在你可以坐下来放松，直到训练完成。

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/6541730427a4c5536fe44b72a6696155.png

放松：你是个好人，亚瑟！

对于我们的训练数据，包括测试数据集和评估数据集的推理时间，训练模型大约需要 3 个小时。

# Load base modelmodel=SentenceTransformer("stsb-distilbert-base")output_dir=f"output/training_mnrl-{datetime.now():%Y-%m-%d_%H-%M-%S}"train_loss=MultipleNegativesRankingLoss(model=model)# Training argumentsargs=SentenceTransformerTrainingArguments(output_dir=output_dir,num_train_epochs=1,per_device_train_batch_size=64,eval_strategy="steps",eval_steps=250,save_steps=250,save_total_limit=2,logging_steps=100,run_name="mnrl")# Train the modeltrainer=SentenceTransformerTrainer(model=model,args=args,train_dataset=splitted_dataset["train"],eval_dataset=splitted_dataset["test"],loss=train_loss,evaluator=ir_evaluator)trainer.train()

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/abf5b0970f3da3936c533e1b7272dbd3.png

笔记本末尾附上的完整结果

结果

为了比较，让我们初始化两个模型实例，一个带有训练过的权重，另一个带有未训练过的权重。

untrained_pubmed_model=SentenceTransformer("stsb-distilbert-base")trained_pubmed_model=SentenceTransformer("/kaggle/input/sentencetransformerpubmedmodel/transformers/default/1/final")

ir_evaluator(untrained_pubmed_model)ir_evaluator(trained_pubmed_model)

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/b46fc0fd6a160851f3a2926e0ba25310.png

结果非常清晰，我们在每个指标上都有惊人的改进。以下是我们在关心的指标上的改进：

• recall@1–78.80 % over untrained model

• recall@3–137.92 % over untrained model

• recall@5–116.36 % over untrained model

• recall@10- 95.09 % over untrained model

分析结果后，很明显，嵌入模型增强了上下文召回，从而显著提高了 RAG 生成的整体准确性。然而，一个缺点是需要监控知识库中文档的添加，并定期重新训练模型。

这可以通过遵循标准的机器学习管道流程来实现，其中我们监控模型是否存在漂移，如果漂移超过某个阈值，则重新启动训练流程。

笔记：www.kaggle.com/code/vigneshboss/embedding-model-training-blog?scriptVersionId=200526579

尝试实现这个想法，并请对性能改进的结果进行评论。

请关注、点赞和分享内容！

除非另有说明，所有图片均为作者所有

参考文献：

1. 领域适应到专有数据适应的想法来源于：GPL：用于密集检索无监督领域适应的生成伪标签

2. RAG 评估 –www.pinecone.io/learn/series/vector-databases-in-production-for-busy-engineers/rag-evaluation/

3. SBERT 训练 –sbert.net/examples/training/ms_marco/cross_encoder_README.html