Spaces:
Paused
Paused
| """Embedding + BM25 cache helpers, as used by Notebooks 2-6. | |
| Cache key = md5(domain | config_name | embedding_model_name | chunk_size | | |
| chunk_overlap)[:12] -- this is the same scheme the notebooks use, so a | |
| combo built by a notebook (and synced into `data/cache/indices/<key>/`) is | |
| picked up here unchanged, and a combo built here is picked up by the | |
| notebooks unchanged. | |
| """ | |
| import hashlib | |
| import json | |
| import pickle | |
| from pathlib import Path | |
| import faiss | |
| import numpy as np | |
| import pandas as pd | |
| from rank_bm25 import BM25Okapi | |
| from sentence_transformers import SentenceTransformer | |
| from src import settings | |
| def make_cache_key(*parts) -> str: | |
| raw = "|".join(str(p) for p in parts) | |
| return hashlib.md5(raw.encode("utf-8")).hexdigest()[:12] | |
| def get_or_build_dense_index( | |
| chunks_df: pd.DataFrame, | |
| model: SentenceTransformer, | |
| model_name: str, | |
| domain: str, | |
| config_name: str, | |
| chunk_size: int, | |
| chunk_overlap: int, | |
| cache_root: Path = None, | |
| docs_df: pd.DataFrame = None, | |
| ): | |
| cache_root = cache_root or settings.INDICES_DIR | |
| cache_key = make_cache_key(domain, config_name, model_name, chunk_size, chunk_overlap) | |
| cache_dir = cache_root / cache_key | |
| emb_path = cache_dir / "embeddings.npy" | |
| index_path = cache_dir / "faiss.index" | |
| if emb_path.exists() and index_path.exists(): | |
| embeddings = np.load(emb_path) | |
| index = faiss.read_index(str(index_path)) | |
| if docs_df is not None and not (cache_dir / "docs.parquet").exists(): | |
| docs_df.to_parquet(cache_dir / "docs.parquet") | |
| return embeddings, index, cache_dir | |
| cache_dir.mkdir(parents=True, exist_ok=True) | |
| texts = chunks_df["text"].tolist() | |
| embeddings = model.encode( | |
| texts, batch_size=64, show_progress_bar=True, convert_to_numpy=True, normalize_embeddings=True | |
| ).astype("float32") | |
| index = faiss.IndexFlatIP(embeddings.shape[1]) | |
| index.add(embeddings) | |
| np.save(emb_path, embeddings) | |
| faiss.write_index(index, str(index_path)) | |
| chunks_df.to_parquet(cache_dir / "chunks.parquet") | |
| if docs_df is not None: | |
| docs_df.to_parquet(cache_dir / "docs.parquet") | |
| with open(cache_dir / "meta.json", "w") as f: | |
| json.dump({ | |
| "model_name": model_name, "chunk_size": chunk_size, "chunk_overlap": chunk_overlap, | |
| "domain": domain, "config": config_name, "n_chunks": len(chunks_df), | |
| }, f, indent=2) | |
| return embeddings, index, cache_dir | |
| def get_or_build_bm25(chunks_df: pd.DataFrame, cache_dir: Path) -> BM25Okapi: | |
| bm25_path = cache_dir / "bm25.pkl" | |
| if bm25_path.exists(): | |
| with open(bm25_path, "rb") as f: | |
| return pickle.load(f) | |
| tokenized = [text.lower().split() for text in chunks_df["text"]] | |
| bm25 = BM25Okapi(tokenized) | |
| with open(bm25_path, "wb") as f: | |
| pickle.dump(bm25, f) | |
| return bm25 | |
| def load_dense_index(cache_dir: Path): | |
| """Load a precomputed FAISS index + chunks.parquet (no rebuilding).""" | |
| index = faiss.read_index(str(cache_dir / "faiss.index")) | |
| chunks_df = pd.read_parquet(cache_dir / "chunks.parquet") | |
| return index, chunks_df | |
| def load_bm25(cache_dir: Path) -> BM25Okapi: | |
| with open(cache_dir / "bm25.pkl", "rb") as f: | |
| return pickle.load(f) | |
| def load_docs(cache_dir: Path) -> pd.DataFrame: | |
| """Full document texts (doc_id, text) -- used to build sentence-window context.""" | |
| return pd.read_parquet(cache_dir / "docs.parquet") | |