自定义 Cohere 重排序器¶

本笔记本提供了关于使用 LlamaIndex 抽象构建 Cohere 自定义重排序器的教程。完成学习后，您将能够创建自定义重排序器，并将其用于增强数据检索。

重要说明： 本笔记本提供的是 Cohere 自定义重排序器的使用指南。教程末尾展示的结果仅针对特定数据集和参数有效。我们建议在将其纳入您的 RAG 流程前，先使用您的数据集和不同参数进行实验验证。

安装设置¶

让我们安装必要的软件包。

In [ ]:

%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-postprocessor-cohere-rerank %pip install llama-index-llms-openai %pip install llama-index-finetuning %pip install llama-index-embeddings-cohere

In [ ]:

!pip install llama-index cohere pypdf

!pip install llama-index cohere pypdf

初始化 API 密钥¶

OpenAI - 用于生成合成数据集。

CohereAI - 用于训练自定义重排序器并通过基础重排序器进行评估。

In [ ]:

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

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

In [ ]:

import os

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

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

In [ ]:

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.node_parser import SimpleNodeParser

# LLM
from llama_index.llms.openai import OpenAI

# Embeddings
from llama_index.embeddings.cohere import CohereEmbedding

# Retrievers
from llama_index.core.retrievers import BaseRetriever, VectorIndexRetriever

# Rerankers
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.postprocessor.cohere_rerank import CohereRerank
from llama_index.core.evaluation import EmbeddingQAFinetuneDataset
from llama_index.finetuning import generate_cohere_reranker_finetuning_dataset

# Evaluator
from llama_index.core.evaluation import generate_question_context_pairs
from llama_index.core.evaluation import RetrieverEvaluator

# Finetuner
from llama_index.finetuning import CohereRerankerFinetuneEngine


from typing import List
import pandas as pd

import nest_asyncio

nest_asyncio.apply()

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader from llama_index.core.node_parser import SimpleNodeParser # LLM from llama_index.llms.openai import OpenAI # Embeddings from llama_index.embeddings.cohere import CohereEmbedding # Retrievers from llama_index.core.retrievers import BaseRetriever, VectorIndexRetriever # Rerankers 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.postprocessor.cohere_rerank import CohereRerank from llama_index.core.evaluation import EmbeddingQAFinetuneDataset from llama_index.finetuning import generate_cohere_reranker_finetuning_dataset # Evaluator from llama_index.core.evaluation import generate_question_context_pairs from llama_index.core.evaluation import RetrieverEvaluator # Finetuner from llama_index.finetuning import CohereRerankerFinetuneEngine from typing import List import pandas as pd import nest_asyncio nest_asyncio.apply()

下载数据¶

我们将使用 Lyft 公司 2021 年 10-K 格式的 SEC 备案文件作为训练数据，并采用 Uber 公司 2021 年 10-K 格式的 SEC 备案文件进行评估。

In [ ]:

!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'

!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'

加载数据¶

In [ ]:

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

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 个词元（每个分块最终将被视为相关段落/困难负样本）。

In [ ]:

# Limit chunk size to 400
node_parser = SimpleNodeParser.from_defaults(chunk_size=400)

# Create nodes
lyft_nodes = node_parser.get_nodes_from_documents(lyft_docs)
uber_nodes = node_parser.get_nodes_from_documents(uber_docs)

# Limit chunk size to 400 node_parser = SimpleNodeParser.from_defaults(chunk_size=400) # Create nodes lyft_nodes = node_parser.get_nodes_from_documents(lyft_docs) uber_nodes = node_parser.get_nodes_from_documents(uber_docs)

我们将使用 gpt-4 从文本块生成问题。

In [ ]:

llm = OpenAI(temperature=0, model="gpt-4")

llm = OpenAI(temperature=0, model="gpt-4")

根据每个节点/文本块生成问题的提示模板。

In [ ]:

