RankLLM：用大语言模型重新排序的 Python 工具包

在信息检索领域，如何从海量数据中精准高效地筛选出与用户查询最相关的文档是个关键问题。近年来，随着大型语言模型（LLM）的兴起，利用这些强大的模型进行重新排序（reranking）成为学术界和工业界的研究热点。今天要给大家介绍的 RankLLM，就是一个基于 Python 的开源工具包，它能让大家轻松地利用各种大型语言模型来优化文档重新排序的效果。

RankLLM 是什么？

RankLLM 是由滑铁卢大学的研究人员开发的一个 Python 工具包，旨在帮助大家在多阶段检索系统中利用大型语言模型进行重新排序。简单来说，它就像是一个桥梁，把各种先进的大型语言模型和文档重新排序任务连接起来，让即使没有很深的机器学习背景的开发者，也能方便地利用这些模型来提升检索系统的性能。

这个工具包具有高度的模块化和可配置性，支持多种专有和开源的大型语言模型。它不仅能和常见的推理框架集成，还提供了一系列实用的功能，比如与 Pyserini 整合进行检索、内置的多阶段流程评估工具，以及对输入提示和模型响应的详细分析模块。这些特点使得 RankLLM 在实际应用中非常灵活和强大。

为什么需要 RankLLM？

在传统的文档检索过程中，通常会先用一些快速的检索方法（比如基于 BM25 的方法）来初步筛选出一批候选文档，然后再用更复杂的模型对这些候选文档进行重新排序，以提高检索结果的相关性。这个重新排序的阶段非常重要，因为它直接决定了最终呈现给用户的文档质量。

以前，重新排序主要依赖于两种技术：交叉编码器（cross-encoders）和双编码器（bi-encoders）。但现在，随着提示解码器（prompt-decoders）这类新型重新排序模型的出现，利用大型语言模型来进行重新排序变得越来越流行。这些模型通过精心设计的提示（prompts）来引导模型对文档进行重新排序，往往能取得比传统方法更好的效果。

然而，研究这些新的重新排序方法通常需要开发者自己编写各种复杂的代码，这不仅增加了开发的难度，也不利于不同方法之间的比较和探索。RankLLM 的出现就是为了改变这种状况。它提供了一个统一的框架，把各种重新排序模型都整合在一起，让大家可以更方便地进行实验和比较，加速这一领域的发展。

RankLLM 的核心特点

接下来，我们详细看看 RankLLM 都有哪些核心特点，让它能在文档重新排序领域大显身手。

多样化的模型支持

RankLLM 一个很大的优势就是它能和各种各样的大型语言模型无缝集成。无论是专有的模型，比如 OpenAI 的 GPT 系列，还是开源的模型，像 LLaMA、Vicuna、Zephyr 和 Mistral 等，都能在 RankLLM 中使用。这意味着你可以根据自己的需求和资源情况，灵活选择不同的模型来进行重新排序任务。

而且，RankLLM 还支持大多数以 RankX 形式存在的模型抽象，通过使用 vLLM 等工具，能够高效地利用开源大型语言模型来进行排名。这种广泛的模型支持能力，让 RankLLM 成为一个非常通用的重新排序平台。

高度的模块化和可配置性

RankLLM 采用了模块化的架构设计，这使得它的各个部分可以像积木一样自由组合和定制。你可以根据自己的具体需求，选择不同的排名方法、大型语言模型、推理框架和提示模板。比如，你可以选择使用点-wise、对-wise 或者列表-wise 的重新排序方法，也可以针对不同的模型调整提示的格式和内容。

这种高度的可配置性，不仅让 RankLLM 能适应各种不同的应用场景，还为研究人员提供了充分的实验空间，方便他们探索不同的模型配置和参数设置对重新排序效果的影响。

辅助的端到端流程组件

除了核心的重新排序功能外，RankLLM 还提供了一些辅助的组件，帮助大家更高效地完成整个检索和重新排序的流程。

