Spaces:
Running
Running
| import faiss | |
| import numpy as np | |
| from sentence_transformers import SentenceTransformer | |
| class VectorStore: | |
| def __init__(self, model_name: str): | |
| # Force CPU & low memory usage | |
| self.model = SentenceTransformer( | |
| model_name, | |
| device="cpu" | |
| ) | |
| self.index = None | |
| self.text_chunks = [] | |
| def build(self, chunks): | |
| self.text_chunks = chunks | |
| # Convert directly to numpy (saves memory) | |
| embeddings = self.model.encode( | |
| chunks, | |
| convert_to_numpy=True, | |
| show_progress_bar=False | |
| ) | |
| dim = embeddings.shape[1] | |
| self.index = faiss.IndexFlatL2(dim) | |
| self.index.add(embeddings) | |
| def search(self, query, k=5): | |
| query_embedding = self.model.encode( | |
| [query], | |
| convert_to_numpy=True | |
| ) | |
| _, idx = self.index.search(query_embedding, k) | |
| return [self.text_chunks[i] for i in idx[0]] | |