iT邦幫忙

2024 iThome 鐵人賽

DAY 17
0

每天的專案會同步到 GitLab 上,可以前往 GitLab 查看,有興趣的朋友歡迎留言 or 來信討論,我的信箱是 nickchen1998@gmail.com

什麼是 ParentDocumentRetriever?

ParentDocumentRetriever 的用途是平衡文件拆分和檢索的需求,適用於當文件過大時,無法直接將整個文件進行檢索的情境。它的主要功能是先將文件拆分成較小的片段來建立向量索引,但在檢索時,會回傳這些小片段的父文件(即原本較大的文件或一部分)。

為什麼我們需要 ParentDocumentRetriever?

與昨天的 metadata 相反,ParentDocumentRetriever 主要是想要解決文本內文不足的問題,由於過長的 chunk 可能會導致搜尋精度降低,因此時常需要將文件拆分成較小的片段來進行檢索。但這樣做可能會使得檢索結果過於分散,無法有效地回答使用者的問題,即便有使用 Recursion Strategy 也仍舊只能保留有限的文本,然而隨著語言模型能支援的 token 長度越來越大,也相對的可以給予更多的內容讓語言模型進行摘要回答。

使用 ParentDocumentRetriever 寫入向量資料

from env_settings import EnvSettings, BASE_DIR
from langchain.retrievers import ParentDocumentRetriever
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain_community.document_loaders import PyMuPDFLoader
from pinecone import Pinecone
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.storage import InMemoryStore


# 準備兩個 Pinecone 連線
env_settings = EnvSettings()
vector_store = PineconeVectorStore(
    index=Pinecone(api_key=env_settings.PINECONE_API_KEY).Index("chunks"),
    embedding=OpenAIEmbeddings(openai_api_key=env_settings.OPENAI_API_KEY)
)

# 準備文本切割器
bic_chunk_spliter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=50)
small_chunk_spliter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=50)

# 讀取文件並建立檢索器
loader = PyMuPDFLoader(file_path=BASE_DIR / "勞動基準法.pdf")
retriever = ParentDocumentRetriever(
    vectorstore=vector_store,
    docstore=InMemoryStore(),
    child_splitter=small_chunk_spliter,
    parent_splitter=bic_chunk_spliter,
)

# 寫入文件
retriever.add_documents(loader.load())

在上面的程式碼當中我們可以看到在準備切割器的時候準備了兩組,分別是 1000 字以及 200 字的,
並且我們呼叫了了 ParentDocumentRetriever 這個類別,並且將 vectorstoredocstore 進行初始化,這邊要特別注意的是,我們先使用了 InMemoryStore 這個類別作為文本的暫存,稍後會為大家說明該如何套用到自家的資料庫上。

使用 ParentDocumentRetriever 進行檢索

在同一段程式碼的結尾我們加上下方這三行程式碼,這樣我們就可以透過問題來檢索段落了。

# 提問並檢索段落
question = "我可以因為不喜歡某個勞工而終止勞動契約嗎?"
retrieved_docs = retriever.invoke(question)
print(retrieved_docs[0].page_content)

下圖為搜尋出來的結果,可以看到我們成功的透過跟昨天一樣的問題,將更長的相關文本檢索出來:

result

持久化 doc_store

剛剛在進行 add_documents() 動作的時候,應該可以發現我們有成功的把向量資料寫入到 Pinecone 中,但是我們的文本資料並沒有被寫入到任何地方,這也是為什麼檢索這段程式是請大家接著寫,而不是另外寫一段,這邊我們可以透過 docstore 來將文本資料持久化到我們的資料庫中,這樣就可以讓我們在下次檢索的時候不要重新讀取文件,而是直接從資料庫中讀取。

from langchain.storage.base import BaseStore
from pymongo import MongoClient
from typing import List, Optional, Sequence, Tuple, Union, Iterator

class MongoDBStore(BaseStore[str, str]):
    def __init__(self, db_name="document_db", collection_name="documents"):
        # 初始化 MongoDB 客戶端和資料庫
        self.client = MongoClient("mongodb://localhost:27017/")
        self.db = self.client[db_name]
        self.collection = self.db[collection_name]

    def mget(self, keys: Sequence[str]) -> List[Optional[str]]:
        """根據鍵值批量檢索文件內容"""
        documents = self.collection.find({"_id": {"$in": keys}})
        result_dict = {doc["_id"]: doc["document"] for doc in documents}
        return [result_dict.get(key) for key in keys]

    def mset(self, key_value_pairs: Sequence[Tuple[str, str]]) -> None:
        """批量添加或更新文件"""
        for key, value in key_value_pairs:
            self.collection.update_one(
                {"_id": key},
                {"$set": {"document": value}},
                upsert=True
            )

    def mdelete(self, keys: Sequence[str]) -> None:
        """批量刪除文件"""
        self.collection.delete_many({"_id": {"$in": keys}})

    def yield_keys(self, prefix: Optional[str] = None) -> Iterator[str]:
        """返回鍵值的迭代器"""
        if prefix:
            cursor = self.collection.find({"_id": {"$regex": f"^{prefix}"}})
        else:
            cursor = self.collection.find()
        for doc in cursor:
            yield doc["_id"]

這邊我們實作了一個簡單的 MongoDBStore 類別,這個類別繼承了 BaseStore 並實作了 mget、mset、mdelete 以及 yield_keys 四個方法,這四個方法分別是用來批量檢索、批量添加或更新、批量刪除以及返回鍵值的迭代器,這樣我們就可以將文本資料持久化到 MongoDB 中了。

如此一來我們就可以將整體的流程拆分成,「計算」、「儲存」、「檢索」三個步驟,讓我們的程式碼更加模組化,也更容易維護。

內容預告

今天我們介紹了 ParentDocumentRetriever 這個檢索器,透過這個檢索器我們可以將文件拆分成較小的片段來進行大段落的檢索,可以補足向量文本過短的問題,明天我們要來介紹 MultiQueryRetriever 這個檢索器。


上一篇
Day 16 - 檔案與 metadata
下一篇
Day 18 - MultiQueryRetriever 多重查詢檢索器
系列文
初探 Langchain 與 LLM:打造簡易問診機器人30
圖片
  直播研討會
圖片
{{ item.channelVendor }} {{ item.webinarstarted }} |
{{ formatDate(item.duration) }}
直播中

尚未有邦友留言

立即登入留言