Faiss – 向量数据库

内容纲要

官网:https://faiss.ai/

Github: https://github.com/facebookresearch/faiss

A library for efficient similarity search and clustering of dense vectors.

用于高效相似性搜索和密集向量聚类的库。

什么是Faiss

Faiss 是一个用于高效相似度搜索和密集向量聚类的库。它包含了在任意大小的向量集合中进行搜索的算法,甚至可以处理那些可能无法完全放入内存的数据集。Faiss 还包含用于评估和参数调优的辅助代码。

Faiss 使用 C++ 编写,并提供了完整的 Python 封装。一些最有用的算法已经在 GPU 上实现。这个库主要由 Meta 的基础 AI 研究团队 FAIR 进行开发。

Faiss 的研究基础

Faiss 基于多年的研究。特别是它实现了以下内容:

  • 来自“Sivic & Zisserman, ICCV 2003”的反向文件。这是大数据集非穷尽搜索的关键。否则,所有搜索都需要扫描索引中的所有元素,即使每个元素的操作很快,这也是不可行的。
  • 来自“Jégou et al., PAMI 2011”的产品量化(PQ)方法。这可以看作是高维向量的有损压缩技术,允许在压缩域中进行相对准确的重建和距离计算。
  • 来自“Tavenard et al., ICASSP'11”的三层量化(IVFADC-R aka IndexIVFPQR)方法。
  • 来自“Babenko & Lempitsky, CVPR 2012”的反向多重索引。这种方法极大地提高了快速/不太精确的操作点的反向索引速度。
  • 来自“He et al., CVPR 2013”的优化 PQ 方法。这种方法可以看作是向量空间的线性变换,使其更适合于产品量化器的索引。
  • 来自“Douze et al., ECCV 2016”的产品量化器距离的预过滤技术。该技术在计算 PQ 距离之前执行二进制过滤阶段。
  • 来自“Johnson et al., ArXiv 1702.08734, 2017”的 GPU 实现和快速 k 选择。
  • 来自“Malkov et al., ArXiv 1603.09320, 2016”的 HNSW 索引方法。
  • 来自“André et al., PAMI'19”中的寄存器内向量比较,也用于“Accelerating Large-Scale Inference with Anisotropic Vector Quantization”中。
  • 来自“Norouzi et al., CVPR'12”的二进制多重索引哈希方法。
  • 来自“Cong Fu et al., VLDB 2019”的基于图的索引方法 NSG。
  • 来自“Julieta Martinez et al. ECCV 2016”的局部搜索量化方法和“Julieta Martinez et al. ECCV 2018”的 LSQ++:在多码本量化中降低运行时间和提高召回率。
  • 来自“Shicong Liu et al, AAAI'15”的残差量化器实现。
  • 来自“Yusuke Matsui et al., ITE transactions on MTA, 2018”的产品量化及相关方法的一般论文:“A Survey of Product Quantization”。

下图来自该论文(点击图片放大):

Faiss 中实现的方法以红色突出显示。

img

图片来源:Yusuke Matsui

什么是相似度搜索?

给定一个维度为 𝑑 的向量集合 𝑥𝑖,Faiss 从中在内存中构建一个数据结构。在构建好该结构后,当给定一个新的维度为 𝑑 的向量 𝑥 时,它可以高效地执行以下操作:

$$j = argmin_i \lVert x - x_i \rVert$$

其中 ∥⋅∥ 是欧几里得距离($$L_2$$)。

在 Faiss 的术语中,该数据结构称为索引(index),这是一个具有添加 $$x_i$$ 向量的方法的对象。需要注意的是,这些 $$x_i$$ 是假设为固定的。

计算 $$arg⁡ min$$ 是在索引上进行的搜索操作。

