Spaces:
No application file
No application file
File size: 1,324 Bytes
d462a5b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | # vector_store.py
import faiss
import numpy as np
import pickle
from typing import List
class SimpleVectorStore:
def __init__(self, dim, index_path=None):
self.dim = dim
self.index = faiss.IndexFlatIP(dim) # cosine if vectors normalized
self.metadata = []
self.index_path = index_path
def add(self, vec: np.ndarray, meta: dict):
if vec.ndim == 1:
vec = vec.reshape(1, -1)
self.index.add(vec.astype("float32"))
self.metadata.append(meta)
def search(self, query_vec: np.ndarray, k=5):
if query_vec.ndim == 1:
query_vec = query_vec.reshape(1, -1)
D, I = self.index.search(query_vec.astype("float32"), k)
results = []
for dist, idx in zip(D[0], I[0]):
if idx == -1 or idx >= len(self.metadata):
continue
results.append((float(dist), self.metadata[idx]))
return results
def save(self, path_prefix):
faiss.write_index(self.index, f"{path_prefix}.index")
with open(f"{path_prefix}.meta", "wb") as f:
pickle.dump(self.metadata, f)
def load(self, path_prefix):
self.index = faiss.read_index(f"{path_prefix}.index")
with open(f"{path_prefix}.meta", "rb") as f:
self.metadata = pickle.load(f)
|