Spaces:
Sleeping
Sleeping
File size: 4,467 Bytes
a4538e5 80b3e44 a4538e5 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | 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, valid_chunks = emb_pipe.embed_chunks(chunks)
if len(valid_chunks) == 0:
print("[WARNING] No valid chunks were embedded. Vector store will not be updated.")
return
metadatas = [{"texts": chunk.page_content} for chunk in valid_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)) |