技术前沿

RankLLM：RAG架构下通过重排序实现精准信息检索

2024 年 4 月 8 日

一、前言

在检索增强生成（Retrieval-Augmented Generation, RAG）的框架下，重排序（Re-Rank）阶段扮演着至关重要的角色。该阶段的目标是对初步检索得到的大量文档进行再次筛选和排序，以确保生成阶段能够优先利用最相关的信息。这一过程类似于传统信息检索中的两阶段排序策略，其中粗排（粗略排序）阶段追求高效的检索速度，而精排（精细排序）阶段则专注于提升结果的相关性和准确性。

在RAG架构中，重排序不仅仅是对文档列表的简单重新排列，它还涉及到对文档相关性的深入理解和评估。有效的重排序技术能够显著提升RAG系统的性能，确保生成的内容不仅准确，而且与用户查询高度相关。然而，现有的重排序方法面临着一系列挑战，包括对输入顺序的过度依赖、缺乏对动态网络信息的适应性，以及模型的不透明性和不可复现性。

针对这些挑战，开源的重排序方法 RankLLM 应运而生。RankLLM 利用大型语言模型（LLMs）的强大能力，通过零样本学习（zero-shot learning）的方式，无需特定任务的训练数据即可执行重排序任务。这种方法不仅提高了重排序的效率和效果，而且通过开源实现了模型的透明性和可访问性，为研究社区提供了一个可复现、可改进的重排序工具。

RankLLM 的引入为RAG架构的发展提供了新的动力，它不仅能够处理静态的、预先定义的查询，还能够适应实时变化的数据和查询，满足不断演变的信息需求。通过这种开源和高效的重排序策略，RankLLM有望成为未来RAG系统中不可或缺的组成部分，推动信息检索和自然语言处理领域的进一步发展。

二、RankLLM 介绍

RankLLM 是一种基于大型语言模型（Large Language Models, LLMs）的重排序方法，它利用了LLM的强大能力来改进信息检索的结果。在信息检索的过程中，初步检索阶段可能会返回大量相关或不相关的文档，RankLLM 的作用是在这些文档中进行再次排序，以提高检索结果的相关性和准确性。RankLLM通过使用LLM作为“提示-解码器”（prompt-decoder），在没有特定任务训练数据的情况下（即零样本设置），对文档列表进行重新排序，优化诸如归一化折扣累积增益（nDCG）等检索指标。

三、RankLLM 原理

RankLLM 重排序的原理基于零样本学习（zero-shot learning），它不需要特定任务的训练数据。RankLLM 使用一种提示（Prompt）来指导LLM如何对文档列表进行排序。这个提示描述了重排序任务，并提供了一个格式，让模型知道如何生成排序后的文档列表。RankLLM 利用了 LLM 的理解和生成能力，通过指令微调（instruction fine-tuning）来学习如何根据相关性对文档进行排序。

在实际操作中，RankLLM接收一系列候选文档和相关的查询，然后使用LLM生成一个按相关性排序的文档列表。这个过程涉及到从大型数据集中提取有用的信息，并将这些信息融入到LLM的生成过程中，以实现更准确的重排序。

四、RankLLM 实现

如何使用 RankLLM 工具对搜索结果进行重新排序，以提升结果的相关性和准确性。RankLLM 提供了一系列专门针对这项任务进行优化的开源大语言模型，例如 RankVicuna 和 RankZephyr。

本文将通过比较 Van Gogh Wiki 页面的搜索结果，分别使用了传统的检索方法和结合 RankLLM 的检索方法。一方面是仅利用检索（通过 llama-index 的 VectorIndexRetriever 实现），另一方面则是结合了 RankLLM 进行的检索加重排序。演示中展现了 RankLLM 的两种模型：

RankVicuna 7B V1
RankZephyr 7B V1 - Full - BF16

依赖项：

目前，RankLLM 的重排序功能需要 CUDA 环境，并且必须安装 rank_llm（通过命令 pip install rank_llm 安装）。
而内置的检索工具则依赖于 Pyserini，需要 JDK11、PyTorch 和 Faiss 环境。

castorini/rank_llm 是一个专门用于利用大语言模型（如 GPT3.5, GPT4, Vicuna 和 Zephyr）进行提示解码的代码库。

4.1、安装依赖

%pip install llama-index-core
%pip install llama-index-llms-openai
%pip install llama-index-postprocessor-rankllm-rerank
%pip install rank-llm

import nest_asyncio

# 启用异步I/O
nest_asyncio.apply()

import logging
import sys

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.postprocessor import LLMRerank
from llama_index.llms.openai import OpenAI
from IPython.display import Markdown, display

import os

OPENAI_API_TOKEN = "sk-"
os.environ["OPENAI_API_KEY"] = OPENAI_API_TOKEN

4.2、加载数据，构建索引

从维基百科（Wikipedia）获取Vincent van Gogh的文本内容，并将这些文本内容保存到本地文件中，最后构建一个索引以便后续的数据检索。

from pathlib import Path
import requests

wiki_titles = [
    "Vincent van Gogh",
]


data_path = Path("data_wiki")

for title in wiki_titles:
    response = requests.get(
        "https://en.wikipedia.org/w/api.php",
        params={
            "action": "query",
            "format": "json",
            "titles": title,
            "prop": "extracts",
            "explaintext": True,
        },
    ).json()
    page = next(iter(response["query"]["pages"].values()))
    wiki_text = page["extract"]

    if not data_path.exists():
        Path.mkdir(data_path)

    with open(data_path / f"{title}.txt", "w") as fp:
        fp.write(wiki_text)

# 加载文档
documents = SimpleDirectoryReader("./data_wiki/").load_data()

# 构建索引
index = VectorStoreIndex.from_documents(
    documents,
)

五、检索与 RankLLM 重排序流程

准备检索器和重排序器

在开始检索流程之前，首先需要准备两个关键组件：检索器（Retriever）和重排序器（Reranker）。检索器负责从大量文档中快速筛选出与查询相关的候选文档集合。常见的检索器包括基于词频的BM25算法、基于神经网络的SPLADE++ ED和RepLLaMA等。这些检索器可以是无监督的、监督的、稀疏的、密集的，或者是它们的混合形式。重排序器则进一步优化检索结果，通过深入理解文档内容和用户查询意图，对候选文档进行精细排序。

对给定的搜索查询进行检索,不加重排序

用户发起搜索查询后，检索器首先在文档集合中进行初步检索，生成一个未经排序的候选文档列表。这一步骤不涉及任何形式的排序，目的是尽可能广泛地覆盖与查询相关的文档，为后续的重排序提供原材料。

使用 RankZephyr 进行重排序后的搜索查询检索

接下来，利用 RankZephyr 对初步检索得到的候选文档列表进行重排序。RankZephyr 是一种基于开源大型语言模型（LLM）的重排序工具，它通过零样本学习（zero-shot learning）的方式，无需特定任务的训练数据即可执行重排序任务。RankZephyr通过指令微调（instruction fine-tuning）来学习如何根据相关性对文档进行排序，能够根据归一化折扣累积增益（nDCG）等检索指标优化文档排序。

使用 RankVicuna 进行重排序后的搜索查询检索

为了进一步提升检索效果，可以采用 RankVicuna 对 RankZephyr 的输出结果进行二次重排序。RankVicuna 是另一种基于开源LLM的重排序工具，它在 RankZephyr 的基础上进一步优化，通过多阶段的重排序策略（progressive reranking）来迭代改进排序结果。RankVicuna 通过引入变量窗口大小和随机输入顺序的训练，增强了模型对于不同重排序场景的适应性和鲁棒性。

通过上述流程，RankLLM 能够有效地结合检索器的快速筛选能力和重排序器的深度理解能力，为用户提供更加精准和相关的搜索结果。

from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core import QueryBundle
from llama_index.postprocessor.rankLLM_rerank import RankLLMRerank

import pandas as pd
from IPython.display import display, HTML


def get_retrieved_nodes(
    query_str,
    vector_top_k=10,
    reranker_top_n=3,
    with_reranker=False,
    with_retrieval=False,
    model="zephyr",
):
    query_bundle = QueryBundle(query_str)
    # 配置检索器
    retriever = VectorIndexRetriever(
        index=index,
        similarity_top_k=vector_top_k,
    )
    retrieved_nodes = retriever.retrieve(query_bundle)

    if with_reranker:
        # 配置reranker
        reranker = RankLLMRerank(
            top_n=reranker_top_n, with_retrieval=with_retrieval, model=model
        )
        retrieved_nodes = reranker.postprocess_nodes(
            retrieved_nodes, query_bundle
        )

    return retrieved_nodes


def pretty_print(df):
    return display(HTML(df.to_html().replace("\\n", "")))


def visualize_retrieved_nodes(nodes) -> None:
    result_dicts = []
    for node in nodes:
        result_dict = {"Score": node.score, "Text": node.node.get_text()}
        result_dicts.append(result_dict)

    pretty_print(pd.DataFrame(result_dicts))

不加重排序的检索前三结果：

预期结果显示：

经过梵高的多次恳求，高更于10月23日抵达阿尔勒，并于11月两人一起作画。高更在他的《向日葵画家》中描绘了梵高。

new_nodes = get_retrieved_nodes(
    "Which date did Paul Gauguin arrive in Arles?",
    vector_top_k=3,
    with_reranker=False,
    model="zephyr",
)

visualize_retrieved_nodes(new_nodes)

正确的结果在未重排序时排名第三。

使用 RankZephyr 重排序前 10 结果并返回前 3：

new_nodes = get_retrieved_nodes(
    "Which date did Paul Gauguin arrive in Arles?",
    vector_top_k=10,
    reranker_top_n=3,
    with_reranker=True,
    with_retrieval=False,
    model="zephyr",
)

visualize_retrieved_nodes(new_nodes)

在 RankZephyr 重排序后，正确的结果升至第一位。

使用 RankVicuna 重排序前 10 结果并返回前 3：

new_nodes = get_retrieved_nodes(
    "Which date did Paul Gauguin arrive in Arles?",
    vector_top_k=10,
    reranker_top_n=3,
    with_reranker=True,
    with_retrieval=False,
    model="vicuna",
)

visualize_retrieved_nodes(new_nodes)

在 RankVicuna 重排序后，正确的结果同样位居第一位。

六、总结

RankLLM是一种创新的重排序方法，它利用了大型语言模型的能力来改进信息检索的结果。通过零样本学习和指令微调，RankLLM能够在不需要特定任务训练数据的情况下，有效地对文档列表进行重新排序。这种方法不仅提高了检索结果的相关性和准确性，而且由于其开源性质，还增加了研究的可重复性和模型的可访问性。RankLLM的引入为解决RAG架构中的一些关键挑战提供了新的可能性，为未来的信息检索和文本生成任务开辟了新的道路。

七、References

[1]. RankLLM Github：https://github.com/castorini/rank_llm

[2]. Pyserini： https://github.com/castorini/pyserini

[3]. Ronak Pradeep, Sahel Sharifymoghaddam, Jimmy Lin, R. (2023). RankZephyr: Effective and Robust Zero-Shot Listwise Reranking is a Breeze: https://arxiv.org/abs/2312.02724

立即咨询 JOTO

JOTO 提供覆盖企业智能体规划与搭建、AI 平台私有化部署、RAG 知识工程、AI 安全治理、FDE 驻场共创及持续运营优化的全周期 AI 落地服务，帮助企业把验证中的 AI 能力转化为安全、可控、可持续迭代的生产力。联系 JOTO 获取 AI 落地咨询

想把这些做法用到你的业务里?

留下你的场景和痛点,我们帮你判断从哪一步开始。

联系我们

RankLLM：RAG架构下通过重排序实现精准信息检索

4.2、加载数据，构建索引

立即咨询 JOTO

想把这些做法用到你的业务里?

开启企业级 AI 落地

填写需求单