Spaces:
Sleeping
Sleeping
| import os | |
| import faiss | |
| import numpy as np | |
| from typing import List, Tuple | |
| def build_or_load_index(vectors: np.ndarray, index_path: str) -> faiss.Index: | |
| """ | |
| 构建或加载FAISS索引。 | |
| """ | |
| if os.path.exists(index_path): | |
| print(f"Loading existing FAISS index from {index_path}") | |
| return faiss.read_index(index_path) | |
| else: | |
| print("Building new FAISS index.") | |
| # === 修改开始 === | |
| # 原来的代码在这里判断 len(vectors) == 0 就报错，导致无法初始化空库 | |
| # 我们现在只判断 vectors 是否为 None | |
| if vectors is None: | |
| raise ValueError("Vectors cannot be None.") | |
| # 获取维度 (例如 512) | |
| # 即使 vectors 是 (0, 512) 的空数组，shape[1] 依然是 512，这是合法的 | |
| vector_dim = vectors.shape[1] | |
| # 创建索引 | |
| index = faiss.IndexFlatL2(vector_dim) | |
| # 只有当向量里真的有数据时，才执行 add | |
| if len(vectors) > 0: | |
| index.add(vectors) | |
| # === 修改结束 === | |
| return index | |
| def save_index(index: faiss.Index, index_path: str) -> None: | |
| """ | |
| 将内存中的FAISS索引保存到硬盘。 | |
| - 参数 index: 需要保存的FAISS索引对象。 | |
| - 参数 index_path: 保存路径。 | |
| """ | |
| print(f"Saving FAISS index to {index_path}") | |
| faiss.write_index(index, index_path) | |
| def search_in_index( | |
| index: faiss.Index, query_vector: np.ndarray, n: int | |
| ) -> Tuple[np.ndarray, np.ndarray]: | |
| """ | |
| 在FAISS索引中执行搜索。 | |
| - 参数 index: 已加载的FAISS索引对象。 | |
| - 参数 query_vector: NumPy数组形式的查询向量 (形状应为 1 x dim)。 | |
| - 参数 n: 希望返回的结果数量。 | |
| - 返回: (距离数组, 索引ID数组) | |
| """ | |
| if not query_vector.flags["C_CONTIGUOUS"]: | |
| query_vector = np.ascontiguousarray(query_vector) | |
| return index.search(query_vector, n) | |
| def add_to_index(index: faiss.Index, vector: np.ndarray) -> None: | |
| """ | |
| 向索引中增量添加一个向量。 | |
| - 参数 index: 需要更新的FAISS索引对象。 | |
| - 参数 vector: NumPy数组形式的新向量 (形状应为 1 x dim)。 | |
| """ | |
| if not vector.flags["C_CONTIGUOUS"]: | |
| vector = np.ascontiguousarray(vector) | |
| index.add(vector) | |