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 pypdf3. 环境配置3.1 获取API密钥要使用本案例中的功能需要获取以下API密钥OpenAI API密钥用于创建合成数据集Cohere API密钥用于训练自定义重排序器和评估基础重排序器openai_api_key YOUR OPENAI API KEY cohere_api_key YOUR COHEREAI API KEY3.2 环境变量配置通过设置环境变量来配置API密钥import os os.environ[OPENAI_API_KEY] openai_api_key os.environ[COHERE_API_KEY] cohere_api_key4. 案例实现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_size400) # 创建节点 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(temperature0, modelgpt-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], llmllm, num_questions_per_chunk1, qa_generate_prompt_tmplqa_generate_prompt_tmpl, ) # 验证数据集 qa_dataset_lyft_val generate_question_context_pairs( lyft_nodes[257:321], llmllm, num_questions_per_chunk1, qa_generate_prompt_tmplqa_generate_prompt_tmpl, ) # 测试数据集 qa_dataset_uber_val generate_question_context_pairs( uber_nodes[:150], llmllm, num_questions_per_chunk1, qa_generate_prompt_tmplqa_generate_prompt_tmpl, )数据格式和要求对于训练和评估目前接受三元组格式的数据每行应包含以下内容query表示问题或目标relevant_passages包含回答查询信息的文档或段落列表。对于每个查询必须至少有一个relevant_passagehard_negatives表示不包含查询答案的块或段落。请注意硬负样本是可选的但提供至少约5个硬负样本将导致有意义的改进创建不同类型的数据集创建三种数据集没有硬负样本的数据集随机选择硬负样本的数据集基于余弦相似度选择硬负样本的数据集from llama_index.finetuning import generate_cohere_reranker_finetuning_dataset from llama_index.embeddings.cohere import CohereEmbedding # 初始化Cohere嵌入模型用于创建硬负样本 embed_model CohereEmbedding( api_keycohere_api_key, model_nameembed-english-v3.0, input_typesearch_document, ) # 没有硬负样本的训练和验证数据集 generate_cohere_reranker_finetuning_dataset( qa_dataset_lyft_train, finetune_dataset_file_nametrain.jsonl ) generate_cohere_reranker_finetuning_dataset( qa_dataset_lyft_val, finetune_dataset_file_nameval.jsonl ) # 随机选择硬负样本的训练和验证数据集 generate_cohere_reranker_finetuning_dataset( qa_dataset_lyft_train, num_negatives5, hard_negatives_gen_methodrandom, finetune_dataset_file_nametrain_5_random.jsonl, embed_modelembed_model, ) generate_cohere_reranker_finetuning_dataset( qa_dataset_lyft_val, num_negatives5, hard_negatives_gen_methodrandom, finetune_dataset_file_nameval_5_random.jsonl, embed_modelembed_model, ) # 基于余弦相似度选择硬负样本的训练和验证数据集 generate_cohere_reranker_finetuning_dataset( qa_dataset_lyft_train, num_negatives5, hard_negatives_gen_methodcosine_similarity, finetune_dataset_file_nametrain_5_cosine_similarity.jsonl, embed_modelembed_model, ) generate_cohere_reranker_finetuning_dataset( qa_dataset_lyft_val, num_negatives5, hard_negatives_gen_methodcosine_similarity, finetune_dataset_file_nameval_5_cosine_similarity.jsonl, embed_modelembed_model, )4.3 模型训练训练自定义重排序器有了训练和验证数据集后我们准备继续训练自定义重排序器过程。请注意此训练预计需要大约25到45分钟。from llama_index.finetuning import CohereRerankerFinetuneEngine # 没有硬负样本的重排序器模型 finetune_model_no_hard_negatives CohereRerankerFinetuneEngine( train_file_nametrain.jsonl, val_file_nameval.jsonl, model_namelyft_reranker_0_hard_negatives, model_typeRERANK, base_modelenglish, ) finetune_model_no_hard_negatives.finetune() # 随机选择5个硬负样本的重排序器模型 finetune_model_random_hard_negatives CohereRerankerFinetuneEngine( train_file_nametrain_5_random.jsonl, val_file_nameval_5_random.jsonl, model_namelyft_reranker_5_random_hard_negatives, model_typeRERANK, base_modelenglish, ) finetune_model_random_hard_negatives.finetune() # 基于余弦相似度选择5个硬负样本的重排序器模型 finetune_model_cosine_hard_negatives CohereRerankerFinetuneEngine( train_file_nametrain_5_cosine_similarity.jsonl, val_file_nameval_5_cosine_similarity.jsonl, model_namelyft_reranker_5_cosine_hard_negatives, model_typeRERANK, base_modelenglish, ) 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_n5) reranker_model_0 finetune_model_no_hard_negatives.get_finetuned_model( top_n5 ) reranker_model_5_random ( finetune_model_random_hard_negatives.get_finetuned_model(top_n5) ) reranker_model_5_cosine ( finetune_model_cosine_hard_negatives.get_finetuned_model(top_n5) )4.4 模型评估测试设置我们将使用Uber的前150个节点进行测试不使用重排序器使用Cohere重排序器未经任何训练使用没有硬负样本的自定义重排序器使用随机选择硬负样本的自定义重排序器使用基于余弦相似度选择硬负样本的自定义重排序器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_keycohere_api_key, model_nameembed-english-v3.0, input_typesearch_document, ) query_embed_model CohereEmbedding( api_keycohere_api_key, model_nameembed-english-v3.0, input_typesearch_query, ) vector_index VectorStoreIndex( uber_nodes[:150], embed_modelindex_embed_model, ) vector_retriever VectorIndexRetriever( indexvector_index, similarity_top_k10, embed_modelquery_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(fRunning 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], retrievercustom_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_indexTrue)查看结果print(results_df)5. 案例效果通过本案例实现的应用开发者可以获得以下效果改进的检索准确性自定义重排序器能够更准确地识别和排序与查询最相关的文档片段领域特定优化通过在特定领域数据上训练重排序器能够更好地理解该领域的术语和概念灵活的训练策略支持使用不同类型的负样本无负样本、随机负样本、基于余弦相似度的负样本进行训练性能评估通过hit_rate和MRR指标量化评估重排序器的性能易于集成训练好的重排序器可以轻松集成到现有的RAG系统中Cohere自定义重排序器带来了改进。重要的是要强调确定最佳硬负样本数量以及是否使用随机或余弦抽样应基于实验结果。本指南提供了一个框架用于使用自定义Cohere重排序器增强检索系统。注意在硬负样本的选择方面有改进的潜力欢迎社区为此做出贡献。6. 案例实现思路核心设计思路本案例的实现基于以下核心思路数据驱动训练通过构建高质量的训练数据集使重排序器能够学习特定领域的检索模式负样本策略探索不同类型的负样本对模型性能的影响包括无负样本、随机负样本和基于余弦相似度的负样本评估导向通过标准化的评估指标hit_rate和MRR量化模型性能为优化提供明确方向端到端流程提供从数据准备、模型训练到评估的完整流程便于开发者快速上手企业级集成利用Cohere的企业级服务和LlamaIndex框架实现可靠的重排序器训练和部署实现步骤步骤1数据准备下载并加载Lyft和Uber的10K SEC文件将其分割为适当大小的节点。步骤2数据集创建使用GPT-4从每个节点生成问题创建问题-上下文对并构建不同类型的训练数据集。步骤3模型训练使用CohereRerankerFinetuneEngine训练三个不同的重排序器模型分别使用不同类型的负样本。步骤4模型评估创建自定义检索器集成不同的重排序器并使用标准评估指标比较它们的性能。7. 扩展建议可能的扩展方向更多负样本策略探索其他负样本选择策略如基于语义相似度、基于TF-IDF等多领域训练使用多个领域的文档进行训练提高重排序器的泛化能力主动学习实现主动学习策略选择最有价值的样本进行标注和训练多语言支持扩展对多语言文档的重排序能力实时训练实现基于用户反馈的在线学习机制持续优化重排序器性能模型压缩研究模型压缩技术在保持性能的同时减少模型大小和推理时间可视化分析开发可视化工具帮助理解重排序器的决策过程混合检索将重排序器与其他检索技术如知识图谱、关键词检索等结合构建更强大的检索系统8. 总结本案例全面展示了如何使用LlamaIndex框架构建和训练Cohere自定义重排序器。通过详细的代码示例和实现说明开发者可以学习到如何准备和构建用于训练重排序器的数据集如何创建不同类型的训练数据集无负样本、随机负样本、基于余弦相似度的负样本如何使用LlamaIndex的CohereRerankerFinetuneEngine训练自定义重排序器如何评估和比较不同重排序器的性能如何将自定义重排序器集成到检索增强生成(RAG)系统中Cohere自定义重排序器提供了一种有效的方法来提高检索系统的准确性特别是在特定领域或专业场景中。通过在领域特定数据上训练重排序器能够更好地理解该领域的术语和概念从而提供更准确的检索结果。本案例为开发者提供了一个全面的起点可以根据具体需求进一步扩展和定制构建出更加强大和智能的检索系统。通过实验不同的负样本策略和训练参数开发者可以找到最适合自己应用场景的重排序器配置。