File size: 3,569 Bytes
d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa | 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 107 108 109 110 111 112 | """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.)"
|