譯者 | 陳峻
審校 | 重樓
作為一種能在先進的 RAG(Retrieval-Augmented Generation,檢索增強生成)模型中實現的方法,父文檔檢索(Parent Document Retrieval,PDR)旨在恢復那些可以從中提取到相關子段落(或片段)的完整父文檔。此類文檔通過將豐富的上下文,傳遞給 RAG 模型,以便對復雜或細微的問題,做出更全面且內涵豐富的回答。通常,在 RAG 模型中檢索出父文檔的主要步驟包括:
- 數據預處理:將冗長的文檔分解為多個可管理的片段。
- 創建嵌入:將片段轉換為數值向量,以實現高效的搜索。
- 用戶查詢:讓用戶提交問題。
- 塊檢索:模型檢索出那些與查詢嵌入最為相似的部分。
- 查找父文檔:檢索原始文檔或從中獲取更大的片段。
- 父文檔檢索:檢索完整的父文檔,為響應提供更為豐富的上下文。
循序漸進的實施
根據上圖,我們可以將實現父文檔檢索的步驟簡單分為如下四個不同的階段:
1. 準備數據
我們首先應為自己的 RAG 系統創建環境并預處理數據,以便對后續的父文檔開展文檔檢索。
(1)導入必要的模塊
我們將從已安裝的庫中導入所需的模塊,以設置我們的 PDR 系統:
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
上述這些庫和模塊正是構成整個過程步驟的主要部分。
(2)設置 OpenAI API 密鑰
接著,我們使用 OpenAI LLM來生成響應,為此我們需要一個 OpenAI 的API 密鑰。該密鑰可被用來設置環境變量: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)定義文本嵌入函數
通過如下方式,我們利用 OpenAI 的嵌入來表示文本數據:
Python
embeddings = OpenAIEmbeddings()
(4)加載文本數據
為了讀取想要檢索的文本文檔,你可以利用類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. 檢索完整的文檔
下面,我們將通過設置系統,來檢索與子段落相關的完整父文檔。
(1)完整文檔的拆分
我們使用RecursiveCharacterTextSplitter將加載的文檔分割成比所需大小更小的文本塊。這些子文檔將使我們能夠有效地搜索相關段落:
Python
child_splitter = RecursiveCharacterTextSplitter(chunk_size=400)
(2)矢量存儲和存儲設置
下面,我們將使用Chroma向量存儲來嵌入各個子文檔,并利用InMemoryStore來跟蹤與子文檔關聯的完整父文檔:
Python
vectorstore = Chroma(
collection_name="full_documents",
embedding_function=OpenAIEmbeddings()
)
store = InMemoryStore()
(3)父文檔檢索器
接著,讓我們從類ParentDocumentRetriever中實例化一個對象。該類主要負責完整父文檔與基于子文檔相似性檢索相關的核心邏輯。
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)相似性搜索和檢索
至此,檢索器已基本實現,你可以在給定查詢的情況下,去檢索相關的子文檔,并獲取相關的完整父文檔:
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. 檢索更大的數據塊
有時,在文檔非常大的情況下,我們可能無法獲取完整的父文檔。對此,可參考如下從父文檔獲取較大片段的方法:
塊和父級文本分割:
- 使用RecursiveCharacterTextSplitter的兩個實例,其中一個用于創建特定大小的較大父文檔。另一個具有較小的塊大小,可用于創建文本片段,即父文檔中的子文檔。
矢量存儲和存儲設置(類似完整的文檔檢索):
- 創建一個向量存儲Chroma,用于索引子文檔的嵌入。
- 使用InMemoryStore保存父文檔的塊。
(1)父文檔檢索器
該檢索器可解決 RAG 中的一個基本問題:由于被檢索的整個文檔過大,而無法包含足夠的上下文。為此,RAG需將文檔切成小塊進行檢索,進而對這些小塊進行索引。不過,在查詢之后,它不會去檢索這些文檔片段,而是檢索整個父文檔,從而為后續的生成提供更為豐富的上下文。
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)相似性搜索和檢索
該過程仍然與完整的文檔檢索類似,我們需要查找相關的子文檔,然后從父文檔中獲取相應的更大文檔塊。
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 集成
至此,我們已經實現了一個父文檔檢索器,你可以將其與RetrievalQA鏈集成,以使用檢索到的父文檔進行各種問答:
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)
小結
綜上所述,PDR 大幅提高了 RAG 模型輸出響應的準確性,而且這些響應都帶有豐富的上下文。而通過對父文檔的全文檢索,我們可以深入準確地回答各種復雜問題,這也是復雜人工智能的基本要求。
譯者介紹
陳峻(Julian Chen),51CTO社區編輯,具有十多年的IT項目實施經驗,善于對內外部資源與風險實施管控,專注傳播網絡與信息安全知識與經驗。
原文標題:Parent Document Retrieval (PDR): Useful Technique in RAG,作者:Intiaz Shaik
鏈接:https://dzone.com/articles/parent-document-retrieval-useful-technique-in-rag。