欢迎使用 RAGatouille
在任何 RAG 流程中轻松使用和训练最先进的检索方法。注重模块化和易用性,由研究支持。
RAGatouille 的主要动机很简单:弥合最先进研究与炼金术式 RAG 流程实践之间的差距。RAG 很复杂,涉及许多活动部件。要获得最佳性能,你需要优化多个组件:其中一个非常重要的组件是用于检索的模型。
密集检索,即使用如 OpenAI 的 text-ada-002
这样的嵌入,是一个不错的基线,但有大量研究表明密集嵌入可能不是最适合你的用例的选择。
信息检索研究领域最近蓬勃发展,像 ColBERT 这样的模型已被证明能更好地泛化到新的或复杂的领域,比密集嵌入更优,数据效率惊人,甚至更适合高效地在数据量较少的非英语语言上进行训练!不幸的是,这些新方法大多不太为人所知,而且比密集嵌入更难使用。
这就是 RAGatouille 的用武之地: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 的一个优势是预训练模型在泛化方面特别出色,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("宫崎骏"), get_wikipedia_page("吉卜力工作室")]
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 参考以了解更多详情!)
🗄️ 索引
要创建索引,你需要加载一个训练好的模型,这可以是你自己的模型,也可以是来自 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("宫崎骏"), get_wikipedia_page("吉卜力工作室")]
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": "巴拉巴拉", "score": 42.424242, "rank": 1, "document_id": "x"},
...,
{"content": "啊啦啊啦", "score": 24.242424, "rank": k, "document_id": "y"},
]
# 多查询结果
[
[
{"content": "巴拉巴拉", "score": 42.424242, "rank": 1, "document_id": "x"},
...,
{"content": "啊啦啊啦", "score": 24.242424, "rank": k, "document_id": "y"},
],
[
{"content": "巴拉巴拉", "score": 42.424242, "rank": 1, "document_id": "x"},
...,
{"content": "啊啦啊啦", "score": 24.242424, "rank": k, "document_id": "y"},
],
],
如果你的索引包含文档元数据,它将作为结果字典中的document_metadata
键返回:
[
{"content": "巴拉巴拉", "score": 42.424242, "rank": 1, "document_id": "x", "document_metadata": {"A": 1, "B": 2}},
...,
{"content": "啊啦啊啦", "score": 24.242424, "rank": k, "document_id": "y", "document_metadata": {"A": 3, "B": 4}},
]
我被说服了,我能将后期交互RAG集成到我的项目中吗?
要开始使用,RAGatouille打包了构建ColBERT原生索引和查询所需的一切。只需查看文档即可!RAGatouille以压缩格式将索引持久化到磁盘上,一个非常可行的生产部署方式是简单地将你需要的索引集成到你的项目中并直接查询它。不要只相信我们的话,这正是Spotify在生产环境中使用他们自己的向量搜索框架所做的,服务于数千万用户:
无状态:Spotify的许多系统在内存中使用最近邻搜索,实现无状态部署(通过Kubernetes),几乎完全消除了维护有状态数据库集群的维护和成本负担。(Spotify,宣布Voyager)
集成
如果你想使用RAGatouille以外的功能,ColBERT有越来越多的集成,它们都完全支持使用RAGatouille训练或微调的模型!
- 官方ColBERT实现内置了一个查询服务器(使用Flask),你可以通过API请求轻松查询,并且支持用RAGatouille生成的索引!这对于大多数小型应用来说应该足够了,只要你能将索引持久化到磁盘上。
- Vespa提供了一个完全托管的支持ColBERT的RAG引擎:它本质上就像一个向量数据库,但有更多的检索选项!对ColBERT模型的完全支持将在接下来几周内发布,使用RAGatouille训练的模型将像从huggingface hub加载一样简单!Vespa是一个经过充分测试、广泛使用的框架,并且在LangChain中得到完全支持,使其成为替换当前RAG管道以使用ColBERT的理想替代品!
- Intel的FastRAG支持ColBERT模型用于RAG,并且与RAGatouille训练的模型完全兼容。
- LlamaIndex正在构建ColBERT集成,并且已经有了早期的ColBERT支持,正在积极开发中。