欢迎使用 RAGatouille

轻松在任何 RAG 流程中使用和训练最先进的检索方法。模块化设计、易于使用，且以研究为基础。

RAGatouille 的核心动机很简单：弥合最前沿研究与当前 RAG 流程中“炼金术式”实践之间的鸿沟。RAG 很复杂，涉及众多环节。要获得最佳性能，需要优化多个组件，而其中非常关键的一点就是检索模型的选择。

密集检索（如使用 OpenAI 的 text-ada-002 等嵌入）是一个不错的基线，但大量研究表明，密集嵌入可能并非是你特定用例的最佳选择。

信息检索领域近年来蓬勃发展，像 ColBERT 这样的模型已被证明，与密集嵌入相比，它们能更好地泛化到新的或复杂的领域，数据效率极高，甚至更适合在低资源语言上高效训练。可惜的是，这些新方法大多不太为人所知，而且使用起来比密集嵌入要困难得多。

RAGatouille 正是为此而生：它的目标是弥合这一差距，让你无需纠结于细节或多年文献积累，就能在 RAG 流程中轻松使用最先进的方法。目前，RAGatouille 专注于简化 ColBERT 的使用。想了解后续计划，可以查看我们的整体路线图！

想了解更多关于该项目的动机、理念以及 ColBERT 所使用的“晚期交互”方法为何如此有效，请查阅文档中的介绍部分。

想要立即试用？非常简单，只需运行 pip install ragatouille 即可。

⚠️ 运行注意事项和要求：⚠️

• 如果在脚本中运行，必须在 if __name__ == "__main__" 块内执行。
• 不支持 Windows 系统。RAGatouille 似乎在 WSL 之外无法工作，且在 WSL1 上存在问题。部分用户在 WSL2 上成功运行了 RAGatouille。

快速上手

RAGatouille 让使用 ColBERT 变得极其简单！我们希望库能在两个层面发挥作用：

• 强大但可自定义的默认配置：你只需几行代码就能开始使用，还能充分利用 ColBERT 的全部能力；如有需要，你也可以调整任何相关参数。
• 底层强大且简洁的可复用组件：库中的任何部分都可以独立使用。你可以在 RAGPretrainedModel 和 RagTrainer 之外单独使用 DataProcessor 或负样本挖掘器，甚至可以编写自己的负样本挖掘器并集成到流程中。

本节将快速介绍 RAGatouille 的三个核心方面：

• 🚀 训练与微调 ColBERT 模型
• 🗄️ 嵌入与索引文档
• 🔎 检索文档

➡️ 如果你只想看完整可运行的代码示例，请前往示例目录⬅️

🚀 训练与微调

如果你只是做原型开发，不需要训练自己的模型！微调固然有用，但 ColBERT 的一个优势在于，预训练模型特别擅长泛化，ColBERTv2 已被反复证明在新领域的零样本检索中表现出色！

数据处理

RAGatouille 的 RAGTrainer 内置了 TrainingDataProcessor，它可以接收多种形式的检索训练数据，并自动将其转换为训练三元组，同时进行数据增强。处理流程如下：

• 接收成对数据、标记对数据以及各种形式的三元组（字符串或字符串列表）作为输入——透明处理！
• 自动移除所有重复项，并将所有正/负样本映射到对应的查询。
• 默认情况下，会挖掘困难负样本：即生成与正样本难以区分的负样本，从而提高训练效果。

这一切都由 RAGTrainer.prepare_training_data() 处理，只需将数据传递给它即可：

from ragatouille import RAGTrainer

my_data = [
    ("生命的意义是什么？", "生命的意义是42"),
    ("什么是神经搜索？", "神经搜索是指一系列...相关技术的术语"),
    ...
]  # 这里是无标签的对数据
trainer = RAGTrainer()
trainer.prepare_training_data(raw_data=my_data)

ColBERT 倾向于将处理后的训练数据存储到文件中，这也便于通过 wandb 或 dvc 对训练数据进行版本管理。默认情况下，数据会写入 ./data/ 目录，但你可以通过向 prepare_training_data() 传递 data_out_path 参数来覆盖此设置。

与 RAGatouille 的所有功能一样，prepare_training_data 使用了强大的默认配置，也完全可以自定义参数。

执行训练/微调

训练和微调遵循完全相同的流程。实例化 RAGTrainer 时，必须传入 pretrained_model_name。如果该预训练模型是一个 ColBERT 实例，训练器将进入微调模式；如果是其他类型的 Transformer，则会进入训练模式，从模型权重初始化一个新的 ColBERT 开始训练。

from ragatouille import RAGTrainer
from ragatouille.utils import get_wikipedia_page

pairs = [
    ("生命的意义是什么？", "生命的意义是42"),
    ("什么是神经搜索？", "神经搜索是指一系列...相关技术的术语"),
    # 训练需要更多对数据！详情请参阅示例！
    ...
]

my_full_corpus = [get_wikipedia_page("Hayao_Miyazaki"), get_wikipedia_page("Studio_Ghibli")]

trainer = RAGTrainer(model_name="MyFineTunedColBERT",
                     pretrained_model_name="colbert-ir/colbertv2.0")  # 本例进行微调

# 此步骤处理所有数据预处理，详情请参阅示例！
trainer.prepare_training_data(raw_data=pairs,
                              data_out_path="./data/",
                              all_documents=my_full_corpus)

trainer.train(batch_size=32)  # 使用默认超参数进行训练

