速度提高100倍 - 扩展 RAG 应用程序,以实现数十亿个嵌入,并行计算余弦相似度

原文链接:100x Faster — Scaling Your RAG App for Billions of Embeddings

2024 年 2 月 15 日

RAG应用程序最大的问题之一是它们的计算检索时间。想象一下,你有一个向量数据库,包含一万亿条Embedding向量的记录。当您尝试将用户查询与一万亿向量匹配时,检索正确的信息肯定要花费一分钟以上的时间。

我们能否在CPU内核上使用并行处理将检索正确信息的时间缩短到几秒钟?

减少时间包括找到有效的方法来计算用户查询Embedding向量与存储在向量数据库中的百万,十亿,甚至万亿个其他Embedding向量之间的余弦相似度。

Chunkdot,在MIT许可下,是专门为此目的而设计的,为密集矩阵和稀疏矩阵提供多线程矩阵乘法。它适用于对项目矩阵表示进行分段(Embeddings),并使用Numba加速计算,从而计算出大量项目中K个最相似的项目。

HuggingFace上有很多数据集,提供了超过100万个条目的Embedding向量,比如这个来自Qdrant的dataset。你可以用它来测试Chunkdot的性能。然而,对于详细的性能测量,我们将使用NumPy库来生成各种维度的随机Embedding向量。

我们将比较两种方法,一种来自Chunkdot,另一种是余弦相似度的伪代码。我们将观察增加大小和维度对性能的影响。我将使用Kaggle(无GPU)笔记本来完成这项任务,以确保一致性。

这个博客的所有代码都可以在我的GitHub存储库中找到。

目录表

  • 舞台设置

  • 编码伪码算法

  • 编码块点算法

  • 编码计算时间函数

  • 测试10k向量Embeddings

  • 测试100k向量Embeddings

  • 测试100万个向量Embeddings

  • 可视化可扩展性影响

  • Chunkdot功能

  • 下一步是什么

搭建舞台

Chunkdot需要与其他库类似的安装过程。

1
2
# installing chunkdot
pip install chunkdot

在运行任何东西之前,我们必须首先检查Kaggle环境中的可用内存。

1
2
# Checking available memory
!free -h

img

可用内存在Kaggle笔记本

检查可用内存对Chunkdot至关重要。随着向量数据库大小的增加,计算内存也会增加。为了防止超出可用内存,监控硬件中的剩余内存非常重要。在我的情况下,可用空间是25GB,不包括Buff/Cache。

让我们导入必要的库。

1
2
3
4
5
6
7
8
# to matrix generate matrices
import numpy as np# importing cosine similarity module from chunkdot
from chunkdot import cosine_similarity_top_k# to calculate computation time
import timeit

伪代码算法

我们将首先构建一个伪代码算法,计算用户查询向量与其他数百万个可能存储在数据库或本地的向量之间的余弦相似度。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def cosine_pseudocode(query_v, doc_v, num_indices):"""Retrieve indices of the highest cosine similarity values betweenthe query vector and embeddings.Parameters:query_v (numpy.ndarray): Query vector.doc_v (list of numpy.ndarray): List of embedding vectors.num_indices (int): Number of Top indices to retrieve.Returns:list of int: Indices of the highest cosine similarity values."""cosine_similarities = []  # Initialize an empty list to store cosine similaritiesquery_norm = np.linalg.norm(query_v)  # Calculate the norm of the query vector# Iterate over each documents embedding vectors in the listfor vec in doc_v:dot_product = np.dot(vec, query_v.T)  # Calculate dot product between embedding vector and query vectorembedding_norm = np.linalg.norm(vec)  # Calculate the norm of the embedding vectorcosine_similarity = dot_product / (embedding_norm * query_norm)  # Calculate cosine similaritycosine_similarities.append(cosine_similarity)  # Append cosine similarity to the listcosine_similarities = np.array(cosine_similarities)  # Convert the list to a numpy array# Sort the array in descending ordersorted_array = sorted(range(len(cosine_similarities)), key=lambda i: cosine_similarities[i], reverse=True)# Get indices of the top two valuestop_indices = sorted_array[:num_indices]# Return the indices of highest cosine similarity valuesreturn top_indices

这个余弦相似度函数,独立于除NumPy以外的任何库,接受三个输入:

  • query_v用户查询的Embedding向量
  • doc_v存储在某处的文档的Embedding向量
  • num_indices类似top_k结果的文档索引号

Chunkdot算法

