1. 案例目标

本案例展示了如何使用LlamaIndex框架构建和训练Cohere自定义重排序器(Reranker)。通过该案例，开发者可以学习如何：

• 准备和构建用于训练重排序器的数据集
• 创建不同类型的训练数据集（无负样本、随机负样本、基于余弦相似度的负样本）
• 使用LlamaIndex的CohereRerankerFinetuneEngine训练自定义重排序器
• 评估和比较不同重排序器的性能
• 将自定义重排序器集成到检索增强生成(RAG)系统中

注意：本案例提供的是Cohere自定义重排序器的指南。教程末尾展示的结果是特定于所选数据集和参数的。建议在决定将其纳入RAG管道之前，使用自己的数据集和各种参数进行实验。

2. 技术栈与核心依赖

核心技术栈

• LlamaIndex: 用于构建LLM应用的数据框架
• Cohere: 提供企业级大型语言模型和重排序器服务的平台
• OpenAI: 用于生成合成数据集
• Python: 编程语言环境

核心依赖包

%pip install llama-index-postprocessor-cohere-rerank %pip install llama-index-llms-openai %pip install llama-index-finetuning %pip install llama-index-embeddings-cohere !pip install llama-index cohere pypdf

3. 环境配置

3.1 获取API密钥

要使用本案例中的功能，需要获取以下API密钥：

• OpenAI API密钥：用于创建合成数据集
• Cohere API密钥：用于训练自定义重排序器和评估基础重排序器

openai_api_key = "YOUR OPENAI API KEY" cohere_api_key = "YOUR COHEREAI API KEY"

3.2 环境变量配置

通过设置环境变量来配置API密钥：

import os os.environ["OPENAI_API_KEY"] = openai_api_key os.environ["COHERE_API_KEY"] = cohere_api_key

4. 案例实现

4.1 数据准备

下载数据

本案例使用Lyft 2021年10K SEC文件进行训练，使用Uber 2021年10K SEC文件进行评估：

!mkdir -p 'data/10k/' !wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/uber_2021.pdf' -O 'data/10k/uber_2021.pdf' !wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/lyft_2021.pdf' -O 'data/10k/lyft_2021.pdf'

加载数据

from llama_index.core import SimpleDirectoryReader lyft_docs = SimpleDirectoryReader( input_files=["./data/10k/lyft_2021.pdf"] ).load_data() uber_docs = SimpleDirectoryReader( input_files=["./data/10k/uber_2021.pdf"] ).load_data()

创建节点

文档提到Query + Relevant Passage/Query + Hard Negatives应少于510个token。为适应这一点，我们将chunk_size限制为400个token（每个块最终将被视为相关段落/硬负样本）：

from llama_index.core.node_parser import SimpleNodeParser # 将块大小限制为400 node_parser = SimpleNodeParser.from_defaults(chunk_size=400) # 创建节点 lyft_nodes = node_parser.get_nodes_from_documents(lyft_docs) uber_nodes = node_parser.get_nodes_from_documents(uber_docs)

4.2 数据集创建

生成问题-上下文对

使用GPT-4从每个块中生成问题：

from llama_index.llms.openai import OpenAI from llama_index.core.evaluation import generate_question_context_pairs llm = OpenAI(temperature=0, model="gpt-4") # 生成问题的提示模板 qa_generate_prompt_tmpl = """\ 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 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. The questions should not contain options, not start with Q1/ Q2. \ Restrict the questions to the context information provided.\ """

创建训练、验证和测试数据集

训练自定义重排序器至少需要256个(Query + Relevant passage)对（带或不带硬负样本）用于训练，64对用于验证。请注意验证是可选的。

• 训练：使用Lyft的前256个节点创建训练对
• 验证：使用Lyft的接下来的64个节点进行验证
• 测试：使用Uber的150个节点

