annoy

annoy是由 spotify 开发的用于在高维向量中快速查找相似近邻(Approximate Nearest Neighbors, ANN)。他特别适用于大规模数据集上进行向量相似性搜索。

Annoy 原理

Annoy 的核心算法是基于随机投影（Random Projection）和树结构（Tree Structure）来实现的。Annoy 会创建多个二叉树，每个树通过随机投影对数据进行分割，这样当我们进行查询时，便可以从多个方向来进行搜索，从而提高搜索的准确性和效率。

整个构建树的流程如下:

随机选择超平面: 在构建树的过程中，从数据中随机选取一个向量方向(超平面)并计算投影，按照这个方向将数据分成两部分
递归分割: 不断递归进行这种划分，直到每个叶节点中剩下的点数达到预定阈值（例如 10 个点）
构建多棵树: 为了确保查找的准确性，Annoy 会构建多棵树，每棵树使用不同的随机投影方向。查询时会同时遍历多个树，以确保得到更精确的结果

整个搜索过程同样是基于树来进行:

递归搜索树: 从根节点开始，Annoy 通过投影方向判断查询向量落在哪一边的子空间，并递归地向下搜索，直到到达叶节点
从叶节点中选取候选集: 在叶节点，Annoy 会返回所有包含在该节点中的向量作为候选集
遍历多棵树: Annoy 会对所有树进行搜索，并对每棵树的结果进行合并。合并后，对这些候选向量计算精确距离，并返回最近邻结果

通过先构建在搜索 Annoy 就能够快速返回近似最近邻。他使用多棵树来减少随机投影时引入的误差，从而提高结果的准确性。理论上构建的树数量越多(n_trees)，查询的结果准确性也越高，但是伴随的就是构建和内存占用也会增加，搜索速度也会变慢，因此通常通过实验来调整 n_trees 来平衡构建时间、内存消耗和查询精度。

Annoy 使用流程

Annoy 的基本使用流程:

创建高维向量，例如针对于文本的搜索我们通常使用 TF-IDF 来构建稀疏矩阵作来作为输入源
构建索引(AnnoyIndex(n, 'angular'))，对于文本数据通常选择余弦距离(angular)来构建索引
插入数据(index.add_item(key, vector))，注意 Annoy 仅支持整数键
构建树(index.build(n_trees))，这其中的核心就是 n_trees 的确定
为了以后复用通常可以将训练数据保存到磁盘中(index.save("annoy_index.ann"))，之后我们可以通过 index.load('annoy_index.ann') 来加载而不需要重复构建
查询最近邻，index.get_nns_by_vector(query_vector, k: int) 其中 k 表示最接近的 10 个向量

Python

from annoy import AnnoyIndex

# 传入特征值，vectorizer.get_feature_name_out()
# X.shape[1] 也可以
# 第二个是距离算法, angular 就是余弦，可以测试下结果
annoy_index = AnnoyIndex(X.shape[1], 'angular')

# 遍历矩阵，将每一个行数据都添加到 Annoy 索引中
for i in range(X.shape[0]):
    # 注意 TfidfVectorizer 返回的是稀疏矩阵
    # annoy 只能使用密集矩阵，因此这样是它非常耗费内存的原因
    annoy_index.add_item(i, X[i].toarray()[0])

# 创建索引树，树数量越多越精确但是也越慢
annoy_index.build(10)

# 之后就可以查找近邻了
# 查询
query = '需要查找的字符串'
# 注意需要具有同样特征根的 TfidfVectorizer 来 transform 并转换为密集矩阵
query_vector = vectorizer.transform([query]).toarray()[0]# 转换查询为 TF-IDF 向量
n_neighbors = 5
nearest_neighbors = annoy_index.get_nns_by_vector(query_vector, n_neighbors)

# 可以通过 source_data[nearest_neighbors] 来获取最近邻
print("最近邻的索引:", nearest_neighbors)


# 还可以保存索引在下次使用
annoy_index.save('annoy_index.ann')

# 加载索引
loaded_index = AnnoyIndex(n_features, 'euclidean')
loaded_index.load('annoy_index.ann')  # 记得指定相同的距离度量

# 但是需要注意 tf-idf 需要相同的特征根因此还需要保存 vectorizer 对象,这可以通过 pickle 来实现

数据分块

Annoy 主要是设计用于内存中的近似最近邻搜索，并且它只支持密集矩阵，但你可以通过将数据分块并在需要时加载特定部分来实现一种“磁盘支持”搜索的效果。以下是实现方法的步骤:

Python

import numpy as np
from annoy import AnnoyIndex
import os

# 示例数据
n_samples = 10000  # 总样本数量
n_features = 10    # 特征数量
X = np.random.rand(n_samples, n_features)  # 随机生成数据

# 定义块大小
block_size = 2000
n_blocks = n_samples // block_size + (n_samples % block_size > 0)

# 创建和保存多个 Annoy 索引
for block_id in range(n_blocks):
    annoy_index = AnnoyIndex(n_features, 'euclidean')
    start = block_id * block_size
    end = min(start + block_size, n_samples)

    # 添加数据到当前块的索引
    for i in range(start, end):
        annoy_index.add_item(i, X[i])

    # 建立索引并保存到磁盘
    annoy_index.build(10)
    annoy_index.save(f'annoy_index_block_{block_id}.ann')

# 查询
def query_annoy(query_vector):
    nearest_neighbors = []

    # 逐块加载索引并查询
    for block_id in range(n_blocks):
        annoy_index = AnnoyIndex(n_features, 'euclidean')
        annoy_index.load(f'annoy_index_block_{block_id}.ann')

        # 进行查询
        neighbors = annoy_index.get_nns_by_vector(query_vector, 5)
        nearest_neighbors.extend(neighbors)

    return nearest_neighbors

# 示例查询
query_vector = np.random.rand(n_features)
result = query_annoy(query_vector)
print("最近邻的索引:", result)