引言

实现高效的检索增强生成 (RAG) 管道很大程度上依赖于强大的检索性能。正如我们在之前的博客文章中探讨的那样，重排序器（reranker）对提升检索性能有显著影响。但如果我们能更进一步呢？如果我们的重排序器不是任何一个重排序器，而是针对我们特定领域或数据集进行调优的呢？这种专业化能否进一步增强检索性能？

为了回答这些问题，我们转向 CohereAI 的重排序器微调（自定义重排序器）模型测试版。通过将其与 LlamaIndex 集成，我们现在提供了使用我们简化流程构建您自己的 Cohere 自定义重排序器的能力。

在这篇博客文章中，我们将指导您完成使用 LlamaIndex 创建 Cohere 自定义重排序器并评估检索性能的步骤。

如需动手实践，您可以按照 Google Colab Notebook 中的教程操作。

让我们开始使用 LlamaIndex 微调 Cohere 重排序器（自定义重排序器）。

注：本指南用于微调 Cohere 重排序器（自定义重排序器）。本教程末尾呈现的结果仅适用于所选的数据集和参数。我们建议您在使用自己的数据集和各种参数进行实验后再决定是否将其纳入您的 RAG 管道。

环境设置

!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

下载数据

我们将使用 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'

加载数据

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 个标记。为了满足这一点，我们将chunk_size限制为 400 个标记。（每个块最终将被视为一个相关段落/困难负样本）

# 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 从块中创建问题。

llm = OpenAI(api_key=openai_api_key, temperature=0, model='gpt-4')

从每个节点/块生成问题的提示。

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

# Training dataset
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")

# Validation dataset
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")

# Testing dataset
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（相关段落）：这代表包含问题答案信息的文档或段落列表。对于每个查询，必须至少有一个 relevant_passage。

hard_negatives（困难负样本）：这代表不包含问题答案的块或段落。需要注意的是，困难负样本是可选的，但提供至少约 5 个困难负样本将带来显著的改进。

您可以查看文档获取更多详细信息。

我们需要一个 embedding 模型来使用余弦相似度方法创建困难负样本。

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

让我们创建 3 个数据集。

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

# 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 分钟。

# 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_negatives1",
    model_type="RERANK",
    base_model="english",
    api_key = cohere_api_key
)
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_negatives1",
    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_negatives1",
    model_type="RERANK",
    base_model="english",
)
finetune_model_cosine_hard_negatives.finetune()

作业提交后，您可以在仪表板的models部分检查训练状态。您可以在仪表板中检查作业状态，应该会看到类似以下的图像。

然后您需要获取 Cohere 重排序器模型进行测试。

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. 不使用重排序器。
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,
}

创建索引和检索器用于评估。

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

service_context_index = ServiceContext.from_defaults(llm=None, embed_model=index_embed_model)
service_context_query = ServiceContext.from_defaults(llm=None, embed_model=query_embed_model)

vector_index = VectorStoreIndex(uber_nodes[:150], service_context=service_context_index)
vector_retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=10, service_context=service_context_query)

定义一个函数来显示结果

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

遍历不同的重排序器，使用自定义检索器评估检索性能。

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

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

结果

从上表（1 — 不使用重排序器，2 — 使用基础 Cohere 重排序器，3-5：微调重排序器（自定义重排序器））中可以看出，微调重排序器（自定义重排序器）带来了性能提升。至关重要的是，选择最佳数量的困难负样本以及在随机采样或余弦采样之间做出决定都应基于经验证据。本指南提供了一种通过微调 Cohere 重排序器来改进检索系统的结构化方法。

总结

在这篇博客文章中，我们展示了如何使用 LlamaIndex 对 Cohere 重排序器（自定义重排序器）进行微调，这提高了检索性能指标。我们热切期待社区利用这些能力来提升其 RAG 管道中的检索效率。此外，在选择困难负样本方面还有改进空间，我们邀请社区为此做出贡献。