当运行 train() 时，如果是微调，默认会继承其父 ColBERT 的超参数；如果是训练新的 ColBERT，则会使用默认训练参数。你可以根据需要随意修改（更多细节请参考示例和 API 参考文档！）。

🗄️ 索引文档

要创建索引，你需要加载一个训练好的模型，可以是自己训练的，也可以是从 Hugging Face Hub 下载的预训练模型！使用默认配置创建索引只需几行代码：

from ragatouille import RAGPretrainedModel
from ragatouille.utils import get_wikipedia_page

RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")
my_documents = [get_wikipedia_page("Hayao_Miyazaki"), get_wikipedia_page("Studio_Ghibli")]
index_path = RAG.index(index_name="my_index", collection=my_documents)

你还可以在创建索引时选择添加文档 ID 或文档元数据：

document_ids = ["miyazaki", "ghibli"]
document_metadatas = [
    {"entity": "person", "source": "wikipedia"},
    {"entity": "organisation", "source": "wikipedia"},
]
index_path = RAG.index(
    index_name="my_index_with_ids_and_metadata",
    collection=my_documents,
    document_ids=document_ids,
    document_metadatas=document_metadatas,
)

运行完成后，索引将保存到磁盘，随时可以查询！RAGatouille 和 ColBERT 负责所有流程：
- 分割文档
- 对文档进行分词
- 识别每个词项
- 嵌入文档并生成嵌入袋
- 压缩向量并存储到磁盘

好奇这是如何工作的？请查阅晚期交互与 ColBERT 概念解释。

🔎 检索文档

索引创建完成后，查询与创建一样简单！你可以直接从索引的配置中加载所需的模型：

from ragatouille import RAGPretrainedModel

query = "ColBERT，我亲爱的 ColBERT，谁是最美的文档？"
RAG = RAGPretrainedModel.from_index("path_to_your_index")
results = RAG.search(query)

这是推荐的做法，因为每个索引都保存了创建它所用模型的完整配置，你可以轻松地重新加载。

RAG.search 是一个灵活的方法！你可以设置 k 值来控制返回结果的数量（默认为 10），也可以一次性搜索多个查询：

RAG.search([
    "宫崎骏创作了哪些漫画？",
    "吉卜力工作室的创始人是谁？",
    "《千与千寻》的导演是谁？"
])

RAG.search 返回的结果格式为字典列表，如果使用了多个查询，则返回字典列表的列表：

# 单次查询结果
[
    {"content": "blablabla", "score": 42.424242, "rank": 1, "document_id": "x"},
    ...,
    {"content": "albalbalba", "score": 24.242424, "rank": k, "document_id": "y"},
]
# 多次查询结果
[
    [
        {"content": "blablabla", "score": 42.424242, "rank": 1, "document_id": "x"},
        ...,
        {"content": "albalbalba", "score": 24.242424, "rank": k, "document_id": "y"},
    ],
    [
        {"content": "blablabla", "score": 42.424242, "rank": 1, "document_id": "x"},
        ...,
        {"content": "albalbalba", "score": 24.242424, "rank": k, "document_id": "y"},
    ],
]

如果索引中包含文档元数据，它将以字典形式返回在结果字典的 document_metadata 键中：

[
    {"content": "blablabla", "score": 42.424242, "rank": 1, "document_id": "x", "document_metadata": {"A": 1, "B": 2}},
    ...,
    {"content": "albalbalba", "score": 24.242424, "rank": k, "document_id": "y", "document_metadata": {"A": 3, "B": 4}},
]

我动心了，可以把晚期交互 RAG 集成到我的项目中吗？

入门很简单，RAGatouille 集成了构建和查询 ColBERT 原生索引所需的一切。只需查阅文档即可！RAGatouille 将索引以压缩格式持久化到磁盘上，一个可靠的生产部署方式就是直接将所需的索引集成到项目中并进行查询。这不是空口说白话，这就是 Spotify 在生产环境中使用他们自己的向量搜索框架服务数千万用户的方式：

无状态：Spotify 的许多系统在内存中使用最近邻搜索，这使得无状态部署（通过 Kubernetes）成为可能，几乎完全消除了维护有状态数据库集群的维护和成本负担。（Spotify，宣布 Voyager）

集成方案

如果你希望使用更多功能，ColBERT 拥有越来越多的集成方案，它们都完全支持使用 RAGatouille 训练或微调的模型！

• 官方 ColBERT 实现 内置了一个查询服务器（使用 Flask），你可以通过 API 请求轻松查询，并且它支持 RAGatouille 生成的索引！只要你能将索引持久化到磁盘，这应该足以满足大多数小型应用的需求。
• Vespa 提供了一个完全托管的 RAG 引擎，支持 ColBERT：它本质上就像一个向量数据库，但拥有更多的检索选项！对 ColBERT 模型的全面支持将在未来几周内发布，届时使用 RAGatouille 训练的模型就像从 Hugging Face Hub 加载一样简单！Vespa 是一个经过充分测试、广泛使用的框架，并得到 LangChain 的完全支持，因此它是替代当前 RAG 流程引入 ColBERT 的理想选择！
• Intel 的 FastRAG 支持用于 RAG 的 ColBERT 模型，并且与 RAGatouille 训练的模型完全兼容。
• LlamaIndex 正在构建 ColBERT 集成，并且已经拥有早期的 ColBERT 支持，正在积极开发中。