精通重排序模型：优化结果的5个关键步骤

Rifx.Online
Natural Language Processing , Machine Learning , AI Applications
05 Mar, 2025

好久没写博客了，很高兴能再次开始！

在我之前的文章中，我们探讨了微调嵌入模型的世界——这是改进检索系统的关键一步。今天，我们将更进一步，深入研究微调重新排序模型。虽然嵌入可以帮助我们检索相关文档，但重新排序器会优化这些结果，以确保最准确和上下文相关的匹配。

在本文中，我将向您介绍我的方法，重点是准备数据和微调自定义重新排序器。

让我们开始吧。

Cross Encoders

Cross-encoders 是一种神经网络架构，主要用于自然语言处理 (NLP) 中需要理解两段文本之间关系的任务，例如句子对。它们特别适用于语义相似性、问答和自然语言推理等任务。

Cross-Encoders 的工作原理：

它们的工作原理

输入对处理：

将两段文本连接起来，通常用特殊标记（如 [SEP]）分隔。例如，在 BERT 风格的模型中，格式通常是： [CLS] Text1 [SEP] Text2 [SEP]
这种格式允许模型联合处理两个输入。

联合编码：

与独立编码文本的双编码器不同，交叉编码器一起处理连接的序列，允许模型捕获两个文本之间的 token 级交互。
这使得交叉编码器对于需要理解深层关系的任务更准确，但与双编码器相比，计算量也更大。

输出：

输出通常源自 [CLS] token 的最终隐藏状态，它充当整个序列的表示。
对于二元分类等任务，分数或概率可以指示关系（例如，语义相似性）。对于其他任务，额外的层可能会产生更复杂的输出。

用例：

语义文本相似性 (STS)：确定两个句子有多相似。
自然语言推理 (NLI)：对假设是蕴含、矛盾还是中立于前提进行分类。
问答：根据候选答案与问题的相关性对其进行排名。
信息检索：在初始检索步骤后对候选文档或段落进行重新排序。

Re-ranking

在 检索增强生成 (RAG) 中，重新排序是提高检索到的文档或段落的质量的关键步骤，然后再使用它们来生成最终答案。RAG 结合了基于检索的方法（从大型语料库中提取相关文档）和生成模型（根据检索到的内容生成答案）。重新排序有助于确保将最相关和高质量的文档优先用于生成步骤。

为什么重新排序在 RAG 中很重要

对重新排序的需求源于初始检索阶段的局限性。检索器，例如稀疏检索器（如 BM25）或密集检索器（如双编码器），可能会返回大量候选文档，这些文档的与查询的相关性排名并不完美。重新排序通过使用更复杂的模型（例如交叉编码器）来改进检索到的文档的顺序来解决这个问题，以更好地评估每个文档与查询的相关性。通过将最相关的文档馈送到生成模型，最终输出（无论是答案还是摘要）都会变得更加准确和上下文相关。

重新排序在 RAG 中的工作原理

在实践中，重新排序作为多步骤流程的一部分工作。

初始检索：

检索器（例如，BM25、DPR 或密集嵌入模型）根据输入查询从大型语料库中提取一组候选文档或段落。

重新排序：

重新排序器（例如，交叉编码器）评估每个检索到的文档与查询的相关性。
交叉编码器一起处理查询和每个文档，捕获细粒度的交互并产生相关性分数。
文档根据这些分数重新排序。

生成：

将排名前 k 位的重新排序文档传递给生成模型（例如，GPT、LLAMA 或 Qwen）。
生成模型根据最相关的文档生成最终输出（例如，答案）。

微调重新排序模型

微调重新排序模型对于针对特定任务或领域优化其性能至关重要。虽然 BERT 或 RoBERTa 等预训练模型对语言有一般的理解，但它们可能无法针对任务的细微差别进行定制，例如根据与查询的相关性对文档进行排名。

微调 使模型适应任务，提高准确性和相关性评分。它有助于模型学习特定于任务的关系、特定于领域的知识和隐式上下文连接，这对于准确的重新排序至关重要。

微调还可以使模型与目标数据分布保持一致。此过程增强了模型对输入数据中的噪声和变化的鲁棒性，例如拼写错误或释义查询。

