29 lines
939 B
Python
29 lines
939 B
Python
import faiss
|
|
import numpy as np
|
|
import threading
|
|
|
|
class FaissIndexWrapper:
|
|
def __init__(self, index_path: str):
|
|
self.index_lock = threading.Lock()
|
|
self.index = None
|
|
self.index_path = index_path
|
|
self.load_index(index_path)
|
|
|
|
def load_index(self, path: str = None):
|
|
with self.index_lock:
|
|
path = path or self.index_path
|
|
self.index = faiss.read_index(path)
|
|
print(f"[FAISS] Index loaded from {path}")
|
|
|
|
def update_index(self, path: str = None):
|
|
"""热更新替换当前索引"""
|
|
with self.index_lock:
|
|
path = path or self.index_path
|
|
self.index = faiss.read_index(path)
|
|
print(f"[FAISS] Index hot-swapped from {path}")
|
|
|
|
def search(self, vector: np.ndarray, top_k: int = 5):
|
|
with self.index_lock:
|
|
D, I = self.index.search(vector.astype(np.float32), top_k)
|
|
return D[0], I[0]
|