Spaces:
Running on Zero
Running on Zero
| import faiss | |
| import numpy as np | |
| import os | |
| class FaissDB: | |
| def __init__(self, dimension=512): | |
| self.dimension = dimension | |
| self.index = faiss.IndexFlatIP(dimension) | |
| def create(self): | |
| self.index = faiss.IndexFlatIP(self.dimension) | |
| def load(self, path): | |
| if os.path.exists(path): | |
| self.index = faiss.read_index(path) | |
| return True | |
| return False | |
| def save(self, path): | |
| faiss.write_index(self.index, path) | |
| def add(self, embeddings): | |
| embeddings = np.asarray(embeddings, dtype=np.float32) | |
| self.index.add(embeddings) | |
| def search(self, query_embedding, k=10): | |
| embeddings = np.asarray( | |
| query_embedding, dtype=np.float32).reshape(1, -1) | |
| scores, indices = self.index.search(embeddings, k) | |
| return scores[0], indices[0] | |
| def size(self): | |
| return self.index.ntotal | |