使用RankLLM模型对搜索结果进行重新排序的方法

最新推荐文章于 2024-11-14 08:53:59 发布

qq_37836323

最新推荐文章于 2024-11-14 08:53:59 发布

阅读量419

点赞数 5

文章标签： python

本文链接：https://blog.csdn.net/qq_29929123/article/details/140302701

版权

使用RankLLM模型对搜索结果进行重新排序的方法

在自然语言处理（NLP）领域，重新排序（reranking）是提高搜索系统性能的重要技术。本文介绍了如何使用RankLLM模型对搜索结果进行重新排序。RankLLM提供了一套列表式重新排序器，重点在于为任务调整的开源LLM（如RankVicuna和RankZephyr）。

依赖项

RankLLM重新排序器需要CUDA支持
内置的检索器使用Pyserini，需要JDK11、PyTorch和Faiss

安装依赖

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

示例代码

以下是演示如何使用RankLLM对搜索结果进行重新排序的代码示例：

import nest_asyncio
nest_asyncio.apply()

import logging
import sys
import os
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
from pathlib import Path
import requests

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

# 设置OpenAI API
OPENAI_API_TOKEN = "sk-"
os.environ["OPENAI_API_KEY"] = OPENAI_API_TOKEN
from llama_index.core import Settings
Settings.llm = OpenAI(temperature=0, model="gpt-3.5-turbo")
Settings.chunk_size = 512

# 下载数据并构建索引
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)

# 检索和重新排序
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", gpt_model="gpt-3.5-turbo"):
    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 = RankLLMRerank(top_n=reranker_top_n, model=model, with_retrieval=with_retrieval, gpt_model=gpt_model)
        retrieved_nodes = reranker.postprocess_nodes(retrieved_nodes, query_bundle)

    return retrieved_nodes

def visualize_retrieved_nodes(nodes):
    result_dicts = []
    for node in nodes:
        result_dict = {"Score": node.score, "Text": node.node.get_text()}
        result_dicts.append(result_dict)
    display(HTML(pd.DataFrame(result_dicts).to_html().replace("\\n", "<br>")))

# 不进行重新排序的检索
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进行重新排序
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)

# 使用RankVicuna进行重新排序
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)

# 使用RankGPT进行重新排序
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="gpt", gpt_model="gpt-3.5-turbo")
visualize_retrieved_nodes(new_nodes)