RAPTOR 檢索樹再次進一步提升RAG性能的設計思路原創

發布于 2024-7-11 13:29

瀏覽

0收藏

大多數現有方法僅從檢索語料庫中檢索短的連續塊，限制了對整個文檔上下文的整體理解。RAPTOR（Recursive Abstractive Processing for Tree-Organized Retrieval）引入了一種新方法，即遞歸嵌入、聚類和總結文本塊，從下往上構建具有不同總結級別的樹。在推理時，RAPTOR 模型從這棵樹中檢索，整合不同抽象級別的長文檔中的信息。RAPTOR（Recursive Abstractive Processing for Tree-Organized Retrieval）比傳統的檢索增強型 LM 性能與絕對準確度上提高 20%。

開源地址：??https://github.com/parthsarthi03/raptor??

論文地址：??https://arxiv.org/abs/2401.18059??

一、RAPTOR 檢索樹的構建過程

RAPTOR模型的精髓在于其樹狀結構的構建，這一過程通過精心設計的遞歸算法實現，形成了一個層次分明的多層級樹形架構。這種結構不僅顯著提升了模型對長文本信息的處理能力，更使得對復雜文本內容的理解和檢索變得高效而精準。通過遞歸的方式，RAPTOR能夠逐步細化文本信息的粒度，從而確保在多個層次上都能捕捉到關鍵信息，為后續的文本分析和應用奠定了堅實的基礎。

下面我將對具體步驟進行詳細的說明：

RAPTOR 檢索樹再次進一步提升RAG性能的設計思路-AI.x社區

首先，需要對文本進行合理的切片處理。這里需要注意選擇合適的切片算法很重要，建議參考主流框架的實現方案。
然后，RAPTOR 根據其語義embedding遞歸地對文本塊chunk進行聚類，并生成這些聚類的文本摘要。
RAPTOR采用軟聚類方法，允許文本塊跨多個聚類，基于高斯混合模型（GMMs）和UMAP技術進行降維，以捕捉文本數據的復雜結構和關系，從而優化文本聚類效果。
RAPTOR通過遞歸的向量分析，精準地對文本塊進行聚類，并提煉出這些聚類的核心摘要，自下而上地構建出一個結構化的樹形模型。在此樹中，相近的節點自然聚集形成兄弟關系，而它們的父節點則承載著整個集群的概要性文本信息。這種設計確保了文本信息的層次化和結構化表達，便于理解和檢索。

二、RAPTOR 的檢索過程

RAPTOR 檢索樹再次進一步提升RAG性能的設計思路-AI.x社區

這里有兩種方式實現，基于深度檢索樹(樹遍歷)和 廣度檢索樹(折疊樹)

深度檢索樹從樹的根級別開始，根據與查詢向量的余弦相似度檢索頂層的 top-k (這里為 top-1) 節點。在每一層，它根據與查詢向量的余弦相似度從上一層的 top-k 節點的子節點中檢索 top-k 節點。這個過程一直重復，直到達到葉節點。最后，將所有選定節點的文本連接起來形成檢索到的上下文。
廣度檢索樹方法將整個樹壓縮成單一層，然后根據與查詢向量的余弦相似度評估所有層的節點，直到達到設定閾值。

三、開源項目案例應用示例

# 安裝，在使用 RAPTOR 之前，請確保已安裝 Python 3.8+。克隆 RAPTOR 存儲庫并安裝必要的依賴項：
git clone https://github.com/parthsarthi03/raptor.git
cd raptor
pip install -r requirements.txt


# 開始使用 RAPTOR，請按照以下步驟操作：
# 首先，設置您的 OpenAI API 密鑰并初始化 RAPTOR 配置：
import os
os.environ["OPENAI_API_KEY"] = "your-openai-api-key"


from raptor import RetrievalAugmentation


# Initialize with default configuration. For advanced configurations, check the documentation. [WIP]
RA = RetrievalAugmentation()


# 將您的文本文檔添加到 RAPTOR 進行索引：
with open('sample.txt', 'r') as file:
    text = file.read()
RA.add_documents(text)


# 現在可以使用 RAPTOR 根據索引文檔回答問題：
question = "How did Cinderella reach her happy ending?"
answer = RA.answer_question(question=question)
print("Answer: ", answer)


# 將構造好的樹保存到指定路徑：
SAVE_PATH = "demo/cinderella"
RA.save(SAVE_PATH)


# 將保存的樹重新加載到 RAPTOR 中：
RA = RetrievalAugmentation(tree=SAVE_PATH)
answer = RA.answer_question(question=question)