譯者 | 陳峻

審校 | 重樓

作為一種能在先進(jìn)的 RAG（Retrieval-Augmented Generation，檢索增強(qiáng)生成）模型中實(shí)現(xiàn)的方法，父文檔檢索（Parent Document Retrieval，PDR）旨在恢復(fù)那些可以從中提取到相關(guān)子段落（或片段）的完整父文檔。此類文檔通過將豐富的上下文，傳遞給 RAG 模型，以便對復(fù)雜或細(xì)微的問題，做出更全面且內(nèi)涵豐富的回答。通常，在 RAG 模型中檢索出父文檔的主要步驟包括：

• 數(shù)據(jù)預(yù)處理：將冗長的文檔分解為多個(gè)可管理的片段。
• 創(chuàng)建嵌入：將片段轉(zhuǎn)換為數(shù)值向量，以實(shí)現(xiàn)高效的搜索。
• 用戶查詢：讓用戶提交問題。
• 塊檢索：模型檢索出那些與查詢嵌入最為相似的部分。
• 查找父文檔：檢索原始文檔或從中獲取更大的片段。
• 父文檔檢索：檢索完整的父文檔，為響應(yīng)提供更為豐富的上下文。

循序漸進(jìn)的實(shí)施

根據(jù)上圖，我們可以將實(shí)現(xiàn)父文檔檢索的步驟簡單分為如下四個(gè)不同的階段：

1. 準(zhǔn)備數(shù)據(jù)

我們首先應(yīng)為自己的 RAG 系統(tǒng)創(chuàng)建環(huán)境并預(yù)處理數(shù)據(jù)，以便對后續(xù)的父文檔開展文檔檢索。

（1）導(dǎo)入必要的模塊

我們將從已安裝的庫中導(dǎo)入所需的模塊，以設(shè)置我們的 PDR 系統(tǒng)：

Python

from langchain.schema import Document
from langchain.vectorstores import Chroma
from langchain.retrievers import ParentDocumentRetriever
from langchain.chains import RetrievalQA
from langchain_openai import OpenAI
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.storage import InMemoryStore
from langchain.document_loaders import TextLoader
from langchain.embeddings.openai import OpenAIEmbeddings

上述這些庫和模塊正是構(gòu)成整個(gè)過程步驟的主要部分。

（2）設(shè)置 OpenAI API 密鑰

接著，我們使用 OpenAI LLM來生成響應(yīng)，為此我們需要一個(gè) OpenAI 的API 密鑰。該密鑰可被用來設(shè)置環(huán)境變量：OPENAI_API_KEY。

Python

OPENAI_API_KEY = os.environ["OPENAI_API_KEY"] = ""  # Add your OpenAI API key
if OPENAI_API_KEY == "":
raise ValueError("Please set the OPENAI_API_KEY environment variable")

（3）定義文本嵌入函數(shù)

通過如下方式，我們利用 OpenAI 的嵌入來表示文本數(shù)據(jù)：

Python

embeddings = OpenAIEmbeddings()

（4）加載文本數(shù)據(jù)

為了讀取想要檢索的文本文檔，你可以利用類TextLoader來讀取文本文件：

Python

loaders = [
    TextLoader('/path/to/your/document1.txt'),
    TextLoader('/path/to/your/document2.txt'),
]
docs = []
for l in loaders:
docs.extend(l.load())

2. 檢索完整的文檔

下面，我們將通過設(shè)置系統(tǒng)，來檢索與子段落相關(guān)的完整父文檔。

（1）完整文檔的拆分

我們使用RecursiveCharacterTextSplitter將加載的文檔分割成比所需大小更小的文本塊。這些子文檔將使我們能夠有效地搜索相關(guān)段落：

Python

child_splitter = RecursiveCharacterTextSplitter(chunk_size=400)

（2）矢量存儲(chǔ)和存儲(chǔ)設(shè)置

下面，我們將使用Chroma向量存儲(chǔ)來嵌入各個(gè)子文檔，并利用InMemoryStore來跟蹤與子文檔關(guān)聯(lián)的完整父文檔： 

Python

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

（3）父文檔檢索器

接著，讓我們從類ParentDocumentRetriever中實(shí)例化一個(gè)對象。該類主要負(fù)責(zé)完整父文檔與基于子文檔相似性檢索相關(guān)的核心邏輯。

Python

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

（4）添加文檔

然后，這些加載的文檔將使用add_documents方法被饋入ParentDocumentRetriever中，如下代碼所示：

Python