微调对于专业领域（例如，医疗或法律）以及当有标记数据可用时尤其重要。通过微调，重新排序模型可以更好地优化检索到的文档，从而确保为RAG等系统中的问答等下游任务提供更高质量的输入。

数据集格式

在微调重排序模型时，使用两种常见的数据集格式：连续评分制和离散类别制。

连续评分制格式涉及句子或文本对，并带有连续的相关性评分（例如，介于 0 和 1 之间）。例如，使用 sentence-transformers 库中的 InputExample 类，您可以定义类似 ["sentence1", "sentence2"] 的对，并附带标签，如 0.3 或 0.8，表示两个文本之间的相关性或相似程度。这种格式非常适合对相关性进行分级的任务，例如语义文本相似度或文档排序。

示例：

train_samples = [
    InputExample(texts=["sentence1", "sentence2"], label=0.3),
    InputExample(texts=["Another", "pair"], label=0.8),
]

另一方面，离散类别制格式使用预定义的类别来表示句子对之间的关系。例如，在自然语言推理 (NLI) 任务中，句子对被标记为“矛盾”、“蕴含”或“中立”，这些类别被映射为整数值（例如，0、1、2）。这种格式适用于需要类别分类的任务，例如确定句子之间的逻辑关系。这两种格式都广泛用于微调重排序模型，选择哪种格式取决于具体的任务和数据的性质。

示例：

label2int = {"contradiction": 0, "entailment": 1, "neutral": 2}
train_samples = [
    InputExample(texts=["sentence1", "sentence2"], label=label2int["neutral"]),
    InputExample(texts=["Another", "pair"], label=label2int["entailment"]),
]

生成用于微调的合成数据

在微调重排序模型时，拥有高质量的训练数据至关重要。如果您可以访问来自已部署解决方案的真实数据，这是最佳选择。真实数据反映了您的系统处理的实际查询和文档，确保微调后的模型与您的特定用例保持一致。您可以从 Langfuse、LangSmith 或 LlamaTrace 等可观察性工具中收集这些数据，这些工具跟踪 LLM 应用程序中的交互，并提供对查询-文档对及其相关性的见解。

但是，如果无法获得真实数据，您可以生成合成数据来微调您的模型。我使用的方法是不寻常但有效的。以下是这个想法：

我利用一个 检索器（例如，向量存储）来获取一组查询的候选文档。然后，我使用 Cohere 的 rerank-v3.5（一种最先进的重排序模型）来优化检索到的文档并分配相关性评分。通过从这些重新排序的结果中抽样，我创建了一个数据集，该数据集模仿了高质量重排序器的行为。这种方法允许我的微调模型从一个更优越的模型的重排序模式中学习，从而有效地提炼其知识。

此外，我在抽样过程中引入了随机性，使模型能够接触到各种查询-文档对，确保它能够很好地泛化到各种场景。虽然合成数据可能无法完美地复制真实世界的分布，但当真实数据稀缺时，这种方法提供了一种实用的引导微调的方法。

Bge-reranker-base（我们的基础模型）

BAAI/bge-reranker-base 是一个高性能的交叉编码器模型，专为文本重排序任务而设计，由 北京智源人工智能研究院 (BAAI) 开发。它是 BAAI 通用嵌入 (BGE) 系列的一部分，该系列侧重于改进类似 RAG（检索增强生成）的检索增强系统，我们将使用我们的数据微调此模型，然后使用它来重新排序 llama-index 检索到的文档。

代码时间

安装必要的库

首先，您需要安装所需的库。运行以下命令以安装所有依赖项：

pip install sentence-transformers llama-index llama-index-llms-gemini  llama-index-embeddings-gemini llama-index-postprocessor-cohere-rerank

设置环境变量

在深入研究代码之前，请确保您拥有 Google 和 Cohere 的必要 API 密钥。这些密钥用于进行身份验证并使用它们各自的服务。将它们存储为环境变量，以便在整个脚本中安全且轻松地访问：

import os

GOOGLE_API_KEY = "xxxxxxxxx"
os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY

COHERE_API_KEY = "xxxxxxxxxxxxxxx"
os.environ["COHERE_API_KEY"] = COHERE_API_KEY

导入必要的库

