| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
| from typing import Any, Iterable, Mapping, Sequence |
| import math |
| import re |
|
|
| import numpy as np |
| from sklearn.feature_extraction.text import TfidfVectorizer, ENGLISH_STOP_WORDS |
|
|
|
|
| def _clip01(x: float, eps: float = 1e-6) -> float: |
| return min(max(float(x), eps), 1.0) |
|
|
|
|
| def _logsumexp(xs: Sequence[float]) -> float: |
| if not xs: |
| return float('-inf') |
| finite = [x for x in xs if math.isfinite(x)] |
| if not finite: |
| return float('-inf') |
| m = max(finite) |
| return m + math.log(sum(math.exp(x - m) for x in finite)) |
|
|
|
|
| @dataclass(frozen=True) |
| class DocumentRecord: |
| id: str |
| text: str |
| title: str = "" |
| source: str = "unknown" |
| metadata: Mapping[str, Any] = field(default_factory=dict) |
|
|
| @property |
| def retrieval_text(self) -> str: |
| return (self.title + ". " + self.text).strip(". ") |
|
|
| def source_prior(self) -> float: |
| md = self.metadata |
| a = max(float(md.get("reliability_alpha", 1.0)), 1e-6) |
| b = max(float(md.get("reliability_beta", 1.0)), 1e-6) |
| reliability = a / (a + b) |
| authority = _clip01(md.get("authority", 0.5)) |
| freshness = _clip01(md.get("freshness", 0.5)) |
| return reliability * authority * freshness |
|
|
|
|
| @dataclass(frozen=True) |
| class RetrievalHit: |
| document: DocumentRecord |
| retrieval_score: float |
| log_prior: float |
| log_likelihood: float |
| log_posterior_unnormalized: float |
| posterior: float |
| query_coverage: float |
| rank: int |
|
|
| def as_dict(self) -> dict[str, Any]: |
| return { |
| "id": self.document.id, |
| "title": self.document.title, |
| "text": self.document.text, |
| "source": self.document.source, |
| "metadata": dict(self.document.metadata), |
| "retrieval_score": self.retrieval_score, |
| "log_prior": self.log_prior, |
| "log_likelihood": self.log_likelihood, |
| "posterior": self.posterior, |
| "query_coverage": self.query_coverage, |
| "rank": self.rank, |
| } |
|
|
|
|
| @dataclass(frozen=True) |
| class BayesianRetrievalConfig: |
| candidate_pool: int = 12 |
| likelihood_temperature: float = 2.0 |
| likelihood_floor: float = 0.02 |
| reliability_weight: float = 1.0 |
| authority_weight: float = 1.0 |
| freshness_weight: float = 0.8 |
| prior_floor: float = 1e-4 |
| min_relevance: float = 0.06 |
| min_top_posterior: float = 0.30 |
| min_query_coverage: float = 0.30 |
| min_dominance_ratio: float = 1.20 |
|
|
|
|
| class BayesianIndex: |
| """Generalized-Bayesian reranker over TF-IDF retrieval candidates. |
| |
| Candidate document hypotheses are scored by |
| p(d|q) ∝ p(d) exp(tau * sim(q,d)). |
| p(d) is induced by a Beta source-reliability mean plus authority/freshness |
| metadata. The exponential similarity term is a Gibbs/generalized-Bayes |
| likelihood, not claimed to be a literal language-model likelihood. |
| """ |
|
|
| def __init__(self, documents: Sequence[DocumentRecord], config: BayesianRetrievalConfig | None = None): |
| if not documents: |
| raise ValueError("documents must be non-empty") |
| self.documents = list(documents) |
| self.config = config or BayesianRetrievalConfig() |
| self.vectorizer = TfidfVectorizer( |
| lowercase=True, |
| stop_words="english", |
| ngram_range=(1, 2), |
| sublinear_tf=True, |
| norm="l2", |
| ) |
| self.matrix = self.vectorizer.fit_transform([d.retrieval_text for d in self.documents]) |
|
|
| def _log_prior(self, d: DocumentRecord) -> float: |
| md = d.metadata |
| a = max(float(md.get("reliability_alpha", 1.0)), 1e-6) |
| b = max(float(md.get("reliability_beta", 1.0)), 1e-6) |
| rel = _clip01(a / (a + b)) |
| auth = _clip01(md.get("authority", 0.5)) |
| fresh = _clip01(md.get("freshness", 0.5)) |
| cfg = self.config |
| prior = max( |
| cfg.prior_floor, |
| (rel ** cfg.reliability_weight) |
| * (auth ** cfg.authority_weight) |
| * (fresh ** cfg.freshness_weight), |
| ) |
| return math.log(prior) |
|
|
| def lexical_scores(self, query: str) -> np.ndarray: |
| q = self.vectorizer.transform([str(query)]) |
| |
| return (self.matrix @ q.T).toarray().ravel() |
|
|
| @staticmethod |
| def _query_coverage(query: str, text: str) -> float: |
| qterms = {t for t in re.findall(r"[A-Za-z0-9_-]+", query.lower()) if len(t) > 1 and t not in ENGLISH_STOP_WORDS} |
| if not qterms: |
| return 0.0 |
| dterms = set(re.findall(r"[A-Za-z0-9_-]+", text.lower())) |
| return len(qterms & dterms) / len(qterms) |
|
|
| def search(self, query: str, k: int = 4) -> list[RetrievalHit]: |
| if k <= 0: |
| return [] |
| scores = self.lexical_scores(query) |
| pool_n = min(max(k, self.config.candidate_pool), len(self.documents)) |
| |
| cand = np.argsort(-scores, kind="stable")[:pool_n] |
| logu: list[float] = [] |
| rows: list[tuple[int, float, float, float]] = [] |
| for i in cand: |
| d = self.documents[int(i)] |
| s = max(float(scores[int(i)]), 0.0) |
| lp = self._log_prior(d) |
| ll = self.config.likelihood_temperature * math.log(s + self.config.likelihood_floor) |
| u = lp + ll |
| rows.append((int(i), s, lp, ll)) |
| logu.append(u) |
| lse = _logsumexp(logu) |
| probs = [math.exp(u - lse) if math.isfinite(lse) else 1.0 / len(logu) for u in logu] |
| ranked = sorted( |
| zip(rows, probs), |
| key=lambda rp: (rp[1], rp[0][1]), |
| reverse=True, |
| ) |
| hits: list[RetrievalHit] = [] |
| for rank, ((i, s, lp, ll), pr) in enumerate(ranked[:k], start=1): |
| hits.append(RetrievalHit( |
| document=self.documents[i], |
| retrieval_score=s, |
| log_prior=lp, |
| log_likelihood=ll, |
| log_posterior_unnormalized=lp + ll, |
| posterior=float(pr), |
| query_coverage=self._query_coverage(query, self.documents[i].retrieval_text), |
| rank=rank, |
| )) |
| return hits |
|
|
| def should_abstain(self, hits: Sequence[RetrievalHit]) -> bool: |
| if not hits: |
| return True |
| top = hits[0] |
| if ( |
| top.retrieval_score < self.config.min_relevance |
| or top.posterior < self.config.min_top_posterior |
| or top.query_coverage < self.config.min_query_coverage |
| ): |
| return True |
| if len(hits) > 1: |
| second = hits[1] |
| dominance = top.posterior / max(second.posterior, 1e-12) |
| if dominance < self.config.min_dominance_ratio: |
| return True |
| return False |
|
|
|
|
| class ExtractiveGroundedGenerator: |
| """Dependency-free grounded answer synthesizer for tests and offline use. |
| |
| It selects the query-relevant sentence from the highest-posterior evidence |
| and returns that evidence sentence verbatim with its document id. It never |
| invents content not present in the retrieved context. |
| """ |
|
|
| abstain_text = "I do not know based on the provided documents." |
|
|
| @staticmethod |
| def _sentences(text: str) -> list[str]: |
| parts = re.split(r"(?<=[.!?])\s+|\n+", text.strip()) |
| return [p.strip() for p in parts if p.strip()] |
|
|
| def generate(self, question: str, hits: Sequence[RetrievalHit]) -> tuple[str, list[str]]: |
| if not hits: |
| return self.abstain_text, [] |
| |
| top = hits[0] |
| sents = self._sentences(top.document.text) |
| if not sents: |
| return self.abstain_text, [] |
| if len(sents) == 1: |
| sentence = sents[0] |
| else: |
| v = TfidfVectorizer(lowercase=True, stop_words="english", ngram_range=(1, 2)) |
| mat = v.fit_transform([question] + sents) |
| q = mat[0] |
| sims = (mat[1:] @ q.T).toarray().ravel() |
| sentence = sents[int(np.argmax(sims))] |
| return f"{sentence} [{top.document.id}]", [top.document.id] |
|
|
|
|
| class PlainTfidfIndex: |
| """Conventional lexical baseline used by the end-to-end RAG benchmark.""" |
| def __init__(self, documents: Sequence[DocumentRecord]): |
| self.documents = list(documents) |
| self.vectorizer = TfidfVectorizer(lowercase=True, stop_words="english", ngram_range=(1, 2), sublinear_tf=True, norm="l2") |
| self.matrix = self.vectorizer.fit_transform([d.retrieval_text for d in self.documents]) |
|
|
| @staticmethod |
| def _query_coverage(query: str, text: str) -> float: |
| qterms = {t for t in re.findall(r"[A-Za-z0-9_-]+", query.lower()) if len(t) > 1 and t not in ENGLISH_STOP_WORDS} |
| if not qterms: |
| return 0.0 |
| dterms = set(re.findall(r"[A-Za-z0-9_-]+", text.lower())) |
| return len(qterms & dterms) / len(qterms) |
|
|
| def search(self, query: str, k: int = 4) -> list[RetrievalHit]: |
| q = self.vectorizer.transform([query]) |
| scores = (self.matrix @ q.T).toarray().ravel() |
| order = np.argsort(-scores, kind="stable")[:k] |
| |
| vals = [max(float(scores[int(i)]), 0.0) for i in order] |
| z = sum(vals) or 1.0 |
| return [RetrievalHit(self.documents[int(i)], float(scores[int(i)]), 0.0, 0.0, 0.0, vals[j]/z, 0.0, j+1) for j, i in enumerate(order)] |
|
|