这就是 Faiss 的全部功能。此外,它还可以:

  • 返回不仅仅是最近的邻居,还可以返回第 2 近、第 3 近,...,第 k 近的邻居。
  • 同时搜索多个向量而不是一个(批处理)。对于许多索引类型,这比逐个搜索向量要快。
  • 通过牺牲精度来提高速度,即以 10% 的错误率使用速度快 10 倍或内存占用少 10 倍的方法。
  • 执行最大内积搜索 $$arg max_i⟨x,xi⟩$$ 而不是最小欧几里得距离搜索。也有限支持其他距离(L1,Linf 等)。
  • 返回在查询点给定半径内的所有元素(范围搜索)。
  • 将索引存储在磁盘上而不是内存中。
  • 索引二进制向量而不是浮点向量。
  • 根据向量 ID 上的谓词忽略索引向量的子集。

主要功能

基础功能

  1. 最近邻搜索:Faiss能够有效地在大规模数据集中进行最近邻搜索(KNN)。它利用了各种算法和索引结构来提高搜索速度和准确性。
  2. 向量量化:Faiss支持多种向量量化方法,可以显著减少存储需求,同时保持较高的搜索准确度。
  3. GPU**加速**:为了处理更大规模的数据集,Faiss提供了对GPU的支持,使得搜索过程可以充分利用并行计算的优势,大幅提升性能。
  4. 多索引类型:Faiss提供了多种索引类型,包括平面索引(Flat Index)、聚类索引(Clustered Index)和压缩索引(Compressed Index)等,以适应不同的应用需求。

高级功能

  1. 层次化**聚类**:Faiss支持层次化聚类算法,可以用于数据集的分层次聚类分析。
  2. 产品量化:产品量化(PQ)是一种高效的向量压缩技术,Faiss实现了多种PQ变体,适用于不同的应用场景。
  3. 索引合并与分片:Faiss支持将多个索引合并为一个索引,或者将一个大索引分成多个小索引,方便分布式处理。

如何安装

使用pip安装

pip install faiss-cpu  # 安装CPU版本
pip install faiss-gpu  # 安装GPU版本

img

从源码编译安装

git clone https://github.com/facebookresearch/faiss.git
cd faiss
./configure
make

使用Conda安装

推荐通过 Conda 安装 Faiss:

$ conda install -c pytorch faiss-cpu

faiss-gpu 包提供了 CUDA 支持的索引:

$ conda install -c pytorch faiss-gpu

注意,应安装其中一个包,而不是同时安装,因为后者是前者的超集。

超集 Superset

超集是指一个集合中包含另一个集合的所有元素,且可能还包含其他元素。如果一个集合中一定有另一个集合没有的元素,则该集合为真超集。反之,如果一个集合是另一个集合的超集,则另一个集合是该集合的子集。

如何使用

☆ 简单使用示例

以下是一个简单的使用示例,展示了如何使用Faiss进行最近邻搜索:

安装依赖环境

pip install numpy

引入依赖包

import numpy as np
import faiss

生成随机数据

# 生成随机数据
d = 64                           # 向量维度
nb = 100000                      # 数据库中的向量数量
nq = 10000                       # 查询向量数量

np.random.seed(1234)             # 固定随机种子
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000.

建立索引

# 建立索引
index = faiss.IndexFlatL2(d)     # 使用L2距离的平面索引
print(index.is_trained)
index.add(xb)                    # 向索引中添加向量
print(index.ntotal)

img

执行搜索

# 执行搜索
k = 4                            # 搜索最近的4个邻居
D, I = index.search(xq, k)       # 执行搜索
print(I[:5])                     # 输出前5个查询的结果
print(D[:5])                     # 输出前5个查询的距离

img

☆ 基本最近邻搜索

这个示例展示了如何使用 Faiss 进行基本的最近邻搜索。

引入依赖包

import numpy as np
import faiss

生成随机数据、随机生成数据库向量和查询向量

# 生成随机数据
d = 64  # 向量维度
nb = 10000  # 数据库中的向量数量
nq = 5  # 查询向量数量

# 随机生成数据库向量和查询向量
np.random.seed(1234)
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')

创建一个使用L2距离的平面索引

index = faiss.IndexFlatL2(d)
print("Index is trained:", index.is_trained)

img

向索引中添加向量

