Spaces:
Running
Running
| """FAISS vector indexer (IndexFlatIP + L2 normalization = cosine similarity). | |
| Stores the chunks and their metadata, aligned by position. Persistence via | |
| save/load (.faiss / .chunks.json / .meta.json). | |
| """ | |
| import json | |
| from pathlib import Path | |
| import faiss | |
| import numpy as np | |
| class FaissIndexer: | |
| """FAISS index shared by the stacks: exact cosine search over the chunks.""" | |
| def __init__(self, dimension: int = 384): | |
| self.dimension = dimension | |
| self.index = faiss.IndexFlatIP(dimension) | |
| self.chunks: list[str] = [] | |
| self.metadata: list[dict] = [] | |
| def add(self, embeddings: np.ndarray, chunks: list[str], metadata: list[dict]): | |
| """Add chunks and their embeddings (L2-normalized before insertion).""" | |
| if len(embeddings) != len(chunks) or len(embeddings) != len(metadata): | |
| raise ValueError( | |
| f"Length mismatch: {len(embeddings)} embeddings, " | |
| f"{len(chunks)} chunks, {len(metadata)} metadata entries." | |
| ) | |
| if embeddings.shape[1] != self.dimension: | |
| raise ValueError( | |
| f"Embedding dimension {embeddings.shape[1]} does not match " | |
| f"index dimension {self.dimension}." | |
| ) | |
| # L2 normalization: cosine similarity is then obtained via the dot product. | |
| embeddings = embeddings.astype(np.float32).copy() | |
| faiss.normalize_L2(embeddings) | |
| self.index.add(embeddings) | |
| self.chunks.extend(chunks) | |
| self.metadata.extend(metadata) | |
| def size(self) -> int: | |
| """Number of vectors currently in the index.""" | |
| return self.index.ntotal | |
| def save(self, path: str): | |
| """Persist the index to disk: {path}.faiss, .chunks.json, .meta.json.""" | |
| base = Path(path) | |
| base.parent.mkdir(parents=True, exist_ok=True) | |
| faiss.write_index(self.index, str(base.with_suffix(".faiss"))) | |
| with open(base.with_suffix(".chunks.json"), "w", encoding="utf-8") as f: | |
| json.dump(self.chunks, f, ensure_ascii=False) | |
| with open(base.with_suffix(".meta.json"), "w", encoding="utf-8") as f: | |
| json.dump(self.metadata, f, ensure_ascii=False) | |
| def load(self, path: str): | |
| """Reload a saved index (same prefix as save()).""" | |
| base = Path(path) | |
| index_path = base.with_suffix(".faiss") | |
| chunks_path = base.with_suffix(".chunks.json") | |
| meta_path = base.with_suffix(".meta.json") | |
| for p in (index_path, chunks_path, meta_path): | |
| if not p.exists(): | |
| raise FileNotFoundError(f"Required file not found: {p}") | |
| self.index = faiss.read_index(str(index_path)) | |
| self.dimension = self.index.d | |
| with open(chunks_path, encoding="utf-8") as f: | |
| self.chunks = json.load(f) | |
| with open(meta_path, encoding="utf-8") as f: | |
| self.metadata = json.load(f) | |