在信息检索系统中,Retriever通常会返回一系列的Document对象,但默认情况下,这些对象不包含关于检索过程的信息,比如与查询的相似度评分。在本文中,我们将演示如何将检索评分添加到文档的metadata中,主要涉及以下两种情况:
- 从向量存储检索器;
- 从更高阶的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—