如何为Retriever结果添加检索评分_multiqueryretriever 文档排序-优快云博客

本文链接：https://blog.youkuaiyun.com/tt_jishu/article/details/143224381

如何为Retriever结果添加检索评分

在构建智能信息检索系统时，retriever通常会返回一系列的文档对象（Document）。这些对象默认不含任何关于检索过程的信息，例如与查询的相似度评分。在本文中，我们将探讨如何在文档的metadata中添加检索评分。

引言

在信息检索中，了解文档与查询的相似度评分有助于排序和筛选结果。本文将介绍如何通过LangChain库实现这一目标，分为以下两种情况：

从向量存储的retrievers中获取评分；
从LangChain的高阶retrievers中获取评分，如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()
)

从向量存储的Retriever中获取评分

我们可以通过包装底层向量存储的.similarity_search_with_score方法来实现这一点。

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方法来实现评分传递。

from langchain.chains.query_constructor.base import AttributeInfo
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(
    llm, vectorstore, document_content_description, metadata_field_info
)

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

从MultiVectorRetriever中获取评分

通过子类化MultiVectorRetriever并重写_get_relevant_documents方法，我们可以实现多向量检索评分。

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)