Faiss 向量数据库

发表于 2024-09-11 更新于 2024-10-11 分类于 LLM开发

faiss 索引类型、代码实践

Faiss 索引类型

IndexFlatIP 内积：方向；

IndexFlatL2 L2距离：几何距离；

平面索引 FLAT

平面索引就是暴力搜索，将 queries 于 database中的所有向量计算，时间复杂度是 $O({n^2})$。

不同索引计算相似度的方法不同，IndexFlatIP 是点积，IndexFlatL2 是 L2 距离，如果需要计算余弦相似度则需要在点积前 L2 归一化 faiss.normalize_L2(vector)，L2 归一化公式为 $\mathbf{X}=\left(\frac{x_1}{\left|\mathbf{x}\right|},\frac{x_2}{\left|\mathbf{x}\right|},\cdots,\frac{x_n}{\left|\mathbf{x}\right|}\right)$

index = faiss.IndexFlatL2(d)
index.add(database)
D, I = index.search(queries[:5], k)

$D$ 表示具体的相似度D.shape = (query_len, similarity)，取值范围是 [0, 1]

$I$ 表示与 $D$ 相对应的索引 I.shape = (query_len, index)，取值范围是 [0, data_num]

分区索引 IVF

对索引数据进行分区优化，将数据通过“沃罗诺伊图单元”（也被叫做泰森多边形）来进行切割（类似传统数据库分库分表）。在搜索时，仅将查询 queries 所在单元中包含的数据库向量 database 以及一些相邻的向量与查询向量进行比较，找到 queries 向量所属的泰森多边形单元就是在质心集中找到该向量的最近邻居。

index.train()训练了一个聚类模型，如 K-means。生成多个聚类中心，聚类的数量由索引的 nlist 参数决定。

倒排索引（IVF）把“文档→单词”的形式变为“单词→文档”的形式。

search 有两个参数： nlist （单元格数量）和 nprobe （为执行搜索而访问的质心集单元数量，在nlist之外），设置 nprobe = nlist 给出与暴力搜索相同的结果。

quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
index.train(database)
index.add(database)
index.nprobe = 10
D, I = index.search(queries, k)

量化索引 PQ

基于乘积量化器的有损压缩来压缩存储的向量，压缩方法和量化级别参考论文。

建立索引时：

向量被分配到最近的聚类（同 IVF）
每个向量被划分为 m 个子向量，分别在各自的子空间中进行量化，存储为对应的量化代码

查询时：

查询向量通过倒排索引（IVF）来确定最接近的聚类中心
查询向量被划分为 m 个子向量
查询子向量与存储的量化代码进行匹配，基于每个子空间的距离查表来计算近似距离
将所有子空间的距离相加，得到查询向量与候选向量的总距离

这个例子，将 64 个 32 位浮点数压缩为 8 个字节，因此压缩因子为 32。

# d 必须是 m 的倍数
quantizer = faiss.IndexFlatL2(d)
# 量化器数量 m、指定每个子向量编码为 8 byte(32 bit)
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8)
# 可以简写为 index = faiss.index_factory(d, "IVF100,PQ8")
index.train(database)
index.add(database)
index.nprobe = 10
D, I = index.search(queries, k)

功能

索引的量化功能

对向量进行编码量化

index.sa_encode(vectors)

针对非量化索引，也会返回 uint8 类型的更紧凑的离散表示

xq[0][:5]
## array([0.81432974, 0.7409969 , 0.8915324 , 0.02642949, 0.24954738], dtype=float32)
index.sa_encode(xq[:1, :])[0, :5]
## array([ 93, 234, 119,  80,  63], dtype=uint8)

获取原始数据/向量

基于根据 index

index.reconstruct(vector_id)

为了追求更快的查询性能，IndexIVF 和向量 ID IndexBinaryIVF 都存储在倒排列表中，无法通过向量 ID 来反查索引中的内容，如果我们想要得到某个数据的内容，需要手动重建索引

index.make_direct_map()
index.reconstruct(vector_id)

获取原文档

FAISS 内置了 docstore 和 index_to_docstore_id ，在每次加入向量时同时存储 doc 内容，且以字典的形式储存 index 与 doc 的映射关系，方便后续查找原文档。

BM25 实现

额外记录一下 BM25 的实现方法

import jieba
from langchain.schema import Document
from langchain_community.retrievers import BM25Retriever

def create_bm25_retriever(docs, top_k):
    bm25_retriever = BM25Retriever.from_documents(
        docs,
        preprocess_func=jieba.lcut_for_search,
    )
    bm25_retriever.k = top_k
    return bm25_retriever

if __name__ == "__main__":
    docs = [
        Document(page_content="这是一篇关于机器学习的文档。"),
        Document(page_content="自然语言处理是人工智能的一个重要领域。"),
        Document(page_content="深度学习是机器学习的一个子集。"),
    ]
    top_k = 2
    bm25_retriever = create_bm25_retriever(docs, top_k)
    query = "机器学习"
    results = bm25_retriever.get_relevant_documents(query)
    for rank, doc in enumerate(results, start=1):
        print(f"Rank {rank}: {doc}")

Debug 查看分词结果 bm25_retriever.vectorizer.doc_freqs