# Prompt to generate questions
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.\
"""

# Prompt to generate questions 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 组（查询+相关段落）配对数据用于训练（可含或不含困难负样本），以及 64 组配对用于验证。请注意验证环节是可选项。

训练阶段： 我们使用 Lyft 数据集的前 256 个节点创建训练配对。

验证阶段： 我们将使用 Lyft 数据集接下来的 64 个节点进行验证。

测试阶段： 我们将使用 Uber 数据集的 150 个节点进行测试。

In [ ]:

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,
)

# Save [Optional]
qa_dataset_lyft_train.save_json("lyft_train_dataset.json")

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, ) # Save [Optional] qa_dataset_lyft_train.save_json("lyft_train_dataset.json")

In [ ]:

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,
)

# Save [Optional]
qa_dataset_lyft_val.save_json("lyft_val_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, ) # Save [Optional] qa_dataset_lyft_val.save_json("lyft_val_dataset.json")

In [ ]:

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,
)

# Save [Optional]
qa_dataset_uber_val.save_json("uber_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, ) # Save [Optional] qa_dataset_uber_val.save_json("uber_val_dataset.json")

既然我们已经从每个数据块中汇总了问题，接下来将按照训练自定义重排模型所需的规范来格式化数据。

数据格式与要求¶

在训练和评估阶段，当前系统接受三元组格式的数据，每行应包含以下内容：

query（查询）： 表示问题或目标

relevant_passages（相关段落）： 表示包含回答查询信息的文档或段落列表。每个查询必须至少对应一个相关段落

hard_negatives（困难负样本）： 表示不包含查询答案的文本块或段落。需注意困难负样本为可选字段，但提供至少约5个困难负样本将带来显著效果提升。

参考文档

In [ ]:

# Initialize the Cohere embedding model which we use it for creating Hard Negatives.
embed_model = CohereEmbedding(
    api_key=cohere_api_key,
    model_name="embed-english-v3.0",
    input_type="search_document",
)

# Initialize the Cohere embedding model which we use it for creating Hard Negatives. embed_model = CohereEmbedding( api_key=cohere_api_key, model_name="embed-english-v3.0", input_type="search_document", )

让我们创建三个数据集。

1. 不包含困难负样本的数据集。
2. 随机选择困难负样本的数据集。
3. 基于余弦相似度选择困难负样本的数据集。

In [ ]:

# Train and val datasets without hard negatives.
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"
)

# Train and val datasets with hard negatives selected at random.
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,
)

# Train and val datasets with hard negatives selected based on cosine similarity.
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,
)

# Train and val datasets without hard negatives. 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" ) # Train and val datasets with hard negatives selected at random. 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, ) # Train and val datasets with hard negatives selected based on cosine similarity. 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分钟。

In [ ]:

# Reranker model with 0 hard negatives.
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()

# Reranker model with 5 hard negatives selected at random
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()

# Reranker model with 5 hard negatives selected based on cosine similarity
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()

# Reranker model with 0 hard negatives. 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() # Reranker model with 5 hard negatives selected at random 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() # Reranker model with 5 hard negatives selected based on cosine similarity 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()

任务提交后，您可以在仪表盘的 models 部分查看训练状态，访问地址为 https://dashboard.cohere.com/models。

然后你需要获取用于测试的模型 ID。

In [ ]:

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)
)

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

1. 不使用重排序器（Reranker）
2. 使用 Cohere 重排序器（未经任何训练）
3. 使用自定义重排序器（不含困难负样本）
4. 使用自定义重排序器（随机选取困难负样本）
5. 使用自定义重排序器（基于余弦相似度选取困难负样本）

In [ ]:

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,
}

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, }

显示结果的函数

In [ ]:

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

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

In [ ]:

# Initialize the Cohere embedding model, `input_type` is different for indexing and retrieval.
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,
)

# Initialize the Cohere embedding model, `input_type` is different for indexing and retrieval. 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, )

In [ ]:

results_df = pd.DataFrame()

embed_name = "CohereEmbedding"

# Loop over rerankers
for rerank_name, reranker in RERANKERS.items():
    print(f"Running Evaluation for Reranker: {rerank_name}")

    # Define Retriever
    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)

results_df = pd.DataFrame() embed_name = "CohereEmbedding" # Loop over rerankers for rerank_name, reranker in RERANKERS.items(): print(f"Running Evaluation for Reranker: {rerank_name}") # Define Retriever 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)

检查结果¶

In [ ]:

print(results_df)

print(results_df)

Cohere 自定义重排序器已带来显著改进。需要重点强调的是，确定最佳难负例数量以及选择随机采样还是余弦采样，都应基于实验结果。本指南提供了一个框架，帮助通过自定义 Cohere 重排序器优化检索系统。

当前难负例的选取策略仍有改进空间，欢迎社区成员在此领域贡献力量。