# 训练数据集 qa_dataset_lyft_train = generate_question_context_pairs( lyft_nodes[:256], llm=llm, num_questions_per_chunk=1, qa_generate_prompt_tmpl=qa_generate_prompt_tmpl, ) # 验证数据集 qa_dataset_lyft_val = generate_question_context_pairs( lyft_nodes[257:321], llm=llm, num_questions_per_chunk=1, qa_generate_prompt_tmpl=qa_generate_prompt_tmpl, ) # 测试数据集 qa_dataset_uber_val = generate_question_context_pairs( uber_nodes[:150], llm=llm, num_questions_per_chunk=1, qa_generate_prompt_tmpl=qa_generate_prompt_tmpl, )

数据格式和要求

对于训练和评估，目前接受三元组格式的数据，每行应包含以下内容：

• query：表示问题或目标
• relevant_passages：包含回答查询信息的文档或段落列表。对于每个查询，必须至少有一个relevant_passage
• hard_negatives：表示不包含查询答案的块或段落。请注意硬负样本是可选的，但提供至少约5个硬负样本将导致有意义的改进

创建不同类型的数据集

创建三种数据集：

1. 没有硬负样本的数据集
2. 随机选择硬负样本的数据集
3. 基于余弦相似度选择硬负样本的数据集

from llama_index.finetuning import generate_cohere_reranker_finetuning_dataset from llama_index.embeddings.cohere import CohereEmbedding # 初始化Cohere嵌入模型，用于创建硬负样本 embed_model = CohereEmbedding( api_key=cohere_api_key, model_name="embed-english-v3.0", input_type="search_document", ) # 没有硬负样本的训练和验证数据集 generate_cohere_reranker_finetuning_dataset( qa_dataset_lyft_train, finetune_dataset_file_name="train.jsonl" ) generate_cohere_reranker_finetuning_dataset( qa_dataset_lyft_val, finetune_dataset_file_name="val.jsonl" ) # 随机选择硬负样本的训练和验证数据集 generate_cohere_reranker_finetuning_dataset( qa_dataset_lyft_train, num_negatives=5, hard_negatives_gen_method="random", finetune_dataset_file_name="train_5_random.jsonl", embed_model=embed_model, ) generate_cohere_reranker_finetuning_dataset( qa_dataset_lyft_val, num_negatives=5, hard_negatives_gen_method="random", finetune_dataset_file_name="val_5_random.jsonl", embed_model=embed_model, ) # 基于余弦相似度选择硬负样本的训练和验证数据集 generate_cohere_reranker_finetuning_dataset( qa_dataset_lyft_train, num_negatives=5, hard_negatives_gen_method="cosine_similarity", finetune_dataset_file_name="train_5_cosine_similarity.jsonl", embed_model=embed_model, ) generate_cohere_reranker_finetuning_dataset( qa_dataset_lyft_val, num_negatives=5, hard_negatives_gen_method="cosine_similarity", finetune_dataset_file_name="val_5_cosine_similarity.jsonl", embed_model=embed_model, )

4.3 模型训练

训练自定义重排序器

有了训练和验证数据集后，我们准备继续训练自定义重排序器过程。请注意，此训练预计需要大约25到45分钟。

from llama_index.finetuning import CohereRerankerFinetuneEngine # 没有硬负样本的重排序器模型 finetune_model_no_hard_negatives = CohereRerankerFinetuneEngine( train_file_name="train.jsonl", val_file_name="val.jsonl", model_name="lyft_reranker_0_hard_negatives", model_type="RERANK", base_model="english", ) finetune_model_no_hard_negatives.finetune() # 随机选择5个硬负样本的重排序器模型 finetune_model_random_hard_negatives = CohereRerankerFinetuneEngine( train_file_name="train_5_random.jsonl", val_file_name="val_5_random.jsonl", model_name="lyft_reranker_5_random_hard_negatives", model_type="RERANK", base_model="english", ) finetune_model_random_hard_negatives.finetune() # 基于余弦相似度选择5个硬负样本的重排序器模型 finetune_model_cosine_hard_negatives = CohereRerankerFinetuneEngine( train_file_name="train_5_cosine_similarity.jsonl", val_file_name="val_5_cosine_similarity.jsonl", model_name="lyft_reranker_5_cosine_hard_negatives", model_type="RERANK", base_model="english", ) finetune_model_cosine_hard_negatives.finetune()

提交作业后，您可以在https://dashboard.cohere.com/models的仪表板的"models"部分检查训练状态。然后需要获取模型ID进行测试。