接下来，导入您将在脚本中使用的库和模块。其中包括用于文档加载、文本拆分、嵌入生成、重排序和微调的工具。以下是导入列表：

from llama_index.core.evaluation import generate_question_context_pairs
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.postprocessor.cohere_rerank import CohereRerank
from sentence_transformers import CrossEncoder, InputExample
from llama_index.embeddings.gemini import GeminiEmbedding
from llama_index.core.node_parser import SentenceSplitter
from sentence_transformers import InputExample
from llama_index.llms.gemini import Gemini
from torch.utils.data import DataLoader
from llama_index.core import Settings
import pandas as pd
import random

定义模型和 cohere 重排序器

现在，定义您将使用的语言模型 (LLM) 和嵌入模型。在这种情况下，我们同时使用 Gemini 进行文本生成和嵌入。

## define the models
llm = Gemini(
    model="models/gemini-1.5-flash",
)

embed_model = GeminiEmbedding(
    model_name="models/embedding-001", api_key=GOOGLE_API_KEY
)
Settings.embed_model = embed_model
Settings.llm = llm

设置 Cohere 重排序器，它将用于优化检索到的文档。top_n 参数指定要重新排序的文档数量，而 model 参数指定 Cohere 重排序器的版本（例如，rerank-v3.5）。

## define the cohere reranker
cohere_rerank = CohereRerank(api_key=COHERE_API_KEY, top_n=5 , model="rerank-v3.5") # top_n is the number of documents to rerank

准备我们的数据

创建问答数据集

为了生成用于微调的合成数据集，我们首先需要从文档集合中创建一个问答数据集。create_qa_dataset 函数处理此过程。它接收一个文档目录，将其分割成更小的块（例如，256 个 token），并使用语言模型 (LLM) 为每个块生成问题。每个块的问题数量是可自定义的，允许您控制数据集的密度。

def create_qa_dataset(input_dir ,  llm , num_questions_per_chunk):
    """
    Create a question-answer dataset from a directory of documents
    Args:
        input_dir: str
            The directory containing the documents
        llm: llama_index.llms.gemini.Gemini
            The LLM model to use for generating questions
        num_questions_per_chunk: int
            The number of questions to generate per chunk
    Returns:
        list: A list of questions
    """
    documents = SimpleDirectoryReader(input_dir=input_dir).load_data()
    node_parser = SentenceSplitter(chunk_size=256)
    nodes = node_parser.get_nodes_from_documents(documents)
    qa_dataset = generate_question_context_pairs(
          nodes,
          llm=llm,
          num_questions_per_chunk=num_questions_per_chunk
      )
    queries = qa_dataset.queries.values()
    return list(queries) , documents

此函数输出一个查询列表和相应的文档，这些文档作为下一步的基础。

创建评分查询上下文数据集

一旦我们有了查询，我们需要将它们与相关的上下文配对并分配相关性分数。create_score_query_context_dataset 函数通过利用检索器和 Cohere 重新排序器来完成此操作。

def create_score_query_context_dataset(queries , documents , split) :
    """
    Create a dataset of queries, contexts and scores
    Args:
        queries: list
            A list of queries
        documents: list
            A list of documents
        split: str
            The split of the dataset
    Returns:
        pd.DataFrame: A DataFrame containing the queries, contexts and scores
    """
  contexts = []
  scores = []
  # create the index
  index = VectorStoreIndex.from_documents(documents)
  # create the retriever
  retriever = index.as_retriever(verbose=True, similarity_top_k=10)
  for query in queries:
      # retrieve the top 10 documents
      nodes = retriever.retrieve(query)
      # rerank the documents using cohere and get the top 5 document
      response = cohere_rerank.postprocess_nodes(
                    nodes=nodes, query_str=query
                )
      random_number = random.randint(0, len(response)-1)
      contexts.append(response[random_number].text)
      scores.append(response[random_number].score)
  assert len(queries) == len(contexts) == len(scores)
  df = pd.DataFrame({"query": queries, "context": contexts, "score": scores})
  df.to_csv(f"{split}-data.csv", index=False)
  return df

其工作原理如下：

