FlashRank 的 re-ranking 在搜索与检索中的应用

老铁们，今天我们来聊聊在搜索与检索管道中如何利用 FlashRank 进行 re-ranking。FlashRank 是一个超轻量级且超快的 Python 库，适用于为现有的检索管道增加 re-ranking 能力。它基于最先进的 cross-encoders，非常感谢所有模型开发者的贡献。今天，我会通过一个笔记本来展示如何用 FlashRank 进行文档压缩与检索。

1. 技术背景介绍

在面对海量信息时，如何快速检索并获得最相关的结果，始终是一个重要的挑战。传统的向量检索技术能帮助我们缩小信息范围，但它们无法保证结果的最佳相关性。这时候，re-ranking 技术就派上用场了。FlashRank 则提供了一种超轻便、快速的方案，可以很方便地集成到现有系统中。

2. 原理深度解析

FlashRank 核心是利用最新的 cross-encoders 模型来重新排序检索结果，提高检索的精度。说白了，就是在初步的检索结果中再进行一轮排序，通过对上下文的深入理解来确定哪些才是最符合查询需求的内容。

3. 实战代码演示

首先，我们需要安装 FlashRank 和 faiss：

%pip install --upgrade --quiet flashrank
%pip install --upgrade --quiet faiss

# OR  (depending on Python version)

%pip install --upgrade --quiet faiss_cpu

接着，我们构建一个基础的向量检索器，并将 2023 年的国情咨文以块的形式存储：

import getpass
import os

os.environ["OPENAI_API_KEY"] = getpass.getpass()

from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

documents = TextLoader(
    "../../how_to/state_of_the_union.txt",
).load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
texts = text_splitter.split_documents(documents)
for idx, text in enumerate(texts):
    text.metadata["id"] = idx

embedding = OpenAIEmbeddings(model="text-embedding-ada-002")
retriever = FAISS.from_documents(texts, embedding).as_retriever(search_kwargs={"k": 20})

query = "What did the president say about Ketanji Brown Jackson"
docs = retriever.invoke(query)

# Helper function for printing docs
def pretty_print_docs(docs):
    print(
        f"\n{'-' * 100}\n".join(
            [
                f"Document {i+1}:\n\n{d.page_content}\nMetadata: {d.metadata}"
                for i, d in enumerate(docs)
            ]
        )
    )

pretty_print_docs(docs)

4. 优化建议分享

在我们构建的基础检索器之上，我们可以用 FlashRank 进行 re-ranking，以获得更精准的结果。这里我们使用 ContextualCompressionRetriever 来将 FlashRank 作为压缩器：

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import FlashrankRerank
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(temperature=0)

compressor = FlashrankRerank()
compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor, base_retriever=retriever
)

compressed_docs = compression_retriever.invoke(
    "What did the president say about Ketanji Jackson Brown"
)
print([doc.metadata["id"] for doc in compressed_docs])

经过 re-ranking，顶级结果的排列与初步检索结果完全不同，再次展示结果：

pretty_print_docs(compressed_docs)

5. 补充说明和总结

为进一步提高查询结果的相关性，我们还可以将 FlashRank 与 Retrieval QA 搭配使用：

from langchain.chains import RetrievalQA

chain = RetrievalQA.from_chain_type(llm=llm, retriever=compression_retriever)
query_result = chain.invoke(query)

print(query_result)

通过这种方式，我们可以获得更精准的答复：

{
 'query': 'What did the president say about Ketanji Brown Jackson',
 'result': "The President mentioned that Ketanji Brown Jackson is one of the nation's top legal minds and will continue Justice Breyer's legacy of excellence."
}

今天的技术分享就到这里，希望对大家有帮助。开发过程中遇到问题也可以在评论区交流~

—END—