• 检索组件 ：RankLLM 内置了和 Pyserini 的集成，可以方便地利用 Pyserini 进行文档检索。Pyserini 是一个功能强大的信息检索工具，支持多种检索方法，比如 BM25 和各种基于深度学习的稠密检索方法。通过 RankLLM 的检索组件，你可以快速地从大规模文档集合中获取初步的候选文档列表，然后将其输入到重新排序模型中。
• 评估组件 ：为了方便大家评估重新排序的效果，RankLLM 提供了评估工具。它可以使用 TREC 评估运行格式来量化重新排序的有效性，支持多种常用的评估指标，如 nDCG 和 mAP。这样你就可以很方便地比较不同重新排序模型和参数设置的性能，选择最适合你应用的方案。
• 调用分析组件 ：这个组件可以详细记录模型每次调用时的提示和响应信息，包括输入输出的令牌数量等。这对于分析模型的行为、调试提示工程以及优化模型的使用效率都非常有帮助。通过研究模型的响应，你可以更好地了解模型是如何处理不同的查询和文档组合的，从而进一步改进你的重新排序策略。

可复现性特性

在研究和实际应用中，能够复现实验结果是非常重要的。RankLLM 为此提供了一系列的特性，比如预定义的配置、全面的日志记录、示例代码片段、详细的文档字符串和 README 指令，以及两步点击复现（2CR）功能。这些特性确保了实验过程的透明性和结果的可重复性，方便团队协作和知识共享。

RankLLM 的安装和使用

了解了 RankLLM 的特点后，想必大家已经迫不及待想试试这个工具了。下面我们就来看看如何安装和使用 RankLLM。

安装 RankLLM

RankLLM 的安装非常简单，主要有两种方式。

• 通过 PyPI 安装 ：RankLLM 已经发布在 Python Package Index 上，你可以直接使用 pip 命令来安装。打开你的终端或者命令行工具，输入以下命令：

• 从源码安装 ：如果你对 RankLLM 的开发感兴趣，或者想对它进行一些定制化的修改，可以从它的源码仓库进行安装。首先，你需要克隆 RankLLM 的 Git 仓库到本地，然后进入相应的目录，使用 pip 进行安装。具体的命令如下：

环境准备

在安装 RankLLM 之前，你还需要确保你的环境中已经安装了 CUDA 加速的 PyTorch，并且配置好了相应的 GPU 环境。此外，如果你要使用 RankLLM 的检索功能，还需要安装 OpenJDK 21 和 Maven。以下是具体的安装步骤：

• 创建 Conda 环境 ：推荐你为 RankLLM 创建一个独立的 Conda 环境，以避免与其他项目依赖冲突。使用以下命令创建并激活 Conda 环境：

• 安装 PyTorch ：根据你的操作系统（Windows 或 Linux），使用 pip 安装支持 CUDA 的 PyTorch 版本。例如，对于支持 CUDA 12.1 的环境，可以使用以下命令：

• 安装 OpenJDK 和 Maven ：如果你需要使用 RankLLM 的检索功能，还需要安装 OpenJDK 21 和 Maven。可以通过 Conda 进行安装，命令如下：

如果你还想使用 RankLLM 的其他可选功能，比如 SGLang 或 TensorRT-LLM 推理后端，也需要进行相应的安装。具体安装步骤可以参考 RankLLM 的官方文档。

快速入门示例

安装完成后，就可以开始使用 RankLLM 进行文档检索和重新排序了。下面是一个简单的示例，展示了如何使用 RankLLM 对 DL19 数据集中的查询进行检索和重新排序。

from pathlib import Path

from rank_llm.analysis.response_analysis import ResponseAnalyzer
from rank_llm.data import DataWriter
from rank_llm.evaluation.trec_eval import EvalFunction
from rank_llm.rerank import Reranker, get_openai_api_key
from rank_llm.rerank.listwise import (
    SafeOpenai,
    VicunaReranker,
    ZephyrReranker,
)
from rank_llm.retrieve.retriever import Retriever
from rank_llm.retrieve.topics_dict import TOPICS

# -------- 检索 --------

# 默认使用 BM25 检索前 100 个候选文档
dataset_name = "dl19"
retrieved_results = Retriever.from_dataset_with_prebuilt_index(dataset_name)

# 用户可以指定其他检索方法和候选文档数量
# retrieved_results = Retriever.from_dataset_with_prebuilt_index(
#     dataset_name, RetrievalMethod.SPLADE_P_P_ENSEMBLE_DISTIL, k=50
# )
# ---------------------------

# --------- 重新排序 ----------

# 使用 Rank Zephyr 模型
reranker = ZephyrReranker()

# 使用 Rank Vicuna 模型
# reranker = VicunaReranker()

