使用 Google Cloud Vertex AI Reranker 优化搜索结果

在现代信息检索系统中，通过初步搜索获取候选文档后，如何有效地提升搜索结果的相关性是一个重要的挑战。Google Cloud 提供的 Vertex AI Reranker API 能够根据查询的相关性重新排序文档列表，从而提高搜索结果的精确度。这篇文章将带您深入了解这项技术的原理，并通过实战代码演示如何应用。

技术背景介绍

Vertex AI Reranker API 是 Vertex AI Agent Builder 中的一项独立 API。它不同于仅考虑语义相似度的嵌入技术，能够为文档和查询间的匹配程度提供精确的评分。而且，该 API 是无状态的，不需要在调用前对文档进行索引，这使它非常适合优化从任何文档检索器中获取的初步文档。

核心原理解析

Reranker API 的主要功能是根据查询的上下文和内容为每个候选文档生成一个相关性得分。然后根据该得分重新排列文档，从而确保最高相关性的文档在搜索结果的顶部。通过大量的训练数据，Vertex AI Reranker 可以捕获更复杂的文本特征以进行更为准确的评估。

代码实现演示

下面我们演示如何使用 Vertex AI Reranker 对搜索结果进行优化。请确保您已安装必要的 Python 库。

# 安装所需库
%pip install --upgrade --quiet langchain langchain-community langchain-google-community langchain-google-community[vertexaisearch] langchain-google-vertexai langchain-chroma langchain-text-splitters

# 初始化GCP项目
from google.cloud import aiplatform
PROJECT_ID = "your-gcp-project-id"
REGION = "your-region-id"
RANKING_LOCATION_ID = "global"

aiplatform.init(project=PROJECT_ID, location=REGION)

# 加载和预处理数据
from langchain_community.document_loaders import WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_google_vertexai import VertexAIEmbeddings
from langchain_chroma import Chroma

loader = WebBaseLoader("https://en.wikipedia.org/wiki/Google")
data = loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=5)
splits = text_splitter.split_documents(data)

embedding = VertexAIEmbeddings(model_name="textembedding-gecko@003")
vectordb = Chroma.from_documents(documents=splits, embedding=embedding)

# 使用基本检索器和Reranker进行检索
from langchain.retrievers.contextual_compression import ContextualCompressionRetriever
from langchain_google_community.vertex_rank import VertexAIRank

reranker = VertexAIRank(
    project_id=PROJECT_ID,
    location_id=RANKING_LOCATION_ID,
    ranking_config="default_ranking_config",
    title_field="source",
    top_n=5,
)

basic_retriever = vectordb.as_retriever(search_kwargs={"k": 5})
retriever_with_reranker = ContextualCompressionRetriever(
    base_compressor=reranker, base_retriever=basic_retriever
)

# 测试Reranker
query = "how did the name google originate?"
retrieved_docs = basic_retriever.invoke(query)
reranked_docs = retriever_with_reranker.invoke(query)

# 比较未排序和排序后的文档
import pandas as pd

unranked_docs_content = [docs.page_content for docs in retrieved_docs]
ranked_docs_content = [docs.page_content for docs in reranked_docs]

comparison_df = pd.DataFrame(
    {
        "Unranked Documents": unranked_docs_content,
        "Ranked Documents": ranked_docs_content,
    }
)

print(comparison_df)