RAPTOR: 树结构递归摘要处理系统教程

RAPTOR概述

RAPTOR 引入了一种通过构建文档的递归树结构来增强检索的语言模型新方法。该方法解决了阅读中的语义深度和连接问题，通过构建递归树结构，平衡了更广泛的主题理解与细节的把握。这种方法允许基于语义相似性而非文本顺序来分组节点。

树结构的构建

RAPTOR 树的构建从将检索语料库分段为长度为 100 个令牌的短文本开始。如果句子超过 100 个令牌的限制，我们将整个句子移到下一个块，以保持文本的上下文和语义一致性。然后使用 SBERT（一种基于 BERT 的编码器，multi-qa-mpnet-base-cos-v1）对这些文本进行嵌入。

这些文本块及其相应的 SBERT 嵌入形成了我们树结构的叶节点。为了对相似的文本块进行分组，我们采用了一种聚类算法。聚类后，使用语言模型对分组的文本进行摘要。这些摘要文本会被重新嵌入，然后继续进行嵌入、聚类和摘要的循环，直到进一步的聚类变得不可行，从而生成原始文档的多层次树形表示。

安装

在使用 RAPTOR 之前，确保已安装 Python 3.8 及以上版本。克隆 RAPTOR 仓库并安装必要的依赖：

git clone https://github.com/parthsarthi03/raptor.git
cd raptor
pip install -r requirements.txt

基本使用

按照以下步骤开始使用 RAPTOR：

配置 RAPTOR

首先，设置你的 OpenAI API 密钥并初始化 RAPTOR 配置：

import os
os.environ["OPENAI_API_KEY"] = "your-openai-api-key"

from raptor import RetrievalAugmentation

# 使用默认配置进行初始化。有关高级配置，请查看文档。
RA = RetrievalAugmentation()

添加文档到树结构

将你的文本文档添加到 RAPTOR 进行索引：

with open('sample.txt', 'r') as file:
    text = file.read()
RA.add_documents(text)

回答问题

现在可以使用 RAPTOR 基于索引的文档回答问题：

question = "灰姑娘是如何获得幸福结局的？"
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)

拓展 RAPTOR 与其他模型

RAPTOR 设计为灵活的，允许你集成任何用于摘要、问答（QA）和嵌入生成的模型。以下是如何使用你自己的模型扩展 RAPTOR：

自定义摘要模型

如果希望使用不同的语言模型进行摘要，可以通过扩展 BaseSummarizationModel 类来实现。实现 summarize 方法来集成你的自定义摘要逻辑：

from raptor import BaseSummarizationModel

class CustomSummarizationModel(BaseSummarizationModel):
    def __init__(self):
        # 初始化你的模型
        pass

    def summarize(self, context, max_tokens=150):
        # 实现你的摘要逻辑
        summary = "你的摘要"
        return summary

自定义问答模型

对于自定义的问答模型，扩展 BaseQAModel 类并实现 answer_question 方法。该方法应返回模型在给定上下文和问题的情况下找到的最佳答案：

from raptor import BaseQAModel

class CustomQAModel(BaseQAModel):
    def __init__(self):
        # 初始化你的模型
        pass

    def answer_question(self, context, question):
        # 实现你的问答逻辑
        answer = "你的答案"
        return answer

自定义嵌入模型

要使用不同的嵌入模型，扩展 BaseEmbeddingModel 类。实现 create_embedding 方法，该方法应返回输入文本的向量表示：

from raptor import BaseEmbeddingModel

class CustomEmbeddingModel(BaseEmbeddingModel):
    def __init__(self):
        # 初始化你的模型
        pass

    def create_embedding(self, text):
        # 实现你的嵌入逻辑
        embedding = [0.0] * embedding_dim  # 替换为实际的嵌入逻辑
        return embedding

集成自定义模型与 RAPTOR

在实现自定义模型后，将它们与 RAPTOR 集成：

from raptor import RetrievalAugmentation, RetrievalAugmentationConfig

# 初始化你的自定义模型
custom_summarizer = CustomSummarizationModel()
custom_qa = CustomQAModel()
custom_embedding = CustomEmbeddingModel()

# 创建包含自定义模型的配置
custom_config = RetrievalAugmentationConfig(
    summarization_model=custom_summarizer,
    qa_model=custom_qa,
    embedding_model=custom_embedding
)

# 使用你的自定义配置初始化 RAPTOR
RA = RetrievalAugmentation(config=custom_config)

附录

RAPTOR 树的扩展

基于长文本通常呈现子主题和层次结构的想法，RAPTOR 通过构建递归树结构解决了阅读中的语义深度和连接问题。构建 RAPTOR 树始于将检索语料库分割成长度为 100 的连续文本。如果句子超过 100 个令牌限制，我们会将整个句子移到下一个块，以保持文本的上下文和语义一致性。这些文本会使用基于 BERT 的编码器（multi-qa-mpnet-base-cos-v1）进行嵌入。

聚类算法

聚类在构建 RAPTOR 树中起着关键作用，将文本片段组织成紧密的组。我们的聚类算法基于高斯混合模型（GMM），该方法既灵活又提供了概率框架。GMM 假设数据点是由多个高斯分布的混合生成的。

在高维空间中使用距离度量来测量相似性可能会出现问题。为了缓解这一问题，我们采用了均匀流形逼近和投影（UMAP），这是一种用于降维的流形学习技术。我们算法的独特之处在于使用软聚类，即节点可以属于多个聚类，而不需要固定的聚类数量。我们的算法首先识别全局聚类，然后在这些全局聚类内执行局部聚类。

模型摘要

在使用高斯混合模型对节点进行聚类后，这些节点会被送到语言模型进行摘要。我们使用 gpt-3.5-turbo 生成摘要。摘要步骤将大量的检索信息压缩成可管理的大小。

查询策略

我们引入了两种不同的查询策略：树遍历和折叠树。树遍历方法逐层遍历树，在每个层级修剪和选择最相关的节点。折叠树方法则在所有层级中集体评估节点，以找到最相关的节点。

这份教程为你提供了一个从零开始使用 RAPTOR 的全面指南。希望这能帮助你更好地理解和使用这项技术！