优化RAG管道的先进检索工具
RAGatouille是一个开源的检索增强生成(RAG)工具包,专注于将先进的检索方法应用于RAG管道。它集成了ColBERT等最新研究成果,提供简单易用的API接口用于模型训练、文档索引和检索。RAGatouille的设计理念是模块化和易用性,同时保持高度可定制性。通过优化检索性能,该工具包旨在提升RAG系统的整体效果,促进信息检索技术在实际应用中的发展。
在任何 RAG 流程中轻松使用和训练最先进的检索方法。注重模块化和易用性,由研究支持。
<p align="center"><img width=500 alt="RAGatouille 标志,是一只快乐的老鼠在笔记本电脑上工作(电脑上贴着一块被咬了一口的奶酪),旁边堆着一堆他正在查阅信息的书籍。" src="https://yellow-cdn.veclightyear.com/835a84d5/604de8af-b046-495a-99c9-e1e8ad868d14.png"/></p>RAGatouille 的主要动机很简单:弥合最先进研究与炼金术式 RAG 流程实践之间的差距。RAG 很复杂,涉及许多活动部件。要获得最佳性能,你需要优化多个组件:其中一个非常重要的组件是用于检索的模型。
密集检索,即使用如 OpenAI 的 text-ada-002
这样的嵌入,是一个不错的基线,但有大量研究表明密集嵌入可能不是最适合你的用例的选择。
信息检索研究领域最近蓬勃发展,像 ColBERT 这样的模型已被证明能更好地泛化到新的或复杂的领域,比密集嵌入更优,数据效率惊人,甚至更适合高效地在数据量较少的非英语语言上进行训练!不幸的是,这些新方法大多不太为人所知,而且比密集嵌入更难使用。
这就是 RAGatouille 的用武之地:RAGatouille 的目的是弥合这一差距:让在 RAG 流程中使用最先进的方法变得容易,无需担心细节或多年的文献!目前,RAGatouille 专注于简化 ColBERT 的使用。如果你想了解接下来的计划,可以查看我们的大致路线图!
如果你想了解更多关于动机、理念以及为什么 ColBERT 使用的后期交互方法如此有效的信息,请查看文档中的介绍。
想尝试一下吗?很简单,只需运行 pip install ragatouille
就可以开始了!
⚠️ 运行注意事项/要求: ⚠️
if __name__ == "__main__"
内运行RAGatouille 让使用 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 在这里处理了所有事情: