| """Hybrid BM25 and FAISS retrieval over the uploaded profile.""" |
|
|
| import bm25s |
| import faiss |
| import numpy as np |
| import Stemmer |
|
|
| from dataloader import SEPARATOR |
|
|
| WORDS_PER_CHUNK = 100 |
| CHUNK_OVERLAP_WORDS = 15 |
| RRF_K = 60 |
|
|
|
|
| def chunk_documents( |
| profile_text, |
| words_per_chunk=WORDS_PER_CHUNK, |
| overlap=CHUNK_OVERLAP_WORDS, |
| ): |
| """Split each source document into word windows with exact overlap.""" |
| if words_per_chunk <= 0 or overlap < 0 or overlap >= words_per_chunk: |
| raise ValueError("Require words_per_chunk > overlap >= 0") |
|
|
| chunks = [] |
| step = words_per_chunk - overlap |
|
|
| for document in profile_text.split(SEPARATOR): |
| document = document.strip() |
| if not document: |
| continue |
|
|
| lines = document.splitlines() |
| has_source = lines[0].startswith("# Source:") |
| source = lines[0].removeprefix("# Source:").strip() if has_source else "" |
| body = "\n".join(lines[1:]).strip() if has_source else document |
| words = body.split() |
|
|
| for start in range(0, len(words), step): |
| chunk = " ".join(words[start : start + words_per_chunk]) |
| if chunk: |
| chunks.append(_label(source, chunk)) |
| if start + words_per_chunk >= len(words): |
| break |
|
|
| return chunks |
|
|
|
|
| def _label(source, text): |
| return f"[{source}]\n{text}" if source else text |
|
|
|
|
| def _reciprocal_rank_fusion(*rankings, k=RRF_K): |
| scores = {} |
| for ranking in rankings: |
| for rank, chunk_id in enumerate(ranking, start=1): |
| if chunk_id >= 0: |
| scores[chunk_id] = scores.get(chunk_id, 0.0) + 1.0 / (k + rank) |
| return sorted(scores, key=scores.get, reverse=True) |
|
|
|
|
| class ProfileIndex: |
| def __init__(self, chunks, embed_fn): |
| self.chunks = list(chunks) |
| self.embed_fn = embed_fn |
| self.stemmer = Stemmer.Stemmer("english") |
|
|
| if self.chunks: |
| tokens = bm25s.tokenize( |
| self.chunks, |
| stopwords="en", |
| stemmer=self.stemmer, |
| show_progress=False, |
| ) |
| self.bm25 = bm25s.BM25() |
| self.bm25.index(tokens, show_progress=False) |
|
|
| vectors = np.ascontiguousarray(embed_fn(self.chunks), dtype=np.float32) |
| faiss.normalize_L2(vectors) |
| self.faiss_index = faiss.IndexFlatIP(vectors.shape[1]) |
| self.faiss_index.add(vectors) |
| else: |
| self.bm25 = None |
| self.faiss_index = None |
|
|
| def search(self, query, k=5): |
| """Return chunks ranked by reciprocal-rank fusion of BM25 and cosine.""" |
| if not self.chunks or not query.strip() or k <= 0: |
| return [] |
|
|
| candidate_count = min(k, len(self.chunks)) |
| query_tokens = bm25s.tokenize( |
| [query], |
| stopwords="en", |
| stemmer=self.stemmer, |
| show_progress=False, |
| ) |
| bm25_ids, _ = self.bm25.retrieve( |
| query_tokens, |
| k=candidate_count, |
| show_progress=False, |
| ) |
| dense_query = np.ascontiguousarray(self.embed_fn([query]), dtype=np.float32) |
| faiss.normalize_L2(dense_query) |
| _, faiss_ids = self.faiss_index.search(dense_query, candidate_count) |
|
|
| fused_ids = _reciprocal_rank_fusion(bm25_ids[0], faiss_ids[0]) |
| return [self.chunks[i] for i in fused_ids[:candidate_count]] |
|
|
| def search_text(self, query, k=5): |
| found = self.search(query, k) |
| return "\n\n---\n\n".join(found) if found else "(No relevant excerpts found.)" |
|
|