索引文档： 使用向量存储对文档进行索引，从而可以高效地检索相关块。
检索和重新排序： 对于每个查询，检索器会提取前 k 个文档（例如，前 10 个）。然后使用 Cohere 重新排序器对这些文档进行重新排序以优化结果。
随机抽样： 为了引入多样性，从重新排序的结果中随机选择一个文档。这确保了模型接触到各种查询-上下文对。
数据集创建： 查询、上下文及其相应的相关性分数被编译成一个 DataFrame，并保存为 CSV 文件以供进一步使用。

创建训练数据集

为了准备训练数据集，我们首先使用 create_qa_dataset 函数生成一个问答数据集。该函数处理指定目录 (/content/train_data) 中的文档，将其分割成块，并使用 Gemini 语言模型为每个块生成问题。每个块的问题数量设置为 1，确保数据集的重点和可管理性。

## create the the train dataset
queries , documents = create_qa_dataset("/content/train_data", llm , 1 )
train_data = create_score_query_context_dataset(queries , documents , "train")
train_samples = [
    InputExample(texts=[row['query'], row['context']], label=row['score'])
    for _, row in train_data.iterrows()
]
train_dataloader = DataLoader(train_samples, shuffle=True, batch_size=8)

一旦查询和文档准备就绪，我们使用 create_score_query_context_dataset 函数创建评分查询上下文数据集。此函数检索并重新排序每个查询的文档，分配相关性分数，并将数据编译成结构化格式。然后，将生成的数据集转换为 InputExample 对象的列表，这些对象用于创建用于训练的 DataLoader。DataLoader 对数据进行洗牌并以批次为 8 的方式处理数据，从而优化训练过程。

创建验证数据集

验证数据集的创建过程类似。

## create the validation dataset
queries , documents = create_qa_dataset("/content/val_data", llm , 1)
val_data = create_score_query_context_dataset(queries , documents , "validation")
val_samples = [
    InputExample(texts=[row['query'], row['context']], label=row['score'])
    for _, row in val_data.iterrows()
]
val_dataloader = DataLoader(val_samples, shuffle=True, batch_size=3)

初始化模型

为了开始微调，我们初始化一个预先训练的 Cross-Encoder 模型。在这种情况下，我们使用 BAAI/bge-reranker-base model 作为基础模型。我们之前讨论过它。

## Initialize a pre-trained Cross-Encoder model
model = CrossEncoder('BAAI/bge-reranker-base')

创建自定义评估器

为了在训练期间监控模型的性能，我们定义了一个名为 MSEEval 的自定义评估器。该评估器计算模型预测的相关性分数与真实分数之间的均方误差 (MSE)。它记录结果并将其保存到 CSV 文件中，以便于跟踪。

评估器使用验证数据集 (val_dataloader) 以固定间隔评估模型的性能。

from sentence_transformers.evaluation import SentenceEvaluator
import torch
from torch.utils.data import DataLoader
import logging
from sentence_transformers.util import batch_to_device
import os
import csv
from sentence_transformers import CrossEncoder
from tqdm.autonotebook import tqdm

logger = logging.getLogger(__name__)

