使用 Annoy 和 Sentence Transformers 构建本地向量数据库索引

如果您一直依赖云向量数据库来管理和搜索高维数据，让我向您展示另一种选择：在本地部署您自己的向量数据库。这不仅让您拥有完全的控制权，还能消除重复的云成本和潜在的数据隐私问题。我将使用 Annoy 和 Sentence Transformers 向您展示如何构建一个快速高效的向量数据库，以适应您自己的嵌入模型。

---

什么是向量数据库索引？

向量数据库索引用于组织高维数据（例如，嵌入），以实现高效的相似性搜索。它是语义搜索、推荐引擎和个性化系统等应用程序的重要组成部分。

在本指南中，我将演示如何：

• 使用 Annoy 进行近似最近邻 (ANN) 搜索。
• 使用 Sentence Transformers 生成文本嵌入。

结果呢？一个功能强大的本地设置，可以在您自己的机器上完全运行。

---

先决条件

要跟着操作，请确保您已具备：

• 已安装 Python。
• 基本的 Python 编程技能。
• 所需的库：

  pip install annoy sentence-transformers
  

---

第一步：生成嵌入

嵌入是任何向量数据库的支柱。在这里，我使用 Sentence Transformers 模型从文本生成嵌入。

from sentence_transformers import SentenceTransformer

# Load the embeddings model
model = SentenceTransformer("jinaai/jina-embeddings-v2-base-en", trust_remote_code=True)

# Optional: Set the maximum sequence length
model.max_seq_length = 1024

def get_embedding(texts):
    """
    Generate embeddings for a list of texts.
    """
    embeddings = model.encode(texts)
    return embeddings.tolist()

---

第二步：管理元数据和文件路径

每个嵌入都与元数据（如标题或内容）相关联，这些元数据存储在 JSON 文件中，以便于管理。

import os
import json

def get_file_paths(email, db_name):
    """
    Generate paths for the index and metadata files.
    """
    base_dir = os.path.join('user_data', email, db_name)
    index_path = os.path.join(base_dir, 'vector_index.ann')
    metadata_path = os.path.join(base_dir, 'metadata.json')
    return base_dir, index_path, metadata_path

def load_metadata(email, db_name):
    """
    Load metadata from a JSON file.
    """
    _, _, metadata_path = get_file_paths(email, db_name)
    if os.path.exists(metadata_path):
        with open(metadata_path, 'r') as f:
            return json.load(f)
    return {}

def save_metadata(email, db_name, metadata):
    """
    Save metadata to a JSON file.
    """
    base_dir, _, metadata_path = get_file_paths(email, db_name)
    os.makedirs(base_dir, exist_ok=True)
    with open(metadata_path, 'w') as f:
        json.dump(metadata, f, indent=2)

---

第三步：使用 Annoy 构建索引

Annoy 允许您为高效的近似最近邻搜索创建索引。以下是如何在本地构建和保存索引。

from annoy import AnnoyIndex

num_trees = 10
num_dimensions = 768  # Match the dimension of your embeddings

def build_index(email, db_name):
    """
    Build an Annoy index from metadata and save it locally.
    """
    base_dir, index_path, _ = get_file_paths(email, db_name)
    index = AnnoyIndex(num_dimensions, 'angular')
    metadata = load_metadata(email, db_name)

    for item_id, item in metadata.items():
        try:
            index.add_item(int(item_id), item['vector'])
        except KeyError as e:
            print(f"Error adding item {item_id} to index: {e}")

    index.build(num_trees)
    index.save(index_path)
    print('Index built and saved.')

---

第四步：查询索引

一旦构建了索引，您就可以使用嵌入来查询它，以找到最相似的项目。

def query_vector(email, db_name, vector, num_neighbors=5):
    """
    Query the Annoy index with a vector and return the closest neighbors.
    """
    index = load_index(email, db_name)
    metadata = load_metadata(email, db_name)

    if index.get_n_items() == 0:
        print('Index is empty or not loaded.')
        return []

    neighbors = index.get_nns_by_vector(vector, num_neighbors, include_distances=True)
    results = []
    for idx, distance in zip(neighbors[0], neighbors[1]):
        results.append({
            'id': idx,
            'title': metadata.get(str(idx), {}).get('title', 'N/A'),
            'content': metadata.get(str(idx), {}).get('content', 'N/A'),
            'distance': distance
        })
    return results

---

第五步：运行完整流程

让我向您展示这些步骤是如何协同工作的。

1. 生成嵌入:

   texts = ["Example text 1", "Another example text"]
   vectors = get_embedding(texts)
   

2. 添加数据:

   email = "user@example.com"
   db_name = "my_database"
   
   for vector, text in zip(vectors, texts):
       add_vector(email, db_name, vector, title="Title", content=text)
   

3. 构建索引:

   build_index(email, db_name)
   

4. 查询索引:

   query_vector(email, db_name, vectors[0], num_neighbors=3)
   

---

为何选择 Annoy 而非基于云的解决方案？

使用 Annoy 进行本地索引有几个优点：

• 成本效益：无需为云向量数据库服务支付经常性费用。
• 隐私：您的数据永远不会离开您的机器。
• 性能：针对快速、内存中的近似最近邻搜索进行了优化。

这种方法特别适用于可扩展性和控制至关重要的项目。

---

本指南的扩展

您可以调整此设置以适应更复杂的用例：

1. 使用特定领域模型进行嵌入。
2. 集成 SQLite 等持久化存储来管理元数据。
3. 尝试 Annoy 参数（例如 num_trees）以获得更好的性能。

---

通过部署您自己的向量数据库索引，您可以控制整个流程——从嵌入生成到高效检索。此设置非常适合希望构建可扩展且注重隐私的应用程序，而无需依赖第三方云解决方案的任何人。借助这些工具，您现在拥有了支持语义搜索、推荐引擎或任何基于嵌入的应用程序的基础。