现在我们已经编写了伪代码算法,下一步是编写Chunkdot余弦相似度函数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def cosine_chunkdot(query_v, doc_v, num_indices, max_memory):"""Calculate cosine similarity using the chunkdot library.Parameters:query_v (numpy.ndarray): Query vector.doc_v (numpy.ndarray): List of Embedding vectors.num_indices (int): Number of top indices to retrieve.max_memory (float): Maximum memory to use.Returns:numpy.ndarray: Top k indices."""# Calculate Cosine Similaritycosine_array = cosine_similarity_top_k(embeddings=query_v, embeddings_right=doc_v, top_k=num_indices, max_memory=max_memory)  # Calculate cosine similarity using chunkdot# Get indices of the top valuestop_indices = cosine_array.nonzero()[1]# return the top similar resultsreturn top_indices

这个Chunkdot函数有四个输入:

  • query_v用户查询的Embedding向量
  • doc_v存储在某处的文档的Embedding向量
  • num_indices类似top_k结果的文档索引号
  • max_memory表示计算的可用内存,其值以字节为单位。例如,1E9表示1GB, 10E9表示10GB,以此类推。

让我们在一个样本数据集上测试这两个函数,观察它们的输出。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
doc_embeddings = np.random.randn(10, 100) # 10 document embeddings (100 dim)user_query = np.random.rand(1,100) # 1 user query (100 dim)top_indices = 1 # number of top indices to retrievemax_memory = 5E9 # maximum memory to use (5GB)# retrieve indices of the highest cosine similarity values using pseudocode
print("top indices using pseudocode:", cosine_pseudocode(user_query, doc_embeddings, top_indices))# retrieve indices of the highest cosine similarity values using chunkdot
print("top indices using chunkdot:", cosine_chunkdot(user_query, doc_embeddings, top_indices, max_memory))
### OUTPUT ###
top indices using pseudocode: [4]
top indices using chunkdot: [4]
### OUTPUT ###

我为文档Embeddings生成了10个随机Embedding向量,每个向量的维度为100,还有一个用户查询,它是具有相同维度的单个Embedding向量。’ top_indices ‘参数设置为1,这意味着它将根据最高余弦相似度返回文档Embeddings中仅一个相似项的索引。内存使用率设置为5E9,等于5GB。我们的两个函数都返回相同的索引4,这表明我们对两个函数都进行了准确的编码。

编码计算时间函数

我们还需要创建一个计时函数,它可以测量这两个函数输出结果所花费的计算时间。

1
2
3
4
5
6
7
8
9
10
11
12
# calculate time taken
def calculate_execution_time(query_v, doc_v, num_indices, max_memory, times):# calculate time taken to execute the pseudocode functionpseudocode_time = round(timeit.timeit(lambda: cosine_pseudocode(query_v, doc_v, num_indices), number=times), 5)# calculate time taken to execute the chunkdot functionchunkdot_time = round(timeit.timeit(lambda: cosine_chunkdot(query_v, doc_v, num_indices, max_memory), number=times), 5)# print the time takenprint("Time taken for pseudocode function:", pseudocode_time, "seconds")print("Time taken for chunkdot function:", chunkdot_time, "seconds")

我们已经回顾了传递给这个函数的参数。这里唯一的新参数是times,它告诉函数你想运行代码多少次。让我们在更大的规模上测试Chunkdot性能的效率。

测试10k向量Embeddings

我们将从合理数量的文档Embeddings开始,10000个,这相当于一个小规模的特定于领域的RAG应用程序。我将每个Embedding向量的维度设置为1536,这相当于OpenAIEmbedding模型text-embedding-3-small

让我们通过运行100次来计算每种方法的计算时间。

1
2
3
4
5
6
7
8
9
10
doc_embeddings = np.random.randn(10000, 1536) # 10K document embeddings (1536 dim)user_query = np.random.rand(1,1536) # user query (1536 dim)top_indices = 1 # number of top indices to retrieve max_memory = 5E9 # maximum memory set to 5GB# compute the time taken to execute the functions
calculate_execution_time(user_query, doc_embeddings, top_indices, max_memory, 100)

对于10k个文档Embeddings,维度为1536,两种算法运行100次,对比如下:

img

10k文档计算时间

与我们的伪代码相比,Chunkdot需要更多的时间。这是因为它首先创建块,并在合并它们之前对每个块执行计算。因此,对于这个小规模的例子,它可能不是一个合适的解决方案。但是,当我们稍后使用更大的示例时,您将看到Chunkdot的好处。

测试100k向量Embeddings

对于10K,我们的伪代码方法获胜,但是现在让我们将文档Embedding向量增加到100K向量,这与中等规模的RAG应用程序相当。

