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 中实现的方法以红色突出显示。
图片来源: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 上的谓词忽略索引向量的子集。
主要功能
基础功能
- 最近邻搜索:Faiss能够有效地在大规模数据集中进行最近邻搜索(KNN)。它利用了各种算法和索引结构来提高搜索速度和准确性。
- 向量量化:Faiss支持多种向量量化方法,可以显著减少存储需求,同时保持较高的搜索准确度。
- GPU**加速**:为了处理更大规模的数据集,Faiss提供了对GPU的支持,使得搜索过程可以充分利用并行计算的优势,大幅提升性能。
- 多索引类型:Faiss提供了多种索引类型,包括平面索引(Flat Index)、聚类索引(Clustered Index)和压缩索引(Compressed Index)等,以适应不同的应用需求。
高级功能
- 层次化**聚类**:Faiss支持层次化聚类算法,可以用于数据集的分层次聚类分析。
- 产品量化:产品量化(PQ)是一种高效的向量压缩技术,Faiss实现了多种PQ变体,适用于不同的应用场景。
- 索引合并与分片:Faiss支持将多个索引合并为一个索引,或者将一个大索引分成多个小索引,方便分布式处理。
如何安装
使用pip
安装
pip install faiss-cpu # 安装CPU版本
pip install faiss-gpu # 安装GPU版本
从源码编译安装
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)
执行搜索
# 执行搜索
k = 4 # 搜索最近的4个邻居
D, I = index.search(xq, k) # 执行搜索
print(I[:5]) # 输出前5个查询的结果
print(D[:5]) # 输出前5个查询的距离
☆ 基本最近邻搜索
这个示例展示了如何使用 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)
向索引中添加向量
# 向索引中添加向量
index.add(xb)
print("Total number of vectors in index:", index.ntotal)
执行搜索,寻找每个查询向量的4个最近邻
k = 4
D, I = index.search(xq, k)
输出搜索结果
# 输出搜索结果
print("Indices of nearest neighbors:\n", I)
print("Distances to nearest neighbors:\n", D)
☆ 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])
☆ 范围搜索
这个示例展示了如何使用 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])
☆ 批处理搜索
批处理
批处理是一种计算机操作方式,它允许用户按照一定的顺序和规则将一组计算机指令集合起来,然后一次性提交给计算机系统执行,而无需用户进行干预。批处理可以帮助用户自动化和简化重复性的任务,提高计算机系统的效率和可靠性。
这个示例展示了如何使用 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])
场景示例
如何使用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')
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)
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)
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)
一个基于Faiss的RAG示例
此示例展示了如何使用Faiss进行文档检索,并将结果与生成模型结合,以回答复杂的查询。
当然,还有许多有趣且实用的示例可以用在产品内部生产中,比如使用Faiss和RAG(Retrieval-Augmented Generation)技术进行文本检索和增强的示例。以下是一个完整的RAG示例,展示了如何使用Faiss进行文档检索,并将结果与生成模型结合,以回答复杂的查询。
1. 安装必要的库
首先,安装所需的Python库,包括transformers
、faiss
和PyMuPDF
。
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])
3. 构建Faiss索引
将向量添加到Faiss索引中。
# 创建Faiss索引
d = vectors.shape[1]
index = faiss.IndexFlatL2(d)
# 添加向量到索引
index.add(vectors)
print("Total number of vectors in index:", index.ntotal)
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)
# 提取相似段落
similar_paragraphs = [paragraphs[i] for i in indices[0]]
similar_paragraphs
5. 使用生成模型进行RAG
使用检索到的文档内容作为上下文,结合生成模型(如GPT-3或T5)回答用户查询。
1、使用modelscope下载 ChatGLM-6B 模型
# 使用modelscope下载模型
from modelscope import snapshot_download
model_dir = snapshot_download('ZhipuAI/ChatGLM-6B')
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}"
创建pipeline并生成回答
pipe = pipeline(task=Tasks.chat, model=model, max_length=512)
inputs = {'text': input_text, 'history': []}
result = pipe(inputs)
print("Generated Answer:", result['response'])