full_doc_retriever.add_documents(docs)
print(list(store.yield_keys()))  # List document IDs in the store

（5）相似性搜索和檢索

至此，檢索器已基本實(shí)現(xiàn)，你可以在給定查詢的情況下，去檢索相關(guān)的子文檔，并獲取相關(guān)的完整父文檔：

Python

sub_docs = vectorstore.similarity_search("What is LangSmith?", k=2)
print(len(sub_docs))
print(sub_docs[0].page_content)  
retrieved_docs = full_doc_retriever.invoke("What is LangSmith?")
print(len(retrieved_docs[0].page_content)) 
print(retrieved_docs[0].page_content)

3. 檢索更大的數(shù)據(jù)塊

有時(shí)，在文檔非常大的情況下，我們可能無法獲取完整的父文檔。對此，可參考如下從父文檔獲取較大片段的方法：

塊和父級(jí)文本分割：

• 使用RecursiveCharacterTextSplitter的兩個(gè)實(shí)例，其中一個(gè)用于創(chuàng)建特定大小的較大父文檔。另一個(gè)具有較小的塊大小，可用于創(chuàng)建文本片段，即父文檔中的子文檔。

矢量存儲(chǔ)和存儲(chǔ)設(shè)置（類似完整的文檔檢索）：

• 創(chuàng)建一個(gè)向量存儲(chǔ)Chroma，用于索引子文檔的嵌入。
• 使用InMemoryStore保存父文檔的塊。

（1）父文檔檢索器

該檢索器可解決 RAG 中的一個(gè)基本問題：由于被檢索的整個(gè)文檔過大，而無法包含足夠的上下文。為此，RAG需將文檔切成小塊進(jìn)行檢索，進(jìn)而對這些小塊進(jìn)行索引。不過，在查詢之后，它不會(huì)去檢索這些文檔片段，而是檢索整個(gè)父文檔，從而為后續(xù)的生成提供更為豐富的上下文。

Python

parent_splitter = RecursiveCharacterTextSplitter(chunk_size=2000)  
child_splitter = RecursiveCharacterTextSplitter(chunk_size=400)   
vectorstore = Chroma(
    collection_name="split_parents",
    embedding_function=OpenAIEmbeddings()
)
store = InMemoryStore()
big_chunks_retriever = ParentDocumentRetriever(
    vectorstore=vectorstore,
    docstore=store,
    child_splitter=child_splitter,
    parent_splitter=parent_splitter
)
# Adding documents
big_chunks_retriever.add_documents(docs)
print(len(list(store.yield_keys())))  # List document IDs in the store

（2）相似性搜索和檢索

該過程仍然與完整的文檔檢索類似，我們需要查找相關(guān)的子文檔，然后從父文檔中獲取相應(yīng)的更大文檔塊。

Python

sub_docs = vectorstore.similarity_search("What is LangSmith?", k=2)
print(len(sub_docs))
print(sub_docs[0].page_content)  
retrieved_docs = big_chunks_retriever.invoke("What is LangSmith?")
print(len(retrieved_docs))
print(len(retrieved_docs[0].page_content)) 
print(retrieved_docs[0].page_content)

4. 與 RetrievalQA 集成

至此，我們已經(jīng)實(shí)現(xiàn)了一個(gè)父文檔檢索器，你可以將其與RetrievalQA鏈集成，以使用檢索到的父文檔進(jìn)行各種問答：

Python

qa = RetrievalQA.from_chain_type(llm=OpenAI(),
                                chain_type="stuff",
                                retriever=big_chunks_retriever)
query = "What is LangSmith?"
response = qa.invoke(query)
print(response)

小結(jié)

綜上所述，PDR 大幅提高了 RAG 模型輸出響應(yīng)的準(zhǔn)確性，而且這些響應(yīng)都帶有豐富的上下文。而通過對父文檔的全文檢索，我們可以深入準(zhǔn)確地回答各種復(fù)雜問題，這也是復(fù)雜人工智能的基本要求。

譯者介紹

陳峻（Julian Chen），51CTO社區(qū)編輯，具有十多年的IT項(xiàng)目實(shí)施經(jīng)驗(yàn)，善于對內(nèi)外部資源與風(fēng)險(xiǎn)實(shí)施管控，專注傳播網(wǎng)絡(luò)與信息安全知識(shí)與經(jīng)驗(yàn)。

原文標(biāo)題：Parent Document Retrieval (PDR): Useful Technique in RAG，作者：Intiaz Shaik

鏈接：https://dzone.com/articles/parent-document-retrieval-useful-technique-in-rag。