# 向索引中添加向量
index.add(xb)
print("Total number of vectors in index:", index.ntotal)

img

执行搜索,寻找每个查询向量的4个最近邻

k = 4
D, I = index.search(xq, k)

输出搜索结果

# 输出搜索结果
print("Indices of nearest neighbors:\n", I)
print("Distances to nearest neighbors:\n", D)

img

☆ GPU 加速的最近邻搜索

GPU

Graphics processor / Graphics Processing Unit / Graphics Processor / 图形处理器 / 显卡 / 视觉处理器

图形处理器是计算机中的一种专用微处理器,用于处理图像和图形的演算、计算和显示操作。它主要用于视频游戏、3D 渲染、高清影片解码和显示等领域,能够实现复杂的图形和图像处理,比一般的微处理器速度更快。GPU通常与中央处理器(CPU)组合使用,共同处理复杂的计算任务。随着近年来GPU计算技术的发展,GPU不仅仅局限于图像和图形处理,还能承担更多的单机计算任务,如科学计算、深度学习、人工智能等。

这个示例展示了如何在 GPU 上运行 Faiss 以加速最近邻搜索。

import numpy as np
import faiss

# 检查是否有可用的 GPU
if faiss.get_num_gpus() == 0:
    raise RuntimeError("No GPU found. This example requires at least one GPU.")

# 生成随机数据
d = 64  # 向量维度
nb = 100000  # 数据库中的向量数量
nq = 10000  # 查询向量数量

# 随机生成数据库向量和查询向量
np.random.seed(1234)
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')

# 创建一个使用L2距离的 GPU 索引
res = faiss.StandardGpuResources()  # 使用默认资源
index_flat = faiss.IndexFlatL2(d)
index = faiss.index_cpu_to_gpu(res, 0, index_flat)

# 向索引中添加向量
index.add(xb)

# 执行搜索,寻找每个查询向量的4个最近邻
k = 4
D, I = index.search(xq, k)

# 输出搜索结果
print("Indices of nearest neighbors:\n", I[:5])  # 只输出前5个查询向量的结果
print("Distances to nearest neighbors:\n", D[:5])

img

☆ 范围搜索

这个示例展示了如何使用 Faiss 进行范围搜索,找到在给定半径内的所有向量。

import numpy as np
import faiss

# 生成随机数据
d = 16  # 向量维度减少到16
nb = 1000  # 数据库中的向量数量减少到1000,使向量更密集
nq = 50  # 查询向量数量,为了演示搜索出来的向量结果,设置了比较大的数字50,可以设置小一些,如5

# 随机生成数据库向量和查询向量
np.random.seed(1234)
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')

# 创建一个使用L2距离的平面索引
index = faiss.IndexFlatL2(d)

# 向索引中添加向量
index.add(xb)

# 定义搜索半径,适当增大
radius = 0.5  # 增加到0.5,但仍然适中,根据数据调整

# 设置线程数
faiss.omp_set_num_threads(4)

# 执行范围搜索
lims, D, I = index.range_search(xq, radius)

# 格式化输出搜索结果
n_neighbors = np.diff(np.hstack((lims, [len(D)])))  # 计算每个查询向量找到的邻居数
print("Number of neighbors found for each query:", n_neighbors)

for i in range(len(n_neighbors)):
    if n_neighbors[i] == 0:
        print(f"\nQuery {i} has no neighbors within the radius.")
    else:
        start_idx = lims[i]
        end_idx = lims[i + 1]
        print(f"\nQuery {i}:")
        print(" Indices:", I[start_idx:end_idx])
        print(" Distances:", D[start_idx:end_idx])

img

☆ 批处理搜索

批处理

批处理是一种计算机操作方式,它允许用户按照一定的顺序和规则将一组计算机指令集合起来,然后一次性提交给计算机系统执行,而无需用户进行干预。批处理可以帮助用户自动化和简化重复性的任务,提高计算机系统的效率和可靠性。

这个示例展示了如何使用 Faiss 进行批处理搜索,提高多个查询向量的搜索效率。

