当前位置：首页 > news >正文

【finetuning】Cohere自定义重排序器案例分析

news 2026/7/22 0:00:07

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个硬负样本将导致有意义的改进

创建不同类型的数据集

创建三种数据集：

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

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个节点进行测试：

不使用重排序器
使用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_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)