如何为Retriever结果添加相似度评分

在信息检索系统中,Retriever通常会返回一系列的Document对象,但默认情况下,这些对象不包含关于检索过程的信息,比如与查询的相似度评分。在本文中,我们将演示如何将检索评分添加到文档的metadata中,主要涉及以下两种情况:

  1. 从向量存储检索器;
  2. 从更高阶的LangChain检索器,如SelfQueryRetriever或MultiVectorRetriever。

对于第一种情况,我们将为相应的向量存储实现一个简短的包装函数。对于第二种情况,我们将更新相应类的方法。

创建向量存储

首先,我们需要使用一些数据填充一个向量存储。在此示例中,我们将使用PineconeVectorStore,不过本文的指南适用于任何实现了similarity_search_with_score方法的LangChain向量存储。

from langchain_core.documents import Document
from langchain_openai import OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore

# 定义文档数据
docs = [
    Document(
        page_content="A bunch of scientists bring back dinosaurs and mayhem breaks loose",
        metadata={"year": 1993, "rating": 7.7, "genre": "science fiction"},
    ),
    # 其他文档...
]

# 创建并填充向量存储
vectorstore = PineconeVectorStore.from_documents(
    docs, index_name="sample", embedding=OpenAIEmbeddings()
)

为向量存储检索器添加评分

为了从向量存储检索器获得评分,我们将包装底层向量存储的similarity_search_with_score方法,创建一个短函数以将评分打包到相应文档的metadata中。

from typing import List
from langchain_core.documents import Document
from langchain_core.runnables import chain

@chain
def retriever(query: str) -> List[Document]:
    # 执行带有评分的相似性搜索
    docs, scores = zip(*vectorstore.similarity_search_with_score(query))
    for doc, score in zip(docs, scores):
        doc.metadata["score"] = score

    return docs

# 调用检索器
result = retriever.invoke("dinosaur")
print(result)  # 输出添加了评分的文档

SelfQueryRetriever

对于SelfQueryRetriever,我们需要重写_get_docs_with_query方法以使用向量存储的similarity_search_with_score。下面是如何实现这一点的:

from typing import Any, Dict
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain_openai import ChatOpenAI

class CustomSelfQueryRetriever(SelfQueryRetriever):
    def _get_docs_with_query(
        self, query: str, search_kwargs: Dict[str, Any]
    ) -> List[Document]:
        docs, scores = zip(*vectorstore.similarity_search_with_score(query, **search_kwargs))
        for doc, score in zip(docs, scores):
            doc.metadata["score"] = score
        return docs

# 实例化并调用自定义检索器
retriever = CustomSelfQueryRetriever.from_llm(
    ChatOpenAI(temperature=0), vectorstore, "Brief summary of a movie", metadata_field_info
)

result = retriever.invoke("dinosaur movie with rating less than 8")
print(result)

MultiVectorRetriever

为了在MultiVectorRetriever中传递相似度评分,我们需要重写_get_relevant_documents方法,将评分添加到子文档的metadata,并在检索到的父文档的metadata中包含这些子文档列表。

from collections import defaultdict
from langchain.retrievers import MultiVectorRetriever
from langchain_core.callbacks import CallbackManagerForRetrieverRun

class CustomMultiVectorRetriever(MultiVectorRetriever):
    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun) -> List[Document]:
        results = self.vectorstore.similarity_search_with_score(query, **self.search_kwargs)
        id_to_doc = defaultdict(list)
        for doc, score in results:
            doc_id = doc.metadata.get("doc_id")
            if doc_id:
                doc.metadata["score"] = score
                id_to_doc[doc_id].append(doc)
        
        docs = []
        for _id, sub_docs in id_to_doc.items():
            docstore_docs = self.docstore.mget([_id])
            if docstore_docs:
                if doc := docstore_docs[0]:
                    doc.metadata["sub_docs"] = sub_docs
                    docs.append(doc)
        
        return docs

# 实例化并调用自定义多向量检索器
retriever = CustomMultiVectorRetriever(vectorstore=vectorstore, docstore=docstore)
result = retriever.invoke("cat")
print(result)

以上代码展示了如何将相似度评分添加到retriever结果的metadata中,并将其应用于不同类型的检索器。通过这种方式,我们可以更好地了解文档的相关性,提高信息检索的效果。

如果遇到问题欢迎在评论区交流。
—END—

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值