# 使用 RankGPT 模型
# model_coordinator = SafeOpenai("gpt-4o-mini", 4096, keys=get_openai_api_key())
# reranker = Reranker(model_coordinator)

rerank_results = reranker.rerank_batch(requests=retrieved_results)
# ---------------------------

# ------- 评估 --------

# 评估检索结果
ndcg_10_retrieved = EvalFunction.from_results(retrieved_results, TOPICS[dataset_name])
print(ndcg_10_retrieved)

# 评估重新排序结果
ndcg_10_rerank = EvalFunction.from_results(rerank_results, TOPICS[dataset_name])
print(ndcg_10_rerank)

# 默认评估指标为 ndcg@10，也可以指定其他指标
# eval_args = ["-c", "-m", "map_cut.100", "-l2"]
# map_100_rerank = EvalFunction.from_results(rerank_results, topics, eval_args)
# print(map_100_rerank)

# eval_args = ["-c", "-m", "recall.20"]
# recall_20_rerank = EvalFunction.from_results(rerank_results, topics, eval_args)
# print(recall_20_rerank)

# ---------------------------

# --- 分析调用 ---
analyzer = ResponseAnalyzer.from_inline_results(rerank_results)
error_counts = analyzer.count_errors(verbose=True)
print(error_counts)
# ---------------------------

# ------ 保存结果 -------
writer = DataWriter(rerank_results)
Path(f"demo_outputs/").mkdir(parents=True, exist_ok=True)
writer.write_in_jsonl_format(f"demo_outputs/rerank_results.jsonl")
writer.write_in_trec_eval_format(f"demo_outputs/rerank_results.txt")
writer.write_inference_invocations_history(
    f"demo_outputs/inference_invocations_history.json"
)
# ---------------------------

在这个示例中，我们首先使用 BM25 方法从 DL19 数据集中检索出每个查询的前 100 个候选文档。然后，我们选择 Rank Zephyr 模型作为重新排序器，对这些候选文档进行重新排序。接着，我们使用 RankLLM 的评估组件分别对原始检索结果和重新排序后的结果进行评估，输出它们的 nDCG@10 指标值。最后，我们通过调用分析组件统计模型响应中的错误情况，并将重新排序的结果保存为不同的文件格式，方便后续的分析和使用。

RankLLM 的模型选择

RankLLM 支持的模型种类繁多，涵盖了点-wise、对-wise 和列表-wise 重新排序的各种模型。下面我们将重点介绍其中一些比较有代表性的模型。

点-wise 重新排序模型：MonoT5

MonoT5 是一个典型的点-wise 重新排序模型，它基于 T5 架构，通过将查询和文档联合编码，生成一个表示文档相关性的分数。在 RankLLM 中，MonoT5 模型家族被用于点-wise 重新排序任务。它对每个候选文档独立地进行评分，然后根据这些评分对文档进行排序。MonoT5 的特点是简单高效，适合处理大规模的文档集合。

对-wise 重新排序模型：DuoT5

DuoT5 是专门用于对-wise 重新排序的模型。它会同时考虑两个文档与查询的相关性，并对它们进行比较，从而得到一个相对的排序关系。通过对文档对的比较，DuoT5 能够捕捉文档之间的相对信息，进而提高重新排序的准确性。不过，由于需要处理文档对，DuoT5 的计算复杂度相对较高，因此在重新排序大量文档时可能需要更多的计算资源。

列表-wise 重新排序模型

列表-wise 重新排序模型是目前比较受关注的一种方法，它能够同时处理整个候选文档列表，综合考虑文档之间的相互关系，从而生成更优的排序结果。RankLLM 中的列表-wise 模型包括 LiT5、RankVicuna、RankZephyr 等。

• LiT5 ：LiT5 是一个基于 T5 架构的列表-wise 重新排序模型，通过专门设计的提示和训练方法，能够有效地对文档列表进行重新排序。它支持多种变体，包括不同大小的模型和不同的训练配置，以满足不同场景下的需求。
• RankVicuna ：RankVicuna 是基于开源的 Vicuna 模型开发的列表-wise 重新排序模型。它利用 Vicuna 的强大语言理解和生成能力，结合精心设计的提示工程，能够对文档列表进行高质量的重新排序。RankVicuna 的优势在于它的开源特性，使得研究人员和开发者可以方便地进行定制和扩展。
• RankZephyr ：RankZephyr 是另一个开源的列表-wise 重新排序模型，它在多个任务中表现出了良好的性能和鲁棒性。RankZephyr 通过有效的模型架构和训练策略，能够在保证重新排序质量的同时，提高模型的运行效率。