import numpy as np
import faiss

# 生成随机数据
d = 64  # 向量维度
nb = 100000  # 数据库中的向量数量
nq = 10000  # 查询向量数量

# 随机生成数据库向量和查询向量
np.random.seed(1234)
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')

# 创建一个使用L2距离的平面索引
index = faiss.IndexFlatL2(d)

# 向索引中添加向量
index.add(xb)

# 定义批处理大小
batch_size = 1000

# 批处理搜索
all_I = []
all_D = []
for i in range(0, nq, batch_size):
    xq_batch = xq[i:i+batch_size]
    D, I = index.search(xq_batch, 4)
    all_I.append(I)
    all_D.append(D)

# 合并所有批处理结果
all_I = np.vstack(all_I)
all_D = np.vstack(all_D)

# 输出搜索结果
print("Indices of nearest neighbors:\n", all_I[:5])  # 只输出前5个查询向量的结果
print("Distances to nearest neighbors:\n", all_D[:5])

img

场景示例

如何使用Faiss向量化PDF文档并查询

使用Faiss进行PDF文档的向量化并进行查询需要几个步骤,包括读取PDF文档、将文本转换为向量、构建Faiss索引,以及进行查询。以下是一个完整的示例,展示了如何实现这些步骤。

1. 安装必要的库

首先,安装所需的Python库,包括PyMuPDF(用于读取PDF)、transformers(用于将文本转换为向量),以及faiss

如果安装的是faiss-gpu,则需要再安装torch

pip install pymupdf transformers faiss-cpu torch modelscope

2. 读取PDF文档

使用PyMuPDF读取PDF文档并提取文本。

import fitz  # PyMuPDF

def extract_text_from_pdf(pdf_path):
    doc = fitz.open(pdf_path)
    text = ""
    for page_num in range(doc.page_count):
        page = doc.load_page(page_num)
        text += page.get_text()
    return text

pdf_path = "./The_Java_Virtual_Machine_Specification_8.pdf" # path/to/your/document.pdf
text = extract_text_from_pdf(pdf_path)
print(text)

3. 文本向量化

使用transformers库中的预训练模型将文本转换为向量。这里使用sentence-transformers中的模型。

如果网络环境支持,可以使用Huggingface下载

from transformers import AutoTokenizer, AutoModel
import torch

# 加载预训练模型和tokenizer
model_name = "sentence-transformers/all-MiniLM-L6-v2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

