Spaces:
Runtime error
Runtime error
File size: 3,811 Bytes
b74c6ae | 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 | """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
@dataclass
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
|