除了上述这些模型，RankLLM 还支持其他一些专有和开源的列表-wise 重新排序模型，如 Gemini 等。这些模型各有特点，适用于不同的应用场景和需求。

RankLLM 的实验结果

为了验证 RankLLM 的有效性，研究人员在 MS MARCO V1 和 V2 数据集的 DL19–DL23 任务上进行了实验。实验结果表明，使用 RankLLM 中的各种重新排序模型，相较于传统的 BM25 检索方法，能够显著提高检索结果的相关性。

在实验中，BM25 检索方法作为第一阶段的检索器，用于从大规模文档集合中初步筛选出每个查询的前 100 个候选文档。然后，使用 RankLLM 中的不同重新排序模型对这些候选文档进行重新排序。实验结果显示，在多个数据集上，重新排序后的结果在 nDCG@10 指标上都有了很大的提升。

例如，在 DL19 数据集上，使用 MonoT5 模型进行点-wise 重新排序后，nDCG@10 提升到了 0.7174；而使用 RankZephyr 等列表-wise 重新排序模型后，nDCG@10 更是达到了 0.7412。这表明 RankLLM 中的列表-wise 重新排序模型在处理复杂的文档排序任务时，能够更好地捕捉文档之间的关系，从而获得更好的排序效果。

不过，实验中也发现了一些问题。比如，一些即插即用的提示解码器（prompt-decoders）在生成模型响应时，可能会出现格式错误、重复或缺失文档等问题。例如，在使用 GPT4o-mini 模型时，有 28% 到 75% 的响应出现了候选文档 id 缺失的情况。不过，通过 RankLLM 的错误处理和响应后处理机制，这些问题得到了一定程度的缓解，使得这些模型仍然能够在重新排序任务中取得有竞争力的结果。

RankLLM 的优势和应用场景

RankLLM 的优势主要体现在以下几个方面：

• 统一的框架 ：它把各种不同的重新排序模型整合到一个统一的框架中，方便用户进行实验和比较。这对于研究人员来说，可以大大节省开发和测试的时间，提高研究效率。
• 灵活的配置 ：无论是模型选择、提示设计，还是推理框架，RankLLM 都提供了丰富的配置选项。用户可以根据自己的具体需求，灵活调整各个参数，以达到最佳的重新排序效果。
• 强大的可扩展性 ：由于 RankLLM 的模块化设计，用户可以很容易地添加新的模型、提示模板或推理后端。这使得 RankLLM 能够紧跟大型语言模型领域的发展步伐，不断集成新的技术成果。
• 详细的分析和评估工具 ：RankLLM 提供的调用分析和评估组件，可以帮助用户深入了解模型的行为，分析重新排序的结果，从而为进一步优化模型提供依据。

RankLLM 可以应用于多种信息检索相关的场景，比如搜索引擎结果的优化、问答系统中的文档检索和排序、推荐系统中的内容推荐等。在这些场景中，RankLLM 能够利用大型语言模型的强大能力，提升检索结果的相关性和用户体验。

RankLLM 的未来发展方向

RankLLM 目前已经取得了不少成果，但它的发展还远未停止。未来，RankLLM 的开发团队计划继续扩展支持的模型范围，集成更多先进的大型语言模型。同时，他们还将进一步优化工具包的性能，提高其在大规模数据处理和实时应用场景中的效率。

此外，RankLLM 还将加强与其他信息检索和检索增强生成（RAG）框架的整合，为用户提供优质的端到端解决方案。通过与社区的紧密合作，RankLLM 有望不断改进和完善，成为信息检索领域不可或缺的工具之一。

总结

RankLLM 是一个功能强大、灵活易用的 Python 工具包，它为利用大型语言模型进行文档重新排序提供了全面的解决方案。通过 RankLLM，无论是研究人员还是开发者，都可以方便地探索和应用各种重新排序模型，提升信息检索系统的性能。它不仅降低了进入这一领域的门槛，还推动了信息检索技术的发展和创新。

如果你对信息检索和大型语言模型感兴趣，不妨试试 RankLLM，相信它会给你带来意想不到的收获和惊喜。你可以访问 RankLLM 的官方网站 rankllm.ai，获取更多信息和资源，开始你的重新排序之旅吧！