[掌握检索器：如何为检索结果添加相似度分数]-优快云博客

本文链接：https://blog.youkuaiyun.com/sfsdhttdk/article/details/144273000

# 掌握检索器：如何为检索结果添加相似度分数

在处理大规模文档检索系统时，我们常常希望知道每个返回结果的相关性。尽管检索器会返回文档对象的序列，但它们通常不包含关于如何检索这些文档的信息，比如与查询的相似度分数。在本文中，我们将探讨如何为检索结果添加相似度分数，并将其添加到文档的元数据中。

## 主要内容

### 使用向量存储检索器

为向量存储检索器添加相似度分数，我们可以为相应的向量存储实现一个简单的包装函数。这里以`PineconeVectorStore`为例，代码可以兼容所有实现了`similarity_search_with_score`方法的LangChain向量存储。

```python
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方法，将分数添加到文档的元数据中。

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

# 使用API代理服务提高访问稳定性
result = retriever.invoke("dinosaur")
print(result)

使用高级LangChain检索器

对于诸如SelfQueryRetriever或MultiVectorRetriever之类的高级LangChain检索器，我们可以通过更新对应类的方法来实现这一功能。

SelfQueryRetriever

首先，按照指导创建一些元数据用于过滤，然后重载_get_docs_with_query方法，以便利用向量存储的similarity_search_with_score方法。

from langchain.chains.query_constructor.base import AttributeInfo
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain_openai import ChatOpenAI

metadata_field_info = [
    AttributeInfo(
        name="genre",
        description="The genre of the movie.",
        type="string",
    ),
    # 更多元数据字段
]

llm = ChatOpenAI(temperature=0)

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

MultiVectorRetriever

对于MultiVectorRetriever，我们可以重载_get_relevant_documents方法来传播相似度分数。

from collections import defaultdict
from langchain.retrievers import MultiVectorRetriever

class CustomMultiVectorRetriever(MultiVectorRetriever):
    def _get_relevant_documents(self, query: str, *, run_manager) -> 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

# 使用API代理服务提高访问稳定性
retriever = CustomMultiVectorRetriever(vectorstore=vectorstore, docstore=docstore)
retriever.invoke("cat")

常见问题和解决方案

网络限制问题：在某些地区，访问向量存储API可能会受到网络限制，您可以考虑使用API代理服务来提高访问稳定性。
性能问题：在大规模文档集合中，增加相似度分数的计算可能会影响性能。可以考虑对向量存储进行分片或者使用缓存机制。

总结和进一步学习资源

通过为检索结果添加相似度分数，我们可以更好地理解文档的相关性。为了深入了解，可以查看以下资源：

参考资料

LangChain 官方文档
Pinecone 官方文档

如果这篇文章对你有帮助，欢迎点赞并关注我的博客。您的支持是我持续创作的动力！

---END---