BinaryVectorDB - 大数据集的高效搜索

此存储库包含一个用于在大数据集上高效搜索的二进制向量数据库，旨在教育目的。

大多数嵌入模型将它们的向量表示为float32：这些消耗大量内存，在这些上进行搜索非常慢。在Cohere，我们引入了第一个具有本地int8和二进制支持的嵌入模型，这为您提供了极好的搜索质量，成本仅为一小部分

安装

安装非常简单

pip install BinaryVectorDB

要使用以下示例中的某些示例，您需要一个来自 https://cohere.com/ 的 Cohere API密钥（免费或付费）。您必须将此API密钥设置为一个环境变量：export COHERE_API_KEY=your_api_key

用法 - 加载现有的二进制向量数据库

我们将在以后讨论如何构建自己的向量数据库。首先，让我们使用一个预构建的二进制向量数据库。我们在 https://hugging-face.cn/Cohere/BinaryVectorDB 上托管各种预构建的数据库。您可以下载这些数据库并在本地使用。

让我们从维基百科的简单英文版本开始

wget https://hugging-face.cn/datasets/Cohere/BinaryVectorDB/resolve/main/wikipedia-2023-11-simple.zip

然后解压缩此文件

unzip wikipedia-2023-11-simple.zip

加载向量数据库

您可以通过将指针指向上一步中解压缩的文件夹来轻松地加载数据库

from BinaryVectorDB import BinaryVectorDB

# Point it to the unzipped folder from the previous step
# Ensure that you have set your Cohere API key via: export COHERE_API_KEY=<<YOUR_KEY>>
db = BinaryVectorDB("wikipedia-2023-11-simple/")

query = "Who is the founder of Facebook"
print("Query:", query)
hits = db.search(query)
for hit in hits[0:3]:
    print(hit)

数据库有646,424个嵌入项，总大小为962 MB。然而，仅在内存中加载了80 MB的二进制嵌入项。文档及其int8嵌入项保存在磁盘上，仅当需要时才加载。

这种将二进制嵌入项存储在内存中而将int8嵌入项和文档存储在磁盘上的分割，使我们能够在不需要大量内存的情况下扩展到非常大的数据集。

构建自己的二进制向量数据库

构建自己的二进制向量数据库非常简单。

from BinaryVectorDB import BinaryVectorDB
import os
import gzip
import json

simplewiki_file = "simple-wikipedia-example.jsonl.gz"

#If file not exist, download
if not os.path.exists(simplewiki_file):
    cmd = f"wget https://hugging-face.cn/datasets/Cohere/BinaryVectorDB/resolve/main/simple-wikipedia-example.jsonl.gz"
    os.system(cmd)

# Create the vector DB with an empty folder
# Ensure that you have set your Cohere API key via: export COHERE_API_KEY=<<YOUR_KEY>>
db_folder = "path_to_an_empty_folder/"
db = BinaryVectorDB(db_folder)

if len(db) > 0:
    exit(f"The database {db_folder} is not empty. Please provide an empty folder to create a new database.")

# Read all docs from the jsonl.gz file
docs = []
with gzip.open(simplewiki_file) as fIn:
    for line in fIn:
        docs.append(json.loads(line))

#Limit it to 10k docs to make the next step a bit faster
docs = docs[0:10_000]

# Add all documents to the DB
# docs2text defines a function that maps our documents to a string
# This string is then embedded with the state-of-the-art Cohere embedding model
db.add_documents(doc_ids=list(range(len(docs))), docs=docs, docs2text=lambda doc: doc['title']+" "+doc['text'])

文档可以是任何Python可序列化对象。您需要提供一个函数供docs2text使用，将您的文档映射到一个字符串。在上面的示例中，我们将标题和文本字段连接起来。该字符串被发送到嵌入模型以生成所需的文本嵌入。

更新和删除文档

请参阅 examples/add_update_delete.py 以获取如何在数据库中添加/更新/删除文档的示例脚本。

这是一个真正的向量数据库吗？

实际上并不是。该存储库主要用于教育目的，以展示如何扩展到大型数据集的技术。重点更多地放在易用性上，而实现中缺少了一些关键方面，如多进程安全性、回滚等。

如果您真的想要投入生产，请使用像 Vespa.ai 这样的合适向量数据库，它允许您实现类似的结果。