class MSEEval(SentenceEvaluator):
    """
    Evaluate a model based on its accuracy on a labeled dataset

This requires a model with LossFunction.SOFTMAX

The results are written in a CSV. If a CSV already exists, then values are appended.
    """

    def __init__(self,
                 dataloader: DataLoader,
                 name: str = "",
                 show_progress_bar: bool = True,
                 write_csv: bool = True):
        """
        Constructs an evaluator for the given dataset

:param dataloader:
            the data for the evaluation
        """
        self.dataloader = dataloader
        self.name = name
        self.show_progress_bar = show_progress_bar

        if name:
            name = "_"+name

        self.write_csv = write_csv
        self.csv_file = "accuracy_evaluation"+name+"_results.csv"
        self.csv_headers = ["epoch", "steps", "accuracy"]

    def __call__(self, model: CrossEncoder, output_path: str = None, epoch: int = -1, steps: int = -1) -> float:
        model.model.eval()
        total = 0
        loss_total = 0

        if epoch != -1:
            if steps == -1:
                out_txt = " after epoch {}:".format(epoch)
            else:
                out_txt = " in epoch {} after {} steps:".format(epoch, steps)
        else:
            out_txt = ":"

        loss_fnc = torch.nn.MSELoss()
        activation_fnc = torch.nn.Sigmoid()

        logger.info("Evaluation on the "+self.name+" dataset"+out_txt)
        self.dataloader.collate_fn = model.smart_batching_collate
        for features, labels in tqdm(self.dataloader,  desc="Evaluation", smoothing=0.05, disable=not self.show_progress_bar):

            with torch.no_grad():
                model_predictions = model.model(**features, return_dict=True)
                logits = activation_fnc(model_predictions.logits)
                if model.config.num_labels == 1:
                    logits = logits.view(-1)
                loss_value = loss_fnc(logits, labels)

            total += 1 # number of batches
            loss_total += loss_value.cpu().item()
        mse = loss_total/total

        logger.info("MSE: {:.4f} ({}/{})\n".format(mse, loss_total, total))

        if output_path is not None and self.write_csv:
            csv_path = os.path.join(output_path, self.csv_file)
            if not os.path.isfile(csv_path):
                with open(csv_path, newline='', mode="w", encoding="utf-8") as f:
                    writer = csv.writer(f)
                    writer.writerow(self.csv_headers)
                    writer.writerow([epoch, steps, mse])
            else:
                with open(csv_path, newline='', mode="a", encoding="utf-8") as f:
                    writer = csv.writer(f)
                    writer.writerow([epoch, steps, mse])

        return mse

定义评估器后，我们使用验证数据加载器对其进行初始化：

## create the evaluator
evaluator = MSEEval(val_dataloader,show_progress_bar=True , write_csv=True )

训练模型

准备好模型和评估器后，我们开始使用训练数据集微调模型。使用 fit 方法来训练模型，其中包含以下关键参数：

train_dataloader：训练数据，按批处理。
evaluator：用于在训练期间监控性能的自定义评估器。
epochs：训练轮数（在本例中设置为 1）。
warmup_steps：用于稳定训练的预热步数。
evaluation_steps：训练期间评估的频率。
output_path：保存微调模型的目录。
save_best_model：是否根据评估结果保存性能最佳的模型。
use_amp：是否使用自动混合精度 (AMP) 来加快训练速度。
scheduler：学习率调度器（例如，warmupcosine）。
show_progress_bar：是否在训练期间显示进度条。

## Train the model
model.fit(
    train_dataloader=train_dataloader,
    evaluator=evaluator,
    epochs=1,
    warmup_steps=100,
    evaluation_steps=3,
    output_path="my_model",
    save_best_model=True,
    use_amp=True,
    scheduler= 'warmupcosine',
    show_progress_bar=True,
)

登录并将模型推送到 Hub

对模型进行微调后，您可以通过将其推送到 Hugging Face 模型 Hub 来保存和共享它。这允许您对模型进行版本控制、与他人协作以及在生产环境中部署它。

为此，您首先需要使用 notebook_login 函数登录到您的 Hugging Face 帐户。这将提示您输入您的 Hugging Face 凭据。

from huggingface_hub import notebook_login

notebook_login()

登录后，您可以使用 push_to_hub 方法将微调后的模型推送到 Hub。将 user/bge-reranker-base-finetuned 替换为您想要的存储库名称。这将把模型权重、配置和其他必要的文件上传到 Hub。

model.push_to_hub("user/bge-reranker-base-finetuned")

使用微调模型与 llama-index

将模型上传到 Hub 后，您可以轻松地将其集成到您的 LlamaIndex 管道中进行重新排序。 SentenceTransformerRerank 类允许您加载微调模型并使用它来重新排序检索到的文档。 top_n 参数指定重新排序后要返回的文档数量。

from llama_index.core.postprocessor import SentenceTransformerRerank

rerank = SentenceTransformerRerank(
    model="user/bge-reranker-base-finetuned", top_n=3
)

在处理嵌入后，微调重新排序模型是顺理成章的下一步，它是一种改进系统理解和优先处理信息的强大方法。从生成合成数据到训练和评估模型，每一步都带来了其自身的挑战和回报。无论您是使用真实世界的数据还是创建自己的数据集，关键在于进行实验、迭代和改进。我希望本指南能激励您探索在您自己的项目中进行微调的潜力。