让我们计算每种方法的计算时间,但这次我们将“times”参数设置为1(只运行一次代码),因为向量的数量相当大,并且不需要多次执行计算。

1
2
3
4
5
6
7
8
9
10
11
12
doc_embeddings = np.random.randn(100000, 1536) # 100K document embeddings (1536 dim)user_query = np.random.rand(1,1536) # user query (1536 dim)top_indices = 1 # number of top indices to retrieve max_memory = 5E9 # maximum memory set to 5GBtimes = 1 # number of times to execute the functions# compute the time taken to execute the functions
calculate_execution_time(user_query, doc_embeddings, top_indices, max_memory, times)

对于100k的文档Embeddings,维度为1536,运行两种算法一次,下面是比较:

img

100k文档计算时间

与我们的伪代码相比,Chunkdot花费的时间更少,几乎是一半。现在我们看到了Chunkdot带来的积极影响。

一百万向量Embeddings的测试

处理涉及数百万个Embeddings的任务时,您需要检查的第一件事是文档Embedding向量占用了多少内存。

1
2
3
4
5
6
7
8
9
# 1 Million document embeddings (1536 dim)
doc_embeddings = np.random.randn(1000000, 1536)# user query (1536 dim)
user_query = np.random.rand(1,1536)# Check the memory size of doc_embeddings and user_query embedding
print(doc_embeddings.nbytes / (1024 * 1024 * 1024),user_query.nbytes / (1024 * 1024))

img

100万个Embedding向量的内存大小

我们的文档Embeddings大约占用12GB。让我们检查一下剩余的可用空间。

img

检查可用空间

我们有高达17GB的可用内存。为了避免任何内存错误,我们将为max_memory参数设置一个安全值,即12GB。让我们看看结果。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 1 Million document embeddings (1536 dim)
doc_embeddings = np.random.randn(1000000, 1536)# user query (1536 dim)
user_query = np.random.rand(1,1536)top_indices = 1 # number of top indices to retrieve max_memory = 12E9 # maximum memory set to  --- 12GB ---times = 1 # number of times to execute the functions# compute the time taken to execute the functions
calculate_execution_time(user_query, doc_embeddings, top_indices, max_memory, times)

img

100万个文档计算时间

ChunkDot确实有效地减少了计算量。当你打算构建一个严肃的RAG应用程序时,你应该考虑从至少一百万个查询开始。工作与更高维度的Embedding模型,高达4000。这种方法将变得更加有效。

可视化可伸缩性影响

让我们可视化增加文档Embedding向量数量的影响,从10,000开始到一个非常大的数字。

img

不同数量文档的计算时间

我绘制了三种方法,在增加文档Embeddings数量的基础上,Chunkdot是所有方法中最优越的。现在,让我们看看Embedding向量的维数是如何影响计算时间的。

img

不同维度的计算时间

我在增加向量维度的同时使用了100K个文档,在增加文档数量时观察到的行为与我们看到的相同。

Chunkdot的特点

Chunkdot有一个可以显示进度条的功能,它可以帮助您跟踪剩余的计算量。

1
2
3
4
5
6
7
8
9
10
11
12
doc_embeddings = np.random.randn(100000, 1536) # 100K document embeddings (1536 dim)user_query = np.random.rand(1,1536) # user query (1536 dim)top_indices = 100 # number of top indices to retrieve max_memory = 5E9 # maximum memory set to 5GB# with progress bar
output_array = cosine_similarity_top_k(user_query, doc_embeddings, top_k=top_indices, show_progress=True)

img

进度条示例

Chunkdot的输出是一个稀疏矩阵,您可以使用以下命令将其转换为数组:

1
2
# converting the ouput
output_array.toarray()

您可以仅对文档Embeddings使用Chunkdot,它将为文档Embeddings的每个元素返回top_k个最相似的元素。

1
2
3
4
5
6
7
8
9
10
11
12
# total 5 documents embeddings
embeddings = np.random.randn(5, 256)# return top 2 most similar item index for each
cosine_similarity_top_k(embeddings, top_k=2).toarray()
### OUTPUT ###
array([[1.        , 0.        , 0.        , 0.        , 0.09924064],[0.        , 1.        , 0.        , 0.09935381, 0.        ],[0.02358785, 0.        , 1.        , 0.        , 0.        ],[0.        , 0.09935381, 0.        , 1.        , 0.        ],[0.09924064, 0.        , 0.        , 0.        , 1.        ]])
### OUTPUT ###

类似地,您可以通过向top_k参数提供负值来返回最不相似的项

