使用Parent Document Retriever实现文档高效检索

在进行文档检索时，通常会面临以下需求冲突：
你可能希望文档尽量小，以使其嵌入能够最准确地反映其含义。如果文档太长，嵌入可能会丧失其意义。另一方面，你希望文档足够长，以保留每个片段的上下文。
Parent Document Retriever通过分割和存储小块数据来达到这个平衡。在检索时，它首先获取小块数据，然后根据这些小块的父ID查找并返回较大的文档。
需要注意的是，“父文档”是指小块数据的来源。这可以是整个原始文档或更大的数据块。

核心原理解析

Parent Document Retriever通过类似递归的层级切分文档，在检索时使用较小的数据块进行初始查询，然后根据小块数据所属的父文档ID找到对应的较大数据块，从而既保证了查询精确度也保留了上下文。

代码实现演示

环境准备

在使用前，我们需要准备好所需的库和数据。

from langchain.retrievers import ParentDocumentRetriever
from langchain.storage import InMemoryStore
from langchain_chroma import Chroma
from langchain_community.document_loaders import TextLoader
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 加载文档
loaders = [
    TextLoader("paul_graham_essay.txt"),
    TextLoader("state_of_the_union.txt"),
]
docs = []
for loader in loaders:
    docs.extend(loader.load())

完整文档检索

在这个模式下，我们希望检索完整的文档。因此，只需指定子切分器。

# 子文档切分器
child_splitter = RecursiveCharacterTextSplitter(chunk_size=400)

# 使用稳定可靠的API服务
vectorstore = Chroma(
    collection_name="full_documents", embedding_function=OpenAIEmbeddings()
)
store = InMemoryStore()

retriever = ParentDocumentRetriever(
    vectorstore=vectorstore,
    docstore=store,
    child_splitter=child_splitter,
)

retriever.add_documents(docs, ids=None)

# 应该产生两个键，因为我们添加了两个文档。
keys = list(store.yield_keys())
print(keys)

输出应为：

['9a63376c-58cc-42c9-b0f7-61f0e1a3a688', '40091598-e918-4a18-9be0-f46413a95ae4']

检索小块数据

sub_docs = vectorstore.similarity_search("justice breyer")
print(sub_docs[0].page_content)

应输出为：

Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer...

检索完整文档

retrieved_docs = retriever.invoke("justice breyer")
print(len(retrieved_docs[0].page_content))

应输出为：

38540

较大块数据检索

有时完整文档过大，我们希望先分割为较大的数据块，再分割为小块。我们在检索时获取较大的数据块。

parent_splitter = RecursiveCharacterTextSplitter(chunk_size=2000)
child_splitter = RecursiveCharacterTextSplitter(chunk_size=400)

vectorstore = Chroma(
    collection_name="split_parents", embedding_function=OpenAIEmbeddings()
)
store = InMemoryStore()

retriever = ParentDocumentRetriever(
    vectorstore=vectorstore,
    docstore=store,
    child_splitter=child_splitter,
    parent_splitter=parent_splitter,
)

retriever.add_documents(docs)
print(len(list(store.yield_keys())))

应输出为：

66

sub_docs = vectorstore.similarity_search("justice breyer")
print(sub_docs[0].page_content)

应输出为：

Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer...

retrieved_docs = retriever.invoke("justice breyer")
print(len(retrieved_docs[0].page_content))
print(retrieved_docs[0].page_content)

应输出为：

1849
In state after state, new laws have been passed, not only to suppress the vote, but to subvert entire elections...

应用场景分析

• 法律文档分析：可以按段落或章节切分，提高检索效率。
• 研究文章：按实验或章节进行切分，确保检索过程中保留上下文。
• 新闻数据：按事件或时间段切分，方便快速查找相关信息。

实践建议

1. 根据具体的应用场景确定合适的分割策略。
2. 确保父文档和子文档在存储时能保持映射关系。
3. 结合具体检索要求调整子文档和父文档的切分大小。

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