增加FAISS向量数据库检索存储等功能。
This commit is contained in:
zhanghuan
@@ -0,0 +1,516 @@
|
||||
import faiss
|
||||
import numpy as np
|
||||
import pickle
|
||||
import os
|
||||
from typing import List, Tuple, Optional, Union, Dict, Any
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import json
|
||||
from datetime import datetime
|
||||
|
||||
class FAISSManager:
|
||||
"""
|
||||
FAISS向量管理器
|
||||
支持向量的增删改查、余弦相似度搜索、批量操作等功能
|
||||
"""
|
||||
|
||||
def __init__(self, dimension: int, index_type: str = "IndexFlatIP",
|
||||
normalize_vectors: bool = True, metric_type: str = "cosine"):
|
||||
"""
|
||||
初始化FAISS管理器
|
||||
|
||||
Args:
|
||||
dimension: 向量维度
|
||||
index_type: 索引类型 ("IndexFlatIP", "IndexFlatL2", "IndexIVFFlat", "IndexHNSW")
|
||||
normalize_vectors: 是否标准化向量（用于余弦相似度）
|
||||
metric_type: 距离度量类型 ("cosine", "l2", "ip")
|
||||
"""
|
||||
self.dimension = dimension
|
||||
self.index_type = index_type
|
||||
self.normalize_vectors = normalize_vectors
|
||||
self.metric_type = metric_type
|
||||
|
||||
# 初始化索引
|
||||
self.index = self._create_index()
|
||||
|
||||
# 存储向量ID到实际ID的映射
|
||||
self.id_mapping = {} # faiss_id -> actual_id
|
||||
self.reverse_id_mapping = {} # actual_id -> faiss_id
|
||||
self.next_faiss_id = 0
|
||||
|
||||
# 存储向量元数据
|
||||
self.metadata = {} # actual_id -> metadata
|
||||
|
||||
# 日志设置
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
self.logger = logging.getLogger(__name__)
|
||||
|
||||
def _create_index(self) -> 'faiss.Index':
|
||||
"""创建FAISS索引"""
|
||||
if self.index_type == "IndexFlatIP":
|
||||
# 内积索引（适合余弦相似度，需要标准化向量）
|
||||
# 初始化的时候需要告诉维度
|
||||
index = faiss.IndexFlatIP(self.dimension)
|
||||
elif self.index_type == "IndexFlatL2":
|
||||
# L2距离索引
|
||||
index = faiss.IndexFlatL2(self.dimension)
|
||||
elif self.index_type == "IndexIVFFlat":
|
||||
# IVF索引（适合大规模数据）
|
||||
quantizer = faiss.IndexFlatIP(self.dimension) if self.metric_type == "cosine" else faiss.IndexFlatL2(self.dimension)
|
||||
nlist = 100 # 聚类中心数量
|
||||
index = faiss.IndexIVFFlat(quantizer, self.dimension, nlist)
|
||||
elif self.index_type == "IndexHNSW":
|
||||
# HNSW索引（高性能近似搜索）
|
||||
index = faiss.IndexHNSWFlat(self.dimension, 32)
|
||||
else:
|
||||
raise ValueError(f"不支持的索引类型: {self.index_type}")
|
||||
|
||||
return index
|
||||
|
||||
def _normalize_vector(self, vector: np.ndarray) -> np.ndarray:
|
||||
"""标准化向量（用于余弦相似度）"""
|
||||
if self.normalize_vectors:
|
||||
norm = np.linalg.norm(vector, axis=-1, keepdims=True)
|
||||
# 避免除零
|
||||
norm = np.where(norm == 0, 1, norm)
|
||||
return vector / norm
|
||||
return vector
|
||||
|
||||
def add_vector(self, vector: Union[np.ndarray, List[float]],
|
||||
vector_id: str, metadata: Optional[Dict[str, Any]] = None) -> bool:
|
||||
"""
|
||||
添加单个向量
|
||||
|
||||
Args:
|
||||
vector: 向量数据
|
||||
vector_id: 向量唯一标识
|
||||
metadata: 向量元数据
|
||||
|
||||
Returns:
|
||||
bool: 是否添加成功
|
||||
"""
|
||||
try:
|
||||
# 检查向量是否已存在
|
||||
if vector_id in self.reverse_id_mapping:
|
||||
self.logger.warning(f"向量ID {vector_id} 已存在，将更新该向量")
|
||||
return self.update_vector(vector, vector_id, metadata)
|
||||
|
||||
# 转换为numpy数组并标准化
|
||||
vector = np.array(vector, dtype=np.float32).reshape(1, -1)
|
||||
if vector.shape[1] != self.dimension:
|
||||
raise ValueError(f"向量维度不匹配: 期望 {self.dimension}, 实际 {vector.shape[1]}")
|
||||
|
||||
vector = self._normalize_vector(vector)
|
||||
|
||||
# 添加到索引
|
||||
self.index.add(vector)
|
||||
|
||||
# 更新映射关系
|
||||
faiss_id = self.next_faiss_id
|
||||
self.id_mapping[faiss_id] = vector_id
|
||||
self.reverse_id_mapping[vector_id] = faiss_id
|
||||
self.next_faiss_id += 1
|
||||
|
||||
# 存储元数据
|
||||
if metadata:
|
||||
self.metadata[vector_id] = metadata
|
||||
|
||||
self.logger.info(f"成功添加向量: {vector_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"添加向量失败: {e}")
|
||||
return False
|
||||
|
||||
def add_vectors_batch(self, vectors: Union[np.ndarray, List[List[float]]],
|
||||
vector_ids: List[str],
|
||||
metadata_list: Optional[List[Dict[str, Any]]] = None) -> List[bool]:
|
||||
"""
|
||||
批量添加向量
|
||||
|
||||
Args:
|
||||
vectors: 向量数据矩阵
|
||||
vector_ids: 向量ID列表
|
||||
metadata_list: 元数据列表
|
||||
|
||||
Returns:
|
||||
List[bool]: 每个向量的添加结果
|
||||
"""
|
||||
results = []
|
||||
|
||||
# 转换为numpy数组
|
||||
vectors = np.array(vectors, dtype=np.float32)
|
||||
if vectors.ndim == 1:
|
||||
vectors = vectors.reshape(1, -1)
|
||||
|
||||
if vectors.shape[1] != self.dimension:
|
||||
raise ValueError(f"向量维度不匹配: 期望 {self.dimension}, 实际 {vectors.shape[1]}")
|
||||
|
||||
if len(vector_ids) != vectors.shape[0]:
|
||||
raise ValueError("向量数量与ID数量不匹配")
|
||||
|
||||
# 标准化向量
|
||||
vectors = self._normalize_vector(vectors)
|
||||
|
||||
# 检查重复ID
|
||||
new_vectors = []
|
||||
new_ids = []
|
||||
new_metadata = []
|
||||
|
||||
for i, vector_id in enumerate(vector_ids):
|
||||
if vector_id not in self.reverse_id_mapping:
|
||||
new_vectors.append(vectors[i])
|
||||
new_ids.append(vector_id)
|
||||
if metadata_list:
|
||||
new_metadata.append(metadata_list[i])
|
||||
results.append(True)
|
||||
else:
|
||||
self.logger.warning(f"向量ID {vector_id} 已存在，跳过")
|
||||
results.append(False)
|
||||
|
||||
if new_vectors:
|
||||
# 批量添加到索引
|
||||
new_vectors = np.array(new_vectors)
|
||||
self.index.add(new_vectors)
|
||||
|
||||
# 批量更新映射关系
|
||||
for i, vector_id in enumerate(new_ids):
|
||||
faiss_id = self.next_faiss_id + i
|
||||
self.id_mapping[faiss_id] = vector_id
|
||||
self.reverse_id_mapping[vector_id] = faiss_id
|
||||
|
||||
# 存储元数据
|
||||
if new_metadata and i < len(new_metadata):
|
||||
self.metadata[vector_id] = new_metadata[i]
|
||||
|
||||
self.next_faiss_id += len(new_ids)
|
||||
self.logger.info(f"批量添加了 {len(new_ids)} 个向量")
|
||||
|
||||
return results
|
||||
|
||||
def search_similar(self, query_vector: Union[np.ndarray, List[float]],
|
||||
k: int = 10, threshold: Optional[float] = None) -> List[Tuple[str, float, Dict[str, Any]]]:
|
||||
"""
|
||||
搜索相似向量
|
||||
|
||||
Args:
|
||||
query_vector: 查询向量
|
||||
k: 返回最相似的k个结果
|
||||
threshold: 相似度阈值
|
||||
|
||||
Returns:
|
||||
List[Tuple[str, float, Dict]]: (vector_id, similarity_score, metadata)
|
||||
"""
|
||||
if self.index.ntotal == 0:
|
||||
return []
|
||||
|
||||
# 转换并标准化查询向量
|
||||
query_vector = np.array(query_vector, dtype=np.float32).reshape(1, -1)
|
||||
if query_vector.shape[1] != self.dimension:
|
||||
raise ValueError(f"查询向量维度不匹配: 期望 {self.dimension}, 实际 {query_vector.shape[1]}")
|
||||
|
||||
query_vector = self._normalize_vector(query_vector)
|
||||
|
||||
# 搜索
|
||||
k = min(k, self.index.ntotal)
|
||||
distances, indices = self.index.search(query_vector, k)
|
||||
|
||||
results = []
|
||||
for i in range(k):
|
||||
faiss_id = indices[0][i]
|
||||
distance = distances[0][i]
|
||||
|
||||
if faiss_id == -1: # 无效结果
|
||||
continue
|
||||
|
||||
# 转换距离为相似度分数
|
||||
if self.metric_type == "cosine" or self.index_type == "IndexFlatIP":
|
||||
similarity = float(distance) # 内积已经是相似度
|
||||
else:
|
||||
similarity = 1.0 / (1.0 + float(distance)) # L2距离转相似度
|
||||
|
||||
# 应用阈值过滤
|
||||
if threshold is not None and similarity < threshold:
|
||||
continue
|
||||
|
||||
vector_id = self.id_mapping.get(faiss_id)
|
||||
if vector_id:
|
||||
metadata = self.metadata.get(vector_id, {})
|
||||
results.append((vector_id, similarity, metadata))
|
||||
|
||||
return results
|
||||
|
||||
def get_vector_by_id(self, vector_id: str) -> Optional[Tuple[np.ndarray, Dict[str, Any]]]:
|
||||
"""
|
||||
根据ID获取向量
|
||||
|
||||
Args:
|
||||
vector_id: 向量ID
|
||||
|
||||
Returns:
|
||||
Optional[Tuple[np.ndarray, Dict]]: (vector, metadata) 或 None
|
||||
"""
|
||||
faiss_id = self.reverse_id_mapping.get(vector_id)
|
||||
if faiss_id is None:
|
||||
return None
|
||||
|
||||
# FAISS不直接支持根据ID获取向量，这里返回元数据
|
||||
# 如果需要向量数据，建议单独存储
|
||||
metadata = self.metadata.get(vector_id, {})
|
||||
return None, metadata
|
||||
|
||||
def update_vector(self, vector: Union[np.ndarray, List[float]],
|
||||
vector_id: str, metadata: Optional[Dict[str, Any]] = None) -> bool:
|
||||
"""
|
||||
更新向量（先删除再添加）
|
||||
|
||||
Args:
|
||||
vector: 新向量数据
|
||||
vector_id: 向量ID
|
||||
metadata: 新元数据
|
||||
|
||||
Returns:
|
||||
bool: 是否更新成功
|
||||
"""
|
||||
# FAISS不支持直接更新，需要重建索引
|
||||
if vector_id not in self.reverse_id_mapping:
|
||||
self.logger.warning(f"向量ID {vector_id} 不存在")
|
||||
return False
|
||||
|
||||
# 删除旧向量
|
||||
self.delete_vector(vector_id)
|
||||
|
||||
# 添加新向量
|
||||
return self.add_vector(vector, vector_id, metadata)
|
||||
|
||||
def delete_vector(self, vector_id: str) -> bool:
|
||||
"""
|
||||
删除向量
|
||||
|
||||
Args:
|
||||
vector_id: 向量ID
|
||||
|
||||
Returns:
|
||||
bool: 是否删除成功
|
||||
"""
|
||||
if vector_id not in self.reverse_id_mapping:
|
||||
self.logger.warning(f"向量ID {vector_id} 不存在")
|
||||
return False
|
||||
|
||||
# FAISS不支持直接删除，标记为删除
|
||||
faiss_id = self.reverse_id_mapping[vector_id]
|
||||
|
||||
# 从映射中移除
|
||||
del self.id_mapping[faiss_id]
|
||||
del self.reverse_id_mapping[vector_id]
|
||||
|
||||
# 删除元数据
|
||||
if vector_id in self.metadata:
|
||||
del self.metadata[vector_id]
|
||||
|
||||
self.logger.info(f"标记删除向量: {vector_id}")
|
||||
return True
|
||||
|
||||
def delete_vectors_batch(self, vector_ids: List[str]) -> List[bool]:
|
||||
"""
|
||||
批量删除向量
|
||||
|
||||
Args:
|
||||
vector_ids: 向量ID列表
|
||||
|
||||
Returns:
|
||||
List[bool]: 每个向量的删除结果
|
||||
"""
|
||||
results = []
|
||||
for vector_id in vector_ids:
|
||||
results.append(self.delete_vector(vector_id))
|
||||
return results
|
||||
|
||||
def rebuild_index(self) -> bool:
|
||||
"""
|
||||
重建索引（清理已删除的向量）
|
||||
"""
|
||||
try:
|
||||
# 获取所有有效向量的信息
|
||||
valid_vectors = []
|
||||
valid_ids = []
|
||||
valid_metadata = {}
|
||||
|
||||
# 这里需要外部提供向量数据，因为FAISS不能直接获取
|
||||
# 建议在实际使用时维护一个向量存储
|
||||
self.logger.warning("重建索引需要外部提供向量数据")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"重建索引失败: {e}")
|
||||
return False
|
||||
|
||||
def get_index_info(self) -> Dict[str, Any]:
|
||||
"""
|
||||
获取索引信息
|
||||
|
||||
Returns:
|
||||
Dict: 索引统计信息
|
||||
"""
|
||||
return {
|
||||
"total_vectors": self.index.ntotal,
|
||||
"dimension": self.dimension,
|
||||
"index_type": self.index_type,
|
||||
"metric_type": self.metric_type,
|
||||
"normalize_vectors": self.normalize_vectors,
|
||||
"active_vectors": len(self.reverse_id_mapping),
|
||||
"metadata_count": len(self.metadata)
|
||||
}
|
||||
|
||||
def save_index(self, save_path: str) -> bool:
|
||||
"""
|
||||
保存索引到文件
|
||||
|
||||
Args:
|
||||
save_path: 保存路径
|
||||
|
||||
Returns:
|
||||
bool: 是否保存成功
|
||||
"""
|
||||
try:
|
||||
save_path = Path(save_path)
|
||||
save_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 保存FAISS索引
|
||||
index_file = save_path / "faiss_index.bin"
|
||||
faiss.write_index(self.index, str(index_file))
|
||||
|
||||
# 保存映射关系和元数据
|
||||
metadata_file = save_path / "metadata.pkl"
|
||||
with open(metadata_file, 'wb') as f:
|
||||
pickle.dump({
|
||||
'id_mapping': self.id_mapping,
|
||||
'reverse_id_mapping': self.reverse_id_mapping,
|
||||
'next_faiss_id': self.next_faiss_id,
|
||||
'metadata': self.metadata,
|
||||
'dimension': self.dimension,
|
||||
'index_type': self.index_type,
|
||||
'normalize_vectors': self.normalize_vectors,
|
||||
'metric_type': self.metric_type
|
||||
}, f)
|
||||
|
||||
# 保存配置信息
|
||||
config_file = save_path / "config.json"
|
||||
with open(config_file, 'w', encoding='utf-8') as f:
|
||||
json.dump({
|
||||
'dimension': self.dimension,
|
||||
'index_type': self.index_type,
|
||||
'normalize_vectors': self.normalize_vectors,
|
||||
'metric_type': self.metric_type,
|
||||
'save_time': datetime.now().isoformat(),
|
||||
'total_vectors': self.index.ntotal
|
||||
}, f, indent=2, ensure_ascii=False)
|
||||
|
||||
self.logger.info(f"索引已保存到: {save_path}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"保存索引失败: {e}")
|
||||
return False
|
||||
|
||||
def load_index(self, load_path: str) -> bool:
|
||||
"""
|
||||
从文件加载索引
|
||||
|
||||
Args:
|
||||
load_path: 加载路径
|
||||
|
||||
Returns:
|
||||
bool: 是否加载成功
|
||||
"""
|
||||
try:
|
||||
load_path = Path(load_path)
|
||||
|
||||
# 检查文件是否存在
|
||||
index_file = load_path / "faiss_index.bin"
|
||||
metadata_file = load_path / "metadata.pkl"
|
||||
|
||||
if not index_file.exists() or not metadata_file.exists():
|
||||
self.logger.error(f"索引文件不存在: {load_path}")
|
||||
return False
|
||||
|
||||
# 加载FAISS索引
|
||||
self.index = faiss.read_index(str(index_file))
|
||||
|
||||
# 加载映射关系和元数据
|
||||
with open(metadata_file, 'rb') as f:
|
||||
data = pickle.load(f)
|
||||
self.id_mapping = data['id_mapping']
|
||||
self.reverse_id_mapping = data['reverse_id_mapping']
|
||||
self.next_faiss_id = data['next_faiss_id']
|
||||
self.metadata = data['metadata']
|
||||
self.dimension = data['dimension']
|
||||
self.index_type = data['index_type']
|
||||
self.normalize_vectors = data['normalize_vectors']
|
||||
self.metric_type = data['metric_type']
|
||||
|
||||
self.logger.info(f"索引已从 {load_path} 加载")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"加载索引失败: {e}")
|
||||
return False
|
||||
|
||||
def clear_index(self) -> bool:
|
||||
"""
|
||||
清空索引
|
||||
|
||||
Returns:
|
||||
bool: 是否清空成功
|
||||
"""
|
||||
try:
|
||||
# 重新创建索引
|
||||
self.index = self._create_index()
|
||||
|
||||
# 清空映射关系
|
||||
self.id_mapping.clear()
|
||||
self.reverse_id_mapping.clear()
|
||||
self.next_faiss_id = 0
|
||||
|
||||
# 清空元数据
|
||||
self.metadata.clear()
|
||||
|
||||
self.logger.info("索引已清空")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"清空索引失败: {e}")
|
||||
return False
|
||||
|
||||
|
||||
# 使用示例
|
||||
if __name__ == "__main__":
|
||||
# 创建FAISS管理器
|
||||
manager = FAISSManager(dimension=512, index_type="IndexFlatIP")
|
||||
|
||||
# 添加向量
|
||||
vector1 = np.random.random(512).astype(np.float32)
|
||||
manager.add_vector(vector1, "vec_1", {"label": "食物1", "category": "川菜"})
|
||||
|
||||
# 批量添加向量
|
||||
vectors = np.random.random((10, 512)).astype(np.float32)
|
||||
ids = [f"vec_{i}" for i in range(2, 12)]
|
||||
metadata_list = [{"label": f"食物{i}", "category": "川菜"} for i in range(2, 12)]
|
||||
manager.add_vectors_batch(vectors, ids, metadata_list)
|
||||
|
||||
# 搜索相似向量
|
||||
query = np.random.random(512).astype(np.float32)
|
||||
results = manager.search_similar(query, k=5)
|
||||
print("搜索结果:", results)
|
||||
|
||||
# 获取索引信息
|
||||
# 获得统计信息
|
||||
info = manager.get_index_info()
|
||||
print("索引信息:", info)
|
||||
|
||||
# 保存索引
|
||||
manager.save_index("./faiss_index")
|
||||
|
||||
# 加载索引
|
||||
new_manager = FAISSManager(dimension=512)
|
||||
new_manager.load_index("./faiss_index")
|
||||
Reference in New Issue
Block a user