Spaces:
Sleeping
Sleeping
| import os | |
| import faiss | |
| import numpy as np | |
| import pickle | |
| from typing import List, Any | |
| from sentence_transformers import SentenceTransformer | |
| from src.embedding import EmbeddingPipeline | |
| class FaissVectorStore: | |
| def __init__( | |
| self, | |
| persist_dir: str = "faiss_store", | |
| embedding_model: str = "all-MiniLM-L6-v2", | |
| chunk_size: int = 1000, | |
| chunk_overlap: int = 200, | |
| ): | |
| self.persist_dir = persist_dir | |
| os.makedirs(self.persist_dir, exist_ok=True) | |
| self.index = None | |
| self.metadata = [] | |
| self.embedding_model = embedding_model | |
| self.model = SentenceTransformer(embedding_model) | |
| self.chunk_size = chunk_size | |
| self.chunk_overlap = chunk_overlap | |
| print(f"[INFO] Loaded embedding model: {embedding_model}") | |
| def build_from_documents(self, documents: List[Any]): | |
| print(f"[INFO] Building vector store from {len(documents)} raw document(s)...") | |
| emb_pipe = EmbeddingPipeline( | |
| model_name=self.embedding_model, | |
| chunk_size=self.chunk_size, | |
| chunk_overlap=self.chunk_overlap, | |
| ) | |
| chunks = emb_pipe.chunk_documents(documents) | |
| embeddings = emb_pipe.embed_chunks(chunks) | |
| metadatas = [{"texts": chunk.page_content} for chunk in chunks] | |
| self.add_embeddings(np.array(embeddings).astype("float32"), metadatas) | |
| self.save() | |
| print(f"[INFO] Vector Store built and saved to {self.persist_dir}") | |
| def add_embeddings(self, embeddings: np.ndarray, metadatas: List[Any] = None): | |
| # Handle empty embeddings case | |
| if embeddings.size == 0: | |
| print("[WARNING] No embeddings to add. Vector store remains empty.") | |
| return | |
| dim = embeddings.shape[1] | |
| if self.index is None: | |
| self.index = faiss.IndexFlatL2(dim) | |
| self.index.add(embeddings) | |
| if metadatas: | |
| self.metadata.extend(metadatas) | |
| print(f"[INFO] Added {embeddings.shape[0]} vectors to Faiss Index.") | |
| def save(self): | |
| if self.index is None: | |
| print("[WARNING] Cannot save: index is empty. Skipping save operation.") | |
| return | |
| faiss_path = os.path.join(self.persist_dir, "faiss.index") | |
| meta_path = os.path.join(self.persist_dir, "metadata.pkl") | |
| faiss.write_index(self.index, faiss_path) | |
| with open(meta_path, "wb") as f: | |
| pickle.dump(self.metadata, f) | |
| print(f"[INFO] Saved Faiss index and metadata to {self.persist_dir}") | |
| def load(self): | |
| faiss_path = os.path.join(self.persist_dir, "faiss.index") | |
| meta_path = os.path.join(self.persist_dir, "metadata.pkl") | |
| if not (os.path.exists(faiss_path) and os.path.exists(meta_path)): | |
| raise FileNotFoundError(f"Missing index/metadata in {self.persist_dir}. Build the store first.") | |
| self.index = faiss.read_index(faiss_path) | |
| with open(meta_path, "rb") as f: | |
| self.metadata = pickle.load(f) | |
| print(f"[INFO] Loaded Faiss Index and metadata from {self.persist_dir}") | |
| def search(self, query_embeddings: np.ndarray, top_k: int = 5): | |
| if self.index is None: | |
| print("[WARNING] Vector store is empty. No results to return.") | |
| return [] | |
| D, I = self.index.search(query_embeddings, top_k) | |
| results = [] | |
| for idx, dist in zip(I[0], D[0]): | |
| meta = self.metadata[idx] if idx < len(self.metadata) else None | |
| results.append({"index": int(idx), "distance": float(dist), "metadata": meta}) | |
| return results | |
| def query(self, query_text: str, top_k: int = 5): | |
| if self.index is None: | |
| print("[WARNING] Vector store is empty. No results to return.") | |
| return [] | |
| print(f"[INFO] Querying vector store for: '{query_text}'") | |
| query_emb = self.model.encode([query_text]).astype("float32") | |
| return self.search(query_emb, top_k=top_k) | |
| if __name__ == "__main__": | |
| from src.data_loader import load_all_documents | |
| docs = load_all_documents("data") | |
| store = FaissVectorStore("faiss_store") | |
| store.build_from_documents(docs) | |
| store.load() | |
| print(store.query("What is Database Management System?", top_k=3)) |