获取训练好的模型

from llama_index.postprocessor.cohere_rerank import CohereRerank reranker_base = CohereRerank(top_n=5) reranker_model_0 = finetune_model_no_hard_negatives.get_finetuned_model( top_n=5 ) reranker_model_5_random = ( finetune_model_random_hard_negatives.get_finetuned_model(top_n=5) ) reranker_model_5_cosine = ( finetune_model_cosine_hard_negatives.get_finetuned_model(top_n=5) )

4.4 模型评估

测试设置

我们将使用Uber的前150个节点进行测试：

1. 不使用重排序器
2. 使用Cohere重排序器（未经任何训练）
3. 使用没有硬负样本的自定义重排序器
4. 使用随机选择硬负样本的自定义重排序器
5. 使用基于余弦相似度选择硬负样本的自定义重排序器

RERANKERS = { "WithoutReranker": "None", "CohereRerank": reranker_base, "CohereRerank_0": reranker_model_0, "CohereRerank_5_random": reranker_model_5_random, "CohereRerank_5_cosine": reranker_model_5_cosine, }

评估函数

def display_results(embedding_name, reranker_name, eval_results): """Display results from evaluate.""" metric_dicts = [] for eval_result in eval_results: metric_dict = eval_result.metric_vals_dict metric_dicts.append(metric_dict) full_df = pd.DataFrame(metric_dicts) hit_rate = full_df["hit_rate"].mean() mrr = full_df["mrr"].mean() metric_df = pd.DataFrame( { "Embedding": [embedding_name], "Reranker": [reranker_name], "hit_rate": [hit_rate], "mrr": [mrr], } ) return metric_df

创建向量索引和检索器

from llama_index.core import VectorStoreIndex from llama_index.core.retrievers import VectorIndexRetriever # 初始化Cohere嵌入模型，索引和检索的input_type不同 index_embed_model = CohereEmbedding( api_key=cohere_api_key, model_name="embed-english-v3.0", input_type="search_document", ) query_embed_model = CohereEmbedding( api_key=cohere_api_key, model_name="embed-english-v3.0", input_type="search_query", ) vector_index = VectorStoreIndex( uber_nodes[:150], embed_model=index_embed_model, ) vector_retriever = VectorIndexRetriever( index=vector_index, similarity_top_k=10, embed_model=query_embed_model, )

运行评估

from llama_index.core.evaluation import RetrieverEvaluator from llama_index.core import QueryBundle from llama_index.core.indices.query.schema import QueryType from llama_index.core.schema import NodeWithScore from llama_index.core.retrievers import BaseRetriever from typing import List results_df = pd.DataFrame() embed_name = "CohereEmbedding" # 循环遍历重排序器 for rerank_name, reranker in RERANKERS.items(): print(f"Running Evaluation for Reranker: {rerank_name}") # 定义检索器 class CustomRetriever(BaseRetriever): """Custom retriever that performs both Vector search and Knowledge Graph search""" def __init__( self, vector_retriever: VectorIndexRetriever, ) -> None: """Init params.""" self._vector_retriever = vector_retriever super().__init__() def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]: """Retrieve nodes given query.""" retrieved_nodes = self._vector_retriever.retrieve(query_bundle) if reranker != "None": retrieved_nodes = reranker.postprocess_nodes( retrieved_nodes, query_bundle ) else: retrieved_nodes = retrieved_nodes[:5] return retrieved_nodes async def _aretrieve( self, query_bundle: QueryBundle ) -> List[NodeWithScore]: """Asynchronously retrieve nodes given query. Implemented by the user. """ return self._retrieve(query_bundle) async def aretrieve( self, str_or_query_bundle: QueryType ) -> List[NodeWithScore]: if isinstance(str_or_query_bundle, str): str_or_query_bundle = QueryBundle(str_or_query_bundle) return await self._aretrieve(str_or_query_bundle) custom_retriever = CustomRetriever(vector_retriever) retriever_evaluator = RetrieverEvaluator.from_metric_names( ["mrr", "hit_rate"], retriever=custom_retriever ) eval_results = await retriever_evaluator.aevaluate_dataset( qa_dataset_uber_val ) current_df = display_results(embed_name, rerank_name, eval_results) results_df = pd.concat([results_df, current_df], ignore_index=True)

