自定义 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 10K SEC Filings 数据进行训练，使用 Uber 2021 10K SEC Filings 数据进行评估。

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

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

In [ ]

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

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

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

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

现在我们已经从每个 chunk 中收集了问题，我们将按照训练自定义重排序器所需的规范来格式化数据。

数据格式和要求¶

对于训练和评估，目前都接受三元组格式的数据，每一行应包含以下内容

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

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

hard_negatives（困难负样本）：代表不包含查询答案的 chunk 或段落。需要注意的是，困难负样本是可选的，但提供至少约 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", )

让我们创建 3 个数据集。

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

提交任务后，您可以在 https://dashboard.cohere.com/models 的仪表板的 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. 不使用重排序器。
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 自定义重排序器增强检索系统的框架。

困难负样本的选择方面仍有改进的潜力；欢迎社区在此领域做出贡献。