自定义 Cohere 重新排序器¶
这个笔记本提供了使用LlamaIndex抽象构建Cohere自定义重新排序器的教程。完成后,您将能够创建一个自定义的重新排序器,并将其用于增强数据检索。
重要提示: 这个笔记本提供了Cohere自定义重新排序器的指南。本教程末尾呈现的结果是根据所选数据集和参数得出的。我们建议在决定将其纳入您的RAG流程之前,先尝试使用您的数据集和各种参数进行实验。
设置¶
让我们安装必要的包。
%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
openai_api_key = "YOUR OPENAI API KEY"
cohere_api_key = "YOUR COHEREAI API KEY"
import os
os.environ["OPENAI_API_KEY"] = openai_api_key
os.environ["COHERE_API_KEY"] = cohere_api_key
from llama_index.core import VectorStoreIndex, SimpleDirectoryReaderfrom llama_index.core.node_parser import SimpleNodeParser# LLMfrom llama_index.llms.openai import OpenAI# 嵌入from llama_index.embeddings.cohere import CohereEmbedding# 检索器from llama_index.core.retrievers import BaseRetriever, VectorIndexRetriever# 重新排序器from llama_index.core import QueryBundlefrom llama_index.core.indices.query.schema import QueryTypefrom llama_index.core.schema import NodeWithScorefrom llama_index.postprocessor.cohere_rerank import CohereRerankfrom llama_index.core.evaluation import EmbeddingQAFinetuneDatasetfrom llama_index.finetuning import generate_cohere_reranker_finetuning_dataset# 评估器from llama_index.core.evaluation import generate_question_context_pairsfrom llama_index.core.evaluation import RetrieverEvaluator# 微调器from llama_index.finetuning import CohereRerankerFinetuneEnginefrom typing import Listimport pandas as pdimport nest_asyncionest_asyncio.apply()
下载数据¶
我们将使用Lyft 2021年的10K SEC文件进行训练,使用Uber 2021年的10K SEC文件进行评估。
!mkdir -p 'data/10k/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/10k/uber_2021.pdf'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/lyft_2021.pdf' -O 'data/10k/lyft_2021.pdf'
加载数据¶
lyft_docs = SimpleDirectoryReader(
input_files=["./data/10k/lyft_2021.pdf"]
).load_data()
uber_docs = SimpleDirectoryReader(
input_files=["./data/10k/uber_2021.pdf"]
).load_data()
数据整理¶
创建节点。¶
文档提到查询+相关段落/查询+硬负例应该少于510个标记。为了满足这一要求,我们将限制块大小为400个标记。(每个块最终将被视为相关段落/硬负例)
# 将块大小限制为400node_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)
我们将使用GPT-4从文本片段中创建问题。
llm = OpenAI(temperature=0, model="gpt-4")
提示生成每个节点/块的问题。
# 生成问题的提示qa_generate_prompt_tmpl = """\下面是上下文信息。---------------------{context_str}---------------------根据上述查询,考虑到上下文信息和先验知识。你是一名教授。你的任务是为即将到来的测验/考试设置{num_questions_per_chunk} 个问题。问题的性质应该多样化跨文档。问题不应包含选项,也不应以Q1/Q2开头。将问题限制在提供的上下文信息中。\"""
训练自定义重新排名器需要至少256个(查询+相关段落)对,可以包含也可以不包含难负例用于训练,以及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_train.save_json("lyft_train_dataset.json")
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_lyft_val.save_json("lyft_val_dataset.json")
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,)# 保存 [可选]qa_dataset_uber_val.save_json("uber_val_dataset.json")
现在我们已经从每个块中汇编了问题,接下来我们将根据训练自定义重新排序器所需的规范格式化数据。
# 初始化Cohere嵌入模型,我们将用它来创建Hard Negatives。embed_model = CohereEmbedding( cohere_api_key=cohere_api_key, model_name="embed-english-v3.0", input_type="search_document",)
让我们创建3个数据集。
- 不包含难负例的数据集。
- 包含难负例,随机选择的数据集。
- 基于余弦相似度选择难负例的数据集。
# 训练和验证数据集,不包含难例负例。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,)
训练自定义重新排序器。¶
准备好我们的训练和验证数据集后,我们可以开始进行自定义重新排名器的训练过程。请注意,预计这次训练将需要大约25到45分钟。
# 具有0个困难负例的Reranker模型。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个困难负例的Reranker模型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个困难负例的Reranker模型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 的“模型”部分检查训练状态。
然后您需要获取用于测试的模型ID。
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)
)
测试¶
我们将从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): """显示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
# 初始化Cohere嵌入模型,`input_type` 在索引和检索时是不同的。index_embed_model = CohereEmbedding( cohere_api_key=cohere_api_key, model_name="embed-english-v3.0", input_type="search_document",)query_embed_model = CohereEmbedding( cohere_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,)
结果_df = pd.DataFrame()embed_name = "CohereEmbedding"# 遍历rerankersfor rerank_name, reranker in RERANKERS.items(): print(f"运行Reranker的评估: {rerank_name}") # 定义检索器 class CustomRetriever(BaseRetriever): """执行向量搜索和知识图搜索的自定义检索器""" def __init__( self, vector_retriever: VectorIndexRetriever, ) -> None: """初始化参数。""" self._vector_retriever = vector_retriever super().__init__() def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]: """给定查询,检索节点。""" 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]: """异步给定查询,检索节点。 用户实现。 """ 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)
Cohere自定义重新排名器已经带来了改进。重要的是要强调,确定最佳数量的难例和是否使用随机抽样或余弦抽样应该基于实验结果。本指南提供了一个框架,用于通过自定义Cohere重新排名器增强检索系统。
在选择难例方面存在改进的潜力;社区对这一领域的贡献是受欢迎的。