Spaces:
Running
Running
| """Cross-encoder reranking + sentence-window context expansion, as used by | |
| Notebook 5. | |
| `Reranker.rerank` re-scores a candidate list with a cross-encoder and keeps | |
| the top `top_n`. `relevance_score` sigmoid-normalizes those scores into a | |
| 0-1 "context relevance" proxy. `sentence_window_expand` widens a chunk to | |
| include `window` sentences of surrounding context from its source document. | |
| """ | |
| import re | |
| import numpy as np | |
| import pandas as pd | |
| from sentence_transformers import CrossEncoder | |
| _SENTENCE_SPLIT_RE = re.compile(r"(?<=[.!?])\s+") | |
| def sigmoid(x: float) -> float: | |
| return 1 / (1 + np.exp(-x)) | |
| def split_sentences(text: str) -> list: | |
| return [s.strip() for s in _SENTENCE_SPLIT_RE.split(text) if s.strip()] | |
| def build_doc_sentences(docs_df: pd.DataFrame) -> dict: | |
| """Maps doc_id -> list of sentences, used by `sentence_window_expand`.""" | |
| return {row["doc_id"]: split_sentences(row["text"]) for _, row in docs_df.iterrows()} | |
| def relevance_score(reranked: list) -> float: | |
| """Mean sigmoid-normalized cross-encoder relevance of the reranked chunks.""" | |
| if not reranked: | |
| return 0.0 | |
| return float(np.mean([sigmoid(s) for _, s in reranked])) | |
| class Reranker: | |
| def __init__(self, model_name: str): | |
| self.model_name = model_name | |
| self.model = CrossEncoder(model_name) | |
| def rerank(self, query: str, candidate_idxs: list, chunks_df: pd.DataFrame, top_n: int) -> list: | |
| if not candidate_idxs: | |
| return [] | |
| pairs = [(query, chunks_df.iloc[idx]["text"]) for idx, _ in candidate_idxs] | |
| scores = self.model.predict(pairs) | |
| order = np.argsort(scores)[::-1][:top_n] | |
| return [(candidate_idxs[i][0], float(scores[i])) for i in order] | |
| def _locate_in_sentences(sentences: list, target: str, from_end: bool = False) -> int | None: | |
| """Find index of target in sentences list. | |
| Tries exact match first. If that fails, uses a 40-char anchor from the | |
| start (or end when from_end=True) of target to handle the common case | |
| where recursive chunking cuts a sentence in half — the chunk boundary | |
| lands mid-sentence so target is a fragment not present verbatim in the | |
| document sentence list. | |
| """ | |
| for i, s in enumerate(sentences): | |
| if s == target: | |
| return i | |
| anchor = target[-40:].strip() if from_end else target[:40].strip() | |
| if len(anchor) < 10: | |
| return None | |
| for i, s in enumerate(sentences): | |
| if anchor in s or s in target: | |
| return i | |
| return None | |
| def sentence_window_expand(chunk_idx: int, chunks_df: pd.DataFrame, doc_sentences: dict, window: int = 1): | |
| """Returns (expanded_text, (start, end)) -- range is None if the chunk's | |
| sentences couldn't be located in the source document (falls back to the | |
| raw chunk text).""" | |
| chunk = chunks_df.iloc[chunk_idx] | |
| sentences = doc_sentences[chunk["doc_id"]] | |
| chunk_sentences = split_sentences(chunk["text"]) | |
| if not chunk_sentences: | |
| return chunk["text"], None | |
| first_idx = _locate_in_sentences(sentences, chunk_sentences[0], from_end=False) | |
| last_idx = _locate_in_sentences(sentences, chunk_sentences[-1], from_end=True) | |
| if first_idx is None or last_idx is None: | |
| return chunk["text"], None | |
| start = max(0, first_idx - window) | |
| end = min(len(sentences), last_idx + 1 + window) | |
| expanded = " ".join(sentences[start:end]) | |
| return expanded, (start, end) | |