查看结果

print(results_df)

5. 案例效果

通过本案例实现的应用，开发者可以获得以下效果：

• 改进的检索准确性：自定义重排序器能够更准确地识别和排序与查询最相关的文档片段
• 领域特定优化：通过在特定领域数据上训练，重排序器能够更好地理解该领域的术语和概念
• 灵活的训练策略：支持使用不同类型的负样本（无负样本、随机负样本、基于余弦相似度的负样本）进行训练
• 性能评估：通过hit_rate和MRR指标量化评估重排序器的性能
• 易于集成：训练好的重排序器可以轻松集成到现有的RAG系统中

Cohere自定义重排序器带来了改进。重要的是要强调，确定最佳硬负样本数量以及是否使用随机或余弦抽样应基于实验结果。本指南提供了一个框架，用于使用自定义Cohere重排序器增强检索系统。

注意：在硬负样本的选择方面有改进的潜力；欢迎社区为此做出贡献。

6. 案例实现思路

核心设计思路

本案例的实现基于以下核心思路：

1. 数据驱动训练：通过构建高质量的训练数据集，使重排序器能够学习特定领域的检索模式
2. 负样本策略：探索不同类型的负样本对模型性能的影响，包括无负样本、随机负样本和基于余弦相似度的负样本
3. 评估导向：通过标准化的评估指标（hit_rate和MRR）量化模型性能，为优化提供明确方向
4. 端到端流程：提供从数据准备、模型训练到评估的完整流程，便于开发者快速上手
5. 企业级集成：利用Cohere的企业级服务和LlamaIndex框架，实现可靠的重排序器训练和部署

实现步骤

步骤1：数据准备

下载并加载Lyft和Uber的10K SEC文件，将其分割为适当大小的节点。

步骤2：数据集创建

使用GPT-4从每个节点生成问题，创建问题-上下文对，并构建不同类型的训练数据集。

步骤3：模型训练

使用CohereRerankerFinetuneEngine训练三个不同的重排序器模型，分别使用不同类型的负样本。

步骤4：模型评估

创建自定义检索器，集成不同的重排序器，并使用标准评估指标比较它们的性能。

7. 扩展建议

可能的扩展方向

• 更多负样本策略：探索其他负样本选择策略，如基于语义相似度、基于TF-IDF等
• 多领域训练：使用多个领域的文档进行训练，提高重排序器的泛化能力
• 主动学习：实现主动学习策略，选择最有价值的样本进行标注和训练
• 多语言支持：扩展对多语言文档的重排序能力
• 实时训练：实现基于用户反馈的在线学习机制，持续优化重排序器性能
• 模型压缩：研究模型压缩技术，在保持性能的同时减少模型大小和推理时间
• 可视化分析：开发可视化工具，帮助理解重排序器的决策过程
• 混合检索：将重排序器与其他检索技术（如知识图谱、关键词检索等）结合，构建更强大的检索系统

8. 总结

本案例全面展示了如何使用LlamaIndex框架构建和训练Cohere自定义重排序器。通过详细的代码示例和实现说明，开发者可以学习到：

• 如何准备和构建用于训练重排序器的数据集
• 如何创建不同类型的训练数据集（无负样本、随机负样本、基于余弦相似度的负样本）
• 如何使用LlamaIndex的CohereRerankerFinetuneEngine训练自定义重排序器
• 如何评估和比较不同重排序器的性能
• 如何将自定义重排序器集成到检索增强生成(RAG)系统中

Cohere自定义重排序器提供了一种有效的方法来提高检索系统的准确性，特别是在特定领域或专业场景中。通过在领域特定数据上训练，重排序器能够更好地理解该领域的术语和概念，从而提供更准确的检索结果。

本案例为开发者提供了一个全面的起点，可以根据具体需求进一步扩展和定制，构建出更加强大和智能的检索系统。通过实验不同的负样本策略和训练参数，开发者可以找到最适合自己应用场景的重排序器配置。