1
2
3
4
5
6
7
8
9
10
11
12
13
# total 5 documents embeddings
embeddings = np.random.randn(5, 256)# return top 2 most dissimilar item index for each 
# Top_K = -2
cosine_similarity_top_k(embeddings, top_k=-2).toarray()
### OUTPUT ###
array([[ 0.        ,  0.        , -0.04357524,  0.        , -0.05118288],[ 0.        ,  0.        ,  0.        ,  0.01619543, -0.01836534],[-0.04357524,  0.        ,  0.        , -0.02466613,  0.        ],[ 0.        ,  0.01619543, -0.02466613,  0.        ,  0.        ],[-0.05118288, -0.01836534,  0.        ,  0.        ,  0.        ]])
### OUTPUT ###

这可能不是你的情况,但如果你处理的是10K维的稀疏Embeddings,你可以使用“密度”参数来更有效地减少计算。

1
2
3
4
5
6
7
8
9
# for creating sparse embeddings
from scipy import sparse# creating spare matrix with 100K documents (10K dim each)
# defining density of 0.005
embeddings = sparse.rand(100000, 10000, density=0.005)# using all you system's memory
cosine_similarity_top_k(embeddings, top_k=50)

接下来

如果您想了解Chunkdot算法是如何工作的,请查看作者的这个有意思的博客。Chunkdot最大的好处之一是它可以在CPU内核上工作。未来,他们计划集成GPU支持,这将大大减少计算时间。如果你的本地环境没有足够的RAM,你可以使用像Kaggle或GitHub Codespaces这样的平台,与GPU成本相比,云CPU内核和RAM的成本非常低。不要忘记查看官方GitHub存储库和他们的博客,因为它非常好地解释了Chunkdot是如何工作的。

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.mzph.cn/news/700880.shtml

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

idea 打jar包、lib文件夹

idea目录文件 idea四层级结构 idea操作Java文件的基本单位:项目(Project)。对应四级结构 第1层级架构:项目(project) 在 IntelliJ IDEA 中Project是最顶级的结构单元,然后就是Module&#xf…

Llama中文大模型-模型+工具(外延能力)

除了持续增强大模型内在的知识储备、通用理解、逻辑推理和想象能力等,未来,我们也会不断丰富大模型的外延能力,例如知识库检索、计算工具、WolframAlpha、操作软件等。 我们首先集成了LangChain框架,可以更方便地基于Llama2开发文…

MySQL知识点总结(六)——InnoDB底层架构

MySQL知识点总结(六)——InnoDB底层架构 InnoDB底层架构总览InnoDB底层各组件分析Buffer PoolChange BufferLog BufferAdaptive Hash IndexSystem TablesapceUndo TablespacesRedo Log InnoDB底层架构总览 关于InnoDB底层架构,网上有一张非常…

JS基础之常用方法

JS基础之常用方法 目录 JS基础之常用方法字符串数字布尔数组 字符串 length: 返回字符串的长度 var sen " Hello World" console.log(sen) // Hello World console.log(sen.length) // 12trim():去除首尾空格 var sen " Hello World" console.log(se…

电脑常见问题及解决方法

如何速度定位电脑故障? 根据开机报警的声音,判断硬件的问题。AMI BIOS 1 短——内存刷新失败 2 短——内存ECC 校验错误 3 短——系统基本内存[检查失败 4 短——系统时钟出错 5 短——中央处理器错误 6 短——键盘控制器错误 7 短——系统实模式…

Linux之JAVA环境配置Tomcat离线安装与启动

一,安装jdk和Tomcat 1.1上传JDK跟Tomcat 1.2解压 解压tomcat tar -zxvf apache-tomcat-8.5.20.tar.gz 解压jdk tar -zxvf jdk-8u151-linux-x64.tar.gz 1.3.配置环境变量 vim /etc/profile 在最后加上: #java environment export JAVA_HOME/usr/local/ja…

异步框架Celery在Django中的运用

参考博客:https://www.cnblogs.com/pyedu/p/12461819.html 参考视频:01 celery的工作机制_哔哩哔哩_bilibili 定义:简单灵活、处理大量消息的分布式系统,专注于实时处理异步队列,支持任务调度 主要架构: …

Linux的信号

Linux的信号是一种用于进程之间通信的机制。它们用于向进程发送通知,告知进程发生了某种事件或请求进程执行某个操作。信号可以由内核、其他进程或进程自身发送。 信号的作用有以下几个方面: 通知进程某个事件的发生,如进程的终止、挂起、恢…

springboot/ssm来访管理系统Java访客预约来访登记系统web

springboot/ssm来访管理系统Java访客预约来访登记系统web 基于springboot(可改ssm)vue项目 开发语言:Java 框架:springboot/可改ssm vue JDK版本:JDK1.8(或11) 服务器:tomcat 数据库:mysq…

软件实例,物流货运配货单打印模板软件单据打印查询管理系统软件教程,可以同时打印标签或补打

软件实例,物流货运配货单打印模板软件单据打印查询管理系统软件教程,可以同时打印标签或补打 一、前言 以下软件教程以 佳易王物流单打印查询系统V17.1为例说明 软件文件下载可以点击最下方官网卡片——软件下载——试用版软件下载 这个版本在原来基…

Android 跨进程通信技术优劣分析

一、引言 在Android开发中,跨进程通信(IPC)是一项常见的任务,用于在不同的应用程序或组件之间交换数据。有多种IPC机制可供选择,每种机制都有其优势和劣势。以下是对几种常见IPC技术的分析。 二、跨进程通信 2.1、Bi…

【软件使用】postman使用教程

​ 🍎个人博客:个人主页 🏆个人专栏:软件安装及使用 ⛳️ 功不唐捐,玉汝于成 ​ 目录 前言 正文 步骤1:安装Postman 步骤2:发送请求 步骤3:管理环境变量 步骤4&#xff1…

Leetcode 26-30题

删除有序数组中的重复项 给定一个有序数组,要求原地删除重复出现的元素,返回删除后的数组的长度。 这里的原地删除其实可以这样表示,用双指针从前往后扫一遍,遇到新的没出现过的元素就放到前面去,就可以实现删除后的数…

Linux线程同步(2)死锁与互斥锁

死锁(Deadlock)是指两个或两个以上的进程(或线程)在执行过程中,由于竞争资源或者由于彼此通信而造成的一种阻塞的现象,若无外力作用,它们都将无法推进下去。此时称系统处于死锁状态或系统产生了…

定频空调与变频空调的区别

变频空调是通过变频器来改变压缩机运转电压以及频率&#xff0c;从而改变压缩机转速的新型空调&#xff0c;那么他和定频空调有哪些区别呢&#xff1f; 一、控制启停的方式不同 定频空调控制方式为&#xff1a;当T环 < T设 - 2℃&#xff0c;压缩机停机。 当 T环 > T…

Java实现就医保险管理系统 JAVA+Vue+SpringBoot+MySQL

目录 一、摘要1.1 项目介绍1.2 项目录屏 二、功能模块2.1 科室档案模块2.2 医生档案模块2.3 预约挂号模块2.4 我的挂号模块 三、系统展示四、核心代码4.1 用户查询全部医生4.2 新增医生4.3 查询科室4.4 新增号源4.5 预约号源 五、免责说明 一、摘要 1.1 项目介绍 基于JAVAVue…

【前端素材】推荐优质后台管理系统Qovex平台模板(附源码)

一、需求分析 1、定义 后台管理系统是一种用于管理和监控网站、应用程序或系统的在线工具。它通常是通过网页界面进行访问和操作&#xff0c;用于管理网站内容、用户权限、数据分析等。后台管理系统是网站或应用程序的控制中心&#xff0c;管理员可以通过后台系统进行各种管理…

构造函数,原型,实例,类的关系整理

视频来源js原型链、构造函数和类_哔哩哔哩_bilibili 如视频所说&#xff0c;构造函数的prototype指向原型&#xff0c;实例化的对象的__proto__指向原型&#xff0c;原型通过constructor指向构造函数&#xff0c;正如class里面的constructor方法就相当于Person构造函数一样&am…

基于PID-bang-bang控制算法的卫星姿态控制matlab仿真

目录 1.课题概述 2.系统仿真结果 3.核心程序与模型 4.系统原理简介 5.完整工程文件 1.课题概述 基于PID-bang-bang控制算法的卫星姿态控制。仿真输出控制器的控制收敛曲线&#xff0c;卫星姿态调整过程的动画。 2.系统仿真结果 3.核心程序与模型 版本&#xff1a;MATLAB…

连接SFTP服务器报错host key for has changed and you have requested strict checking

背景&#xff1a;今天连接SFTP的时候&#xff0c;报错如下&#xff0c;记录一下 报错如下 WARNING: REMOTE HOST IDENTIFICATION HSS CHANGED!cccececccccceeteeacaaeapepeppeeecpeepeeecepeeeaeeeeeceeeeeeIT IS POSSIBLE THAT SOMEONE IS DOING SOMETHING NASTY!Someone co…