介绍

在本文中，我们将引导你通过创建一个与LlamaIndex配合使用的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

下载数据

我们将使用Lyft 2021年的10-K SEC报告进行训练，并使用Uber 2021年的10-K 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：这代表了问题或目标。

related_passages：这代表一个包含回答查询问题信息的文件或段落列表。对于每个查询，必须至少有一个related_passages。

困难负面：这代表不包含查询答案的块或段落。应该注意的是，困难负面是可选的，但提供至少约5个困难负面将导致有意义的改进。

我们需要有一个嵌入模型来使用余弦相似性方法创建困难负面。

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

作业提交后，你可以在仪表板的模型部分检查培训状态。你可以在面板中检查作业的状态，你应该会看到一个类似于下面的图像。

然后，你需要获取Cohere Reranker模型进行测试。

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

测试

我们将使用以下不同的重新排序器对优步的前150个节点进行测试。

1. 没有重新排序。
2. 相干重新排序。
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

在不同的重新排序器上循环，并使用Custom Retriever评估检索性能。

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-带基本粘性重新排列，3–5：微调重新排列（自定义重新排列））中，我们可以看到微调重新排列的性能有所提高。需要注意的是，硬阴性的最佳数量的选择，以及随机或余弦采样之间的决定，都应该以经验证据为基础。本指南提供了一种结构化的方法，通过对相干重排器进行微调来改进检索系统。

总结

在本文中，我们展示了使用LlamaIndex对Cohere重新排序器（自定义重新排序器）进行微调，这提高了检索性能指标。