def text_to_vector(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
    embeddings = outputs.last_hidden_state.mean(dim=1)
    return embeddings.cpu().numpy()

# 将PDF文本转换为向量
vector = text_to_vector(text)
print(vector)

如果网络环境不支持,使用ModelScope下载向量模型

#SDK模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('bensonpeng/all-MiniLM-L6-v2')

img

from transformers import AutoTokenizer, AutoModel
import torch

# 加载预训练模型和tokenizer
model_name = model_dir # 与上面不同的只有这一行
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

def text_to_vector(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
    embeddings = outputs.last_hidden_state.mean(dim=1)
    return embeddings.cpu().numpy()

# 将PDF文本转换为向量
vector = text_to_vector(text)
print(vector)

img

4. 构建Faiss索引

将向量添加到Faiss索引中。

import faiss

# 创建一个Faiss索引
d = vector.shape[1]
index = faiss.IndexFlatL2(d)

# 添加向量到索引
index.add(vector)
print("Total number of vectors in index:", index.ntotal)

img

5. 查询Faiss索引

向Faiss索引中查询与给定查询文本最相似的向量。

def query_vector(text):
    vector = text_to_vector(text)
    k = 1  # 返回最相似的1个向量
    D, I = index.search(vector, k)
    return D, I

query_text = "your query text here"
distances, indices = query_vector(query_text)
print("Indices of nearest neighbors:", indices)
print("Distances to nearest neighbors:", distances)

img

一个基于Faiss的RAG示例

此示例展示了如何使用Faiss进行文档检索,并将结果与生成模型结合,以回答复杂的查询。

当然,还有许多有趣且实用的示例可以用在产品内部生产中,比如使用Faiss和RAG(Retrieval-Augmented Generation)技术进行文本检索和增强的示例。以下是一个完整的RAG示例,展示了如何使用Faiss进行文档检索,并将结果与生成模型结合,以回答复杂的查询。

1. 安装必要的库

首先,安装所需的Python库,包括transformersfaissPyMuPDF

pip install transformers faiss-cpu pymupdf

2. 读取PDF文档并进行向量化

使用PyMuPDF读取PDF文档,transformers库中的预训练模型将文本转换为向量。

import fitz  # PyMuPDF
from transformers import AutoTokenizer, AutoModel
import torch
import faiss
import numpy as np

# 提取PDF文本
def extract_text_from_pdf(pdf_path):
    doc = fitz.open(pdf_path)
    text = ""
    for page_num in range(doc.page_count):
        page = doc.load_page(page_num)
        text += page.get_text()
    return text

# 按500字符拆分文本,并上下额外覆盖50字符
def split_text(text, chunk_size=500, overlap=50):
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + chunk_size, len(text))
        chunks.append(text[max(0, start - overlap):min(len(text), end + overlap)])
        start += chunk_size
    return chunks

# 文本向量化
def text_to_vector(text, tokenizer, model):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
    embeddings = outputs.last_hidden_state.mean(dim=1)
    return embeddings.cpu().numpy()

# 使用modelscope下载模型
from modelscope import snapshot_download
model_dir = snapshot_download('bensonpeng/all-MiniLM-L6-v2')

# 加载预训练模型和tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_dir)
model = AutoModel.from_pretrained(model_dir)

# 读取PDF并向量化
pdf_path = "./The_Java_Virtual_Machine_Specification_8.pdf"
text = extract_text_from_pdf(pdf_path)

# 拆分文本为段落并向量化
paragraphs = split_text(text)
vectors = np.vstack([text_to_vector(p, tokenizer, model) for p in paragraphs])

img

img

3. 构建Faiss索引

将向量添加到Faiss索引中。

# 创建Faiss索引
d = vectors.shape[1]
index = faiss.IndexFlatL2(d)

# 添加向量到索引
index.add(vectors)
print("Total number of vectors in index:", index.ntotal)

img

4. 查询Faiss索引

使用查询文本从索引中检索最相似的向量。

def query_vector(text, tokenizer, model, index):
    vector = text_to_vector(text, tokenizer, model)
    k = 5  # 返回最相似的5个向量
    D, I = index.search(vector, k)
    return D, I

query_text = "The Java® programming language is a general-purpose," # your query text here
distances, indices = query_vector(query_text, tokenizer, model, index)
print("Indices of nearest neighbors:", indices)
print("Distances to nearest neighbors:", distances)

img

# 提取相似段落
similar_paragraphs = [paragraphs[i] for i in indices[0]]
similar_paragraphs

img

5. 使用生成模型进行RAG

使用检索到的文档内容作为上下文,结合生成模型(如GPT-3或T5)回答用户查询。

1、使用modelscope下载 ChatGLM-6B 模型

# 使用modelscope下载模型
from modelscope import snapshot_download
model_dir = snapshot_download('ZhipuAI/ChatGLM-6B')

img

2、主要流程

from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope import Model

from modelscope import snapshot_download
# 使用modelscope下载模型
model_dir = snapshot_download('ZhipuAI/ChatGLM-6B')

# 加载生成模型
model = Model.from_pretrained(model_dir, device_map='auto', revision='v1.0.19').half().cuda()

# 拼接上下文段落
context = "\n".join(similar_paragraphs)
input_text = f"question: {query_text} context: {context}"

img

创建pipeline并生成回答

pipe = pipeline(task=Tasks.chat, model=model, max_length=512)
inputs = {'text': input_text, 'history': []}
result = pipe(inputs)

img

print("Generated Answer:", result['response'])

img

Leave a Comment

您的电子邮箱地址不会被公开。 必填项已用*标注

close
arrow_upward