Spaces:
Running
Running
File size: 1,760 Bytes
112d847 |
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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
from abc import ABC, abstractmethod
import faiss
import numpy as np
import os
import logging
import sys
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', stream=sys.stdout)
logger = logging.getLogger(__name__)
class BaseVectorStore(ABC):
@abstractmethod
def add_documents(self, documents, embeddings):
pass
@abstractmethod
def search(self, query_vector, num_results=5):
pass
class FaissVectorStore(BaseVectorStore):
def __init__(self, dimension):
self.dimension = dimension
self.index = faiss.IndexFlatL2(dimension)
self.documents = []
self.index_path = "data/faiss_index"
os.makedirs("data", exist_ok=True)
self.load_index()
def load_index(self):
if os.path.exists(self.index_path):
try:
self.index = faiss.read_index(self.index_path)
except Exception as e:
logger.error(f"Error loading FAISS index: {e}")
def save_index(self):
try:
faiss.write_index(self.index, self.index_path)
except Exception as e:
logger.error(f"Error saving FAISS index: {e}")
def add_documents(self, documents, embeddings):
self.index.add(np.array(embeddings))
self.documents.extend(documents)
self.save_index()
def search(self, query_vector, num_results=5):
if len(self.documents) == 0:
return []
D, I = self.index.search(np.array([query_vector]), num_results)
return [self.documents[i] for i in I[0]]
def get_vector_store(config):
"""Factory function to get the appropriate vector store"""
return FaissVectorStore # Always return FAISS vector store |