Spaces:
Paused
Paused
File size: 3,452 Bytes
b96b498 | 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 | """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")
|