likes
comments
collection
share

用Rust一周内编写一个向量数据库

作者站长头像
站长
· 阅读数 3

向量数据库目前在科技界风靡一时,这并非只是炒作。由于人工智能的进步使用了向量嵌入,向量搜索变得越来越重要。这些向量嵌入是单词嵌入、句子或文档的向量表示,它们通过简单地查看向量之间的距离度量,为语义上接近的输入提供语义相似性。

典型的例子来自word2vec,其中"king"(国王)的嵌入与单词"queen"(女王)、"man"(男人)和"woman"(女人)的向量结果非常接近,当按照以下公式排列时:

king - man + woman ≈ queen

这个事实一直让我感到惊奇,但只要我们的嵌入空间维度足够高,它甚至对相当大的文档也适用。使用现代深度学习方法,你可以获得复杂文档的优秀嵌入。

对于TerminusDB,我们需要一种方式来利用这种嵌入来完成我们的用户要求的以下任务:

  • 全文搜索
  • 实体解析(找到可能相同的其他文档进行去重)
  • 相似性搜索(寻找相关内容或推荐系统)
  • 聚类

我们决定使用OpenAI的嵌入来制作原型,但为了获得我们需要的其余特性,我们需要一个向量数据库。

我们需要一些不寻常的特性,包括进行增量索引的能力,以及在提交的基础上进行索引的能力,这样我们就可以准确地知道索引适用于哪个提交。这使我们能够将索引放入我们的CI工作流程中。开源的、具有版本控制的向量数据库在市场上是不存在的。所以我们自己写了一个!

编写一个向量数据库

向量数据库是一个向量的存储,能够使用某种度量标准比较任何两个向量。这个度量可以是很多不同的东西,如欧几里得距离、余弦相似性、出租车几何,或者真正遵守三角形不等式规则的任何东西,这些规则需要定义一个度量空间。为了使这个过程快速,你需要有某种索引结构来快速找到已经相近的候选对象进行比较。否则,许多操作将需要每次都与数据库中的每一件事进行比较。索引向量空间有很多方法,但我们选择了HNSW(Hierarchical Navigable Small World)图(参见Malkov和Yashunin)。HNSW易于理解,并在低维和高维中都提供良好的性能,所以是灵活的。最重要的是,我们找到了一个非常清晰的开源实现 - Rust计算机视觉的HNSW

存储向量

向量存储在一个域中。这有助于分离不需要描述相同向量的不同向量存储。对于TerminusDB,我们有许多不同的提交都涉及到相同的向量,所以我们将它们全部放入同一个域中是很重要的。

            Page
            0            1         2...
            ———————————————————————
Vectors:   | 0 [......]  2 [......]
           | 1 [......]  3 [......]

向量存储是基于页面的,其中每个缓冲区都设计为干净地映射到操作系统页面,但又紧密适应我们使用的向量。我们为每个向量分配一个索引,然后我们可以从索引映射到适当的页面和偏移量。

在HNSW索引内部,我们引用一个LoadedVec。这确保了页面存在于一个缓冲区中,目前已加载,所以我们可以对感兴趣的向量进行度量比较。只要最后一个LoadedVec从缓冲区中删除,该缓冲区就可以被加回到缓冲池中,用于加载新的页面。

创建一个版本化的索引

我们为每个(域+提交)对建立一个HNSW结构。如果开始一个新的索引,我们从一个空的HNSW开始。如果从前一个提交开始一个增量索引,我们从前一个提交加载旧的HNSW,然后开始我们的索引操作。什么是新的,什么是旧的,都保留在TerminusDB中,它知道如何在提交之间找到变化,并可以将它们提交给向量数据库索引器。索引器只需要知道它被要求执行的操作(即,插入,删除,替换)。我们在一个LRU池中维护索引本身,这使我们可以按需加载或者如果索引已经在内存中则使用缓存。由于我们只在提交时执行破坏性操作,所以这种缓存始终是一致的。当我们保存索引时,我们使用原始向量索引作为LoadedVec的替代品来序列化结构,这有助于保持索引的小巧。在未来,我们希望使用我们在TerminusDB中学到的一些技巧来保留索引的层,这样新的层可以被添加,而不需要每个增量索引在序列化时添加一个副本。然而,相比我们存储的向量,索引已经足够小,所以这并没有多大关系。我们有一个针对删除和替换操作的设计,我们认为它将与HNSW很好地配合,并希望在此解释,以防任何技术人员有所启发:

  • 如果我们在HNSW的上层,那么简单地忽略删除操作——这应该不会有太大影响,因为大部分向量都不在上层,而那些在上层的向量,只用于导航。 
  • 如果我们在零层但不在上层,从索引中删除该节点,同时尝试根据接近度替换被删除链接的所有邻居之间的链接。
  • 如果我们在零层但也在上层,将该节点标记为已删除,并将其用于导航,但不将该节点存储在候选池中。 

寻找嵌入

我们使用OpenAI来定义我们的嵌入,当向TerminusDB发出索引请求后,我们将每个文档输入到OpenAI,它返回JSON中的浮点向量列表。事实证明,嵌入对上下文非常敏感。我们最初只提交了TerminusDB的JSON文档,结果并不理想。然而,我们发现,如果我们定义一个GraphQL查询+Handlebars模板,我们可以创建非常高质量的嵌入。对于《星球大战》中的人物,这对在我们的模式中定义,看起来像这样:

{
    "embedding": {
        "query": "query($id: ID){ People(id : $id) { birth_year, created, desc, edited, eye_color, gender, hair_colors, height, homeworld { label }, label, mass, skin_colors, species { label }, url } }",
        "template": "The person's name is {{label}}.{{#if desc}} They are described with the following synopsis: {{#each desc}} *{{this}} {{/each}}.{{/if}}{{#if gender}} Their gender is {{gender}}.{{/if}}{{#if hair_colors}} They have the following hair colours: {{hair_colors}}.{{/if}}{{#if mass}} They have a mass of {{mass}}.{{/if}}{{#if skin_colors}} Their skin colours are {{skin_colors}}.{{/if}}{{#if species}} Their species is {{species.label}}.{{/if}}{{#if homeworld}} Their homeworld is {{homeworld.label}}.{{/if}}"
    }
}

在People对象中每个字段的含义都被渲染成文本,这有助于OpenAI理解我们的意思,提供更好的语义。最终,如果我们可以从模式文档和模式结构的组合中猜出这些句子,这可能也可以使用AI聊天实现!但是现在,这种方法运行得非常好,而且不需要太多的技术精细化。

索引星球大战

那么,当我们实际运行这个东西时,会发生什么呢?好吧,我们在我们的星球大战数据产品上试了一下,看看会发生什么。首先,我们发出一个索引请求,我们的索引器从TerminusDB获取信息:

curl 'localhost:8080/index?commit=o2uq7k1mrun1vp4urktmw55962vlpto&domain=admin/star_wars'

这将返回一个任务ID,我们可以使用它轮询一个端点以获取完成情况。

admin/star_wars和commit o2uq7k1mrun1vp4urktmw55962vlpto的索引文件和向量文件输出为:admin%2Fstar_wars@o2uq7k1mrun1vp4urktmw55962vlpto.hnsw和admin%2Fstar_wars.vecs。

现在我们可以询问在指定commit下的语义索引服务器关于我们的文档。

curl 'localhost:8080/search?commit=o2uq7k1mrun1vp4urktmw55962vlpto&domain=admin/star_wars' -d "Who are the squid people"

我们得到了一些结果,以JSON形式返回,看起来像这样:

[{"id":"terminusdb:///star-wars/Species/8","distance":0.09396297}, ...]

但是我们使用哪个嵌入字符串来产生这个结果呢?这就是为Species/8 id渲染的文本:

"The species name is Mon Calamari. They have the following hair colours: 
none. Their skin colours are red, blue, brown, magenta. They speak the 
Mon Calamarian language."

太神奇了!注意到它从来没有在任何地方说过鱿鱼!这里我们的嵌入做了一些非常令人惊奇的工作。让我们再试一次:

curl 'localhost:8080/search?commit=o2uq7k1mrun1vp4urktmw55962vlpto&domain=admin/star_wars' -d "Wise old man"
"The person's name is Yoda. They are described with the following synopsis:
Yoda is a fictional character in the Star Wars franchise created by George 
Lucas, first appearing in the 1980 film The Empire Strikes Back. In the 
original films, he trains Luke Skywalker to fight against the Galactic 
Empire. In the prequel films, he serves as the Grand Master of the Jedi 
Order and as a high-ranking general of Clone Troopers in the Clone Wars. 
Following his death in Return of the Jedi at the age of 900, Yoda was the 
oldest living character in the Star Wars franchise in canon, until the 
introduction of Maz Kanata in Star Wars: The Force Awakens. Their gender 
is male. They have the following hair colours: white. They have a mass of 
17. Their skin colours are green."

令人惊叹!尽管我们在文本中提到了“最古老”,但我们没有说“智慧”或“男人”! 我希望你能看到这如何能帮助你获取数据的高质量语义索引!

结论

我们还添加了查找邻近文档和搜索整个文档库以找到重复项的端点。后者在一些基准测试上使用过,表现令人钦佩。我们希望很快在这里展示这些实验的结果。虽然在野外有很棒的向量数据库,如Pinecone,但我们希望有一个能与TerminusDB良好集成的副本,可以供主要关心内容的不那么技术化的用户使用,他们不会去自己搭建向量数据库。我们对这个VectorLink的潜力非常兴奋,并且希望大家能看看我们到目前为止所做的工作!请原谅我们的错误处理相对稀疏。我们正在疯狂地进行这方面的工作!

作者:Gavin Mendel-Gleason

更多技术干货请关注公号“云原生数据库

squids.cn,目前可体验全网zui低价RDS,免费的迁移工具DBMotion、SQL开发工具等