Spaces:
Runtime error
Runtime error
| """BM25 over the help corpus: two indexes (title, keywords, headings and questions weigh 1.0; body 0.4) plus | |
| context boosts for the page the writer is on and the docs the bot cited last turn.""" | |
| import re | |
| from collections.abc import Iterable | |
| from dataclasses import dataclass | |
| from rank_bm25 import BM25Okapi | |
| from app.support.corpus import Doc, get_corpus | |
| _STOPWORDS = frozenset( | |
| """the a an and or but if then else of for to in on at by with from as is are was were be been being have has had | |
| do does did can could should would may might will shall must you he she it we they them us our your their my me him | |
| her this that these those there here what which who whom whose how why when where not no yes so up down out about | |
| into over under again just very more most some any all each other than too""".split() | |
| ) | |
| _TOKEN = re.compile(r"[a-z0-9]+") | |
| _SUFFIXES = ("ingly", "edly", "ing", "ies", "ied", "ed", "es", "s", "ly") | |
| HIGH_WEIGHT = 1.0 | |
| BODY_WEIGHT = 0.4 | |
| HISTORY_WEIGHT = 0.5 | |
| ROUTE_BOOST = 3.0 | |
| CONTINUITY_BOOST = 2.0 | |
| def tokenize(text: str) -> list[str]: | |
| out: list[str] = [] | |
| for tok in _TOKEN.findall((text or "").lower()): | |
| if len(tok) < 3 or tok in _STOPWORDS: | |
| continue | |
| for suf in _SUFFIXES: # light stemming: syncing/synced/syncs -> sync | |
| if len(tok) > len(suf) + 2 and tok.endswith(suf): | |
| tok = tok[: -len(suf)] | |
| break | |
| out.append(tok) | |
| return out | |
| class Scored: | |
| doc: Doc | |
| score: float | |
| breakdown: dict[str, float] | |
| class Retriever: | |
| def __init__(self, docs: list[Doc]): | |
| self.docs = docs | |
| high = [tokenize(d.high_signal_text) or [d.id] for d in docs] | |
| body = [tokenize(d.body) or high[i] for i, d in enumerate(docs)] | |
| self._high = BM25Okapi(high) | |
| self._body = BM25Okapi(body) | |
| def _scores(self, index: BM25Okapi, weighted: list[tuple[str, float]]) -> list[float]: | |
| total = [0.0] * len(self.docs) | |
| buckets: dict[float, list[str]] = {} | |
| for tok, w in weighted: # rank-bm25 has no per-token weights, so score each weight bucket and sum | |
| buckets.setdefault(w, []).append(tok) | |
| for w, toks in buckets.items(): | |
| for i, s in enumerate(index.get_scores(toks)): | |
| total[i] += w * float(s) | |
| return total | |
| def retrieve( | |
| self, | |
| query: str, | |
| *, | |
| history: Iterable[str] = (), | |
| page_path: str | None = None, | |
| last_cited: Iterable[str] = (), | |
| top_k: int = 3, | |
| min_score: float = 0.5, | |
| ) -> list[Scored]: | |
| weighted = [(t, 1.0) for t in tokenize(query)] | |
| for prior in list(history)[-2:]: # short follow ups ("and on the free plan?") borrow earlier words | |
| weighted += [(t, HISTORY_WEIGHT) for t in tokenize(prior)] | |
| if not weighted: | |
| return [] | |
| high, body = self._scores(self._high, weighted), self._scores(self._body, weighted) | |
| cited = set(last_cited) | |
| scored = [] | |
| for i, doc in enumerate(self.docs): | |
| parts = { | |
| "high": HIGH_WEIGHT * high[i], | |
| "body": BODY_WEIGHT * body[i], | |
| "route": ROUTE_BOOST if page_path and (page_path == doc.route or page_path in doc.related_paths) else 0.0, | |
| "continuity": CONTINUITY_BOOST if doc.id in cited else 0.0, | |
| } | |
| scored.append(Scored(doc, sum(parts.values()), {k: round(v, 3) for k, v in parts.items()})) | |
| scored.sort(key=lambda s: s.score, reverse=True) | |
| return [s for s in scored[:top_k] if s.score >= min_score] | |
| _default: Retriever | None = None | |
| def get_retriever() -> Retriever: | |
| global _default | |
| if _default is None: | |
| _default = Retriever(get_corpus()) | |
| return _default | |