向量索引工具（Vector Indexing Tools）

# 三个核心库（标注版本号，基于 2026 年 3 月稳定版）
pip install faiss-cpu==1.9.0
pip install hnswlib==0.8.0
pip install annoy==1.17.3
pip install numpy==1.26.4

# 可选：GPU 加速版（需要 CUDA 环境）
pip install faiss-gpu==1.9.0

import numpy as np
import faiss
import hnswlib
from annoy import AnnoyIndex
import time

# 准备数据：1 万个 128 维向量
np.random.seed(42)
d = 128          # 向量维度
nb = 10000       # 库向量数量
nq = 5           # 查询向量数量
xb = np.random.random((nb, d)).astype("float32")
xq = np.random.random((nq, d)).astype("float32")

# === FAISS：用 IVF 倒排索引 ===
nlist = 100  # 聚类数
quantizer = faiss.IndexFlatL2(d)
index_faiss = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
index_faiss.train(xb)   # 训练聚类中心
index_faiss.add(xb)     # 添加向量
index_faiss.nprobe = 10  # 查询时搜索 10 个最近的簇

start = time.time()
distances, indices = index_faiss.search(xq, k=5)
t_faiss = (time.time() - start) * 1000
print(f"FAISS IVF  | 耗时: {t_faiss:.2f}ms | Top-1 索引: {indices[0][0]}")

# === HNSWlib：图索引 ===
index_hnsw = hnswlib.Index(space="l2", dim=d)
index_hnsw.init_index(max_elements=nb, ef_construction=200, M=16)
index_hnsw.add_items(xb, np.arange(nb))
index_hnsw.set_ef(50)  # 查询时的搜索宽度

start = time.time()
labels, dists = index_hnsw.knn_query(xq, k=5)
t_hnsw = (time.time() - start) * 1000
print(f"HNSWlib    | 耗时: {t_hnsw:.2f}ms | Top-1 索引: {labels[0][0]}")

# === Annoy：随机投影树 ===
index_annoy = AnnoyIndex(d, "euclidean")
for i in range(nb):
    index_annoy.add_item(i, xb[i])
index_annoy.build(n_trees=10)  # 构建 10 棵树

start = time.time()
for q in xq:
    index_annoy.get_nns_by_vector(q.tolist(), 5, include_distances=True)
t_annoy = (time.time() - start) * 1000
print(f"Annoy      | 耗时: {t_annoy:.2f}ms | Top-1 索引: {index_annoy.get_nns_by_vector(xq[0].tolist(), 1)[0]}")

FAISS IVF  | 耗时: 1.23ms | Top-1 索引: 5765
HNSWlib    | 耗时: 0.45ms | Top-1 索引: 5765
Annoy      | 耗时: 2.10ms | Top-1 索引: 5765