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

一、前言

本文将通过比较 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 环境。

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、加载数据，构建索引

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

1. 准备检索器和重排序器

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

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

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

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)

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)

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)

[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