Project Icon

RAGatouille

优化RAG管道的先进检索工具

RAGatouille是一个开源的检索增强生成(RAG)工具包,专注于将先进的检索方法应用于RAG管道。它集成了ColBERT等最新研究成果,提供简单易用的API接口用于模型训练、文档索引和检索。RAGatouille的设计理念是模块化和易用性,同时保持高度可定制性。通过优化检索性能,该工具包旨在提升RAG系统的整体效果,促进信息检索技术在实际应用中的发展。

欢迎使用 RAGatouille

在任何 RAG 流程中轻松使用和训练最先进的检索方法。注重模块化和易用性,由研究支持。

GitHub 星标 Python 版本 下载量 文档 Twitter 关注

RAGatouille 标志,是一只快乐的老鼠在笔记本电脑上工作(电脑上贴着一块被咬了一口的奶酪),旁边堆着一堆他正在查阅信息的书籍。


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 的全部功能,如果需要,你还应该能够调整任何相关参数!
  • 强大而简单的可重用底层组件:库的任何部分都应该可以独立使用。你可以在 RAGPretrainedModelRagTrainer 之外使用我们的 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 倾向于将处理后的训练数据存储在文件中,这也使得通过 wandbdvc 对训练数据进行适当的版本控制变得更容易。默认情况下,它会写入到 ./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支持,正在积极开发中
项目侧边栏1项目侧边栏2
推荐项目
Project Cover

豆包MarsCode

豆包 MarsCode 是一款革命性的编程助手,通过AI技术提供代码补全、单测生成、代码解释和智能问答等功能,支持100+编程语言,与主流编辑器无缝集成,显著提升开发效率和代码质量。

Project Cover

AI写歌

Suno AI是一个革命性的AI音乐创作平台,能在短短30秒内帮助用户创作出一首完整的歌曲。无论是寻找创作灵感还是需要快速制作音乐,Suno AI都是音乐爱好者和专业人士的理想选择。

Project Cover

有言AI

有言平台提供一站式AIGC视频创作解决方案,通过智能技术简化视频制作流程。无论是企业宣传还是个人分享,有言都能帮助用户快速、轻松地制作出专业级别的视频内容。

Project Cover

Kimi

Kimi AI助手提供多语言对话支持,能够阅读和理解用户上传的文件内容,解析网页信息,并结合搜索结果为用户提供详尽的答案。无论是日常咨询还是专业问题,Kimi都能以友好、专业的方式提供帮助。

Project Cover

阿里绘蛙

绘蛙是阿里巴巴集团推出的革命性AI电商营销平台。利用尖端人工智能技术,为商家提供一键生成商品图和营销文案的服务,显著提升内容创作效率和营销效果。适用于淘宝、天猫等电商平台,让商品第一时间被种草。

Project Cover

吐司

探索Tensor.Art平台的独特AI模型,免费访问各种图像生成与AI训练工具,从Stable Diffusion等基础模型开始,轻松实现创新图像生成。体验前沿的AI技术,推动个人和企业的创新发展。

Project Cover

SubCat字幕猫

SubCat字幕猫APP是一款创新的视频播放器,它将改变您观看视频的方式!SubCat结合了先进的人工智能技术,为您提供即时视频字幕翻译,无论是本地视频还是网络流媒体,让您轻松享受各种语言的内容。

Project Cover

美间AI

美间AI创意设计平台,利用前沿AI技术,为设计师和营销人员提供一站式设计解决方案。从智能海报到3D效果图,再到文案生成,美间让创意设计更简单、更高效。

Project Cover

稿定AI

稿定设计 是一个多功能的在线设计和创意平台,提供广泛的设计工具和资源,以满足不同用户的需求。从专业的图形设计师到普通用户,无论是进行图片处理、智能抠图、H5页面制作还是视频剪辑,稿定设计都能提供简单、高效的解决方案。该平台以其用户友好的界面和强大的功能集合,帮助用户轻松实现创意设计。

投诉举报邮箱: service@vectorlightyear.com
@2024 懂AI·鲁ICP备2024100362号-6·鲁公网安备37021002001498号