"""Lexical retrieval over the corpus: Okapi BM25 on Russian Snowball stems + synonym expansion. Offline and deterministic (ADR-004): no embeddings, no network. Stop words are dropped before stemming; the section title is indexed `title_weight` times so a match with "3.2 Психофармакотерапия" counts more than a passing mention in the body. Re-ranking: the best `rerank_candidates` chunks by BM25 get a bonus for their densest passage — a window of `window_sentences` consecutive sentences scored by the idf of query terms it contains. A chunk that answers in one place beats a long one that mentions the words in passing; the passage itself is returned with the hit and is what `answer()` quotes from. """ import math import re from collections import Counter from dataclasses import dataclass from functools import lru_cache from app.llm.rag.corpus import Audience, Chunk, RagConfig, load_chunks, load_config, split_sentences from app.llm.rag.stemmer import stem _WORD = re.compile(r"[а-яёa-z0-9]+(?:-[а-яёa-z0-9]+)*", re.IGNORECASE) # `privacy.scrub()` placeholders ("[NAME]", "[PHONE]") are not words of the question. _PLACEHOLDER = re.compile(r"\[[A-Z_]+\]") # Function words and generic KR boilerplate that carry no topical signal. STOP_WORDS = frozenset( """ а без более бы был была были было быть в вам вас ваш весь во вот все всего всех вы где да даже для до его ее если есть еще же за и из или им их к как ко когда кто ли либо мне может можно мой мы на над не него нее нет ни них но ну о об однако он она они оно от по под при про с со так также такой там те то тоже только том тот у уже хотя чего чем что чтобы эта эти это этот я ли ней нем какой какая какое какие каков такое такая такие сколько почему зачем меня мой моя мое мои себя сегодня будет """.split() # noqa: SIM905 — a word list reads better as text ) def tokens(text: str) -> list[str]: """Stems of content words (hyphenated words are kept whole and also split: "когнитивно-поведенческая").""" out: list[str] = [] for word in _WORD.findall(_PLACEHOLDER.sub(" ", text).lower().replace("ё", "е")): parts = [p for p in word.split("-") if p not in STOP_WORDS and (len(p) > 1 or p.isdigit())] if "-" in word and len(parts) > 1: out.append("-".join(stem(p) for p in parts)) # the compound itself, in any case form out.extend(stem(p) for p in parts) return out EXPANSION_WEIGHT = 0.7 SYNONYM_CREDIT = 0.8 PREFIX_LEN = 6 @dataclass(frozen=True) class Hit: chunk: Chunk score: float # BM25 + passage bonus (ranking key) bm25: float passage: str # densest window of the chunk for the query def sentences(text: str) -> list[str]: """Sentences of a chunk; every line (table row, list item) is split on its own.""" return [s for line in text.splitlines() for s in split_sentences(line)] class SynonymExpander: """Adds the stems of every phrase in a synonym group when any phrase of the group is in the query. Added stems weigh `EXPANSION_WEIGHT`: they widen recall but should not outrank the patient's own words. """ def __init__(self, groups: list[list[str]]) -> None: self._groups = [[tokens(phrase) for phrase in group] for group in groups] # stem -> stems of the other phrases of its groups ("отменя" -> {"отмен", "брос", ...}) self.related: dict[str, set[str]] = {} for group in self._groups: stems = {t for phrase in group for t in phrase} for t in stems: self.related.setdefault(t, set()).update(stems - {t}) def expand(self, query_tokens: list[str]) -> dict[str, float]: weights: dict[str, float] = {} for t in query_tokens: weights[t] = weights.get(t, 0.0) + 1.0 present = set(query_tokens) for group in self._groups: if any(phrase and set(phrase) <= present for phrase in group): for phrase in group: for t in phrase: if t not in present: # a 3-word expansion of "птср" is one concept, not three weights[t] = max(weights.get(t, 0.0), EXPANSION_WEIGHT / len(phrase)) return weights class BM25Index: def __init__(self, chunks: tuple[Chunk, ...], cfg: RagConfig) -> None: r = cfg.retrieval self._k1, self._b, self._min_score, self._top_k = r.bm25_k1, r.bm25_b, r.min_score, r.top_k self._candidates, self._window, self._window_weight = r.rerank_candidates, r.window_sentences, r.window_weight self._section_boost = r.section_boost self._doc_cues = {doc: re.compile("|".join(cues), re.IGNORECASE) for doc, cues in r.doc_cues.items() if cues} self._doc_cue_boost = r.doc_cue_boost self.chunks = chunks self._sentences = [[(s, set(tokens(s))) for s in sentences(c.text)] for c in chunks] self._expander = SynonymExpander(cfg.synonyms) self._tf: list[Counter[str]] = [] df: Counter[str] = Counter() for c in chunks: tf = Counter(tokens(c.text) + tokens(c.section) * r.title_weight) self._tf.append(tf) df.update(tf.keys()) self._len = [sum(tf.values()) for tf in self._tf] self._avg_len = sum(self._len) / max(len(self._len), 1) n = len(chunks) self._idf = {t: math.log(1 + (n - f + 0.5) / (f + 0.5)) for t, f in df.items()} def query_terms(self, query: str) -> dict[str, float]: """Query stems -> weight (repeated words count more, synonym expansions less).""" return self._expander.expand(tokens(query)) def _score(self, i: int, terms: dict[str, float]) -> float: tf, norm = self._tf[i], self._k1 * (1 - self._b + self._b * self._len[i] / self._avg_len) score = 0.0 for term, weight in terms.items(): f = tf.get(term) if f: score += self._idf[term] * f * (self._k1 + 1) / (f + norm) * weight return score def boost(self, chunk: Chunk, audience: Audience | None, cued: set[str] | None = None) -> float: """Section multiplier for the audience (`retrieval.section_boost`) × disease-cue multiplier.""" boosts = self._section_boost.get(audience or "", {}) section = next((b for prefix, b in boosts.items() if chunk.section.startswith(prefix)), 1.0) return section * (self._doc_cue_boost if cued and chunk.doc.id in cued else 1.0) def cued_docs(self, query: str) -> set[str]: """Documents whose disease the query names ("ТМС при депрессии" -> {"kr_depression"}). Only a one-sided mention counts: a question naming both disorders favours neither. """ q = _PLACEHOLDER.sub(" ", query).replace("ё", "е") cued = {doc for doc, cue in self._doc_cues.items() if cue.search(q)} return cued if len(cued) < len(self._doc_cues) else set() def overlap(self, text: str, terms: dict[str, float]) -> float: """Sum of idf × weight of the query terms present in `text` (sentence relevance for `answer()`).""" present = set(tokens(text)) return sum(self._idf.get(t, 0.0) * w for t, w in terms.items() if t in present) def coverage(self, text: str, query: str) -> float: """Share of the query's own (not synonym-expanded) idf mass found in `text`, 0..1. A query word whose synonym — or a word with the same `PREFIX_LEN`-letter prefix, for stems Snowball leaves apart ("скрининговый"/"скрининг") — is in the text counts as `SYNONYM_CREDIT` of itself. A word absent from the corpus counts with twice the maximal idf — the strongest sign of an off-topic question: "как установить windows" is then mostly uncovered and `answer()` says there is no answer instead of quoting the nearest paragraph about "установить диагноз". """ own = set(tokens(query)) if not own: return 0.0 unseen = 2 * max(self._idf.values(), default=1.0) total = sum(self._idf.get(t, unseen) for t in own) present = set(tokens(text)) prefixes = {p[:PREFIX_LEN] for p in present if len(p) >= PREFIX_LEN} covered = 0.0 for t in own: if t in present: covered += self._idf[t] elif self._expander.related.get(t, set()) & present or ( len(t) >= PREFIX_LEN and t[:PREFIX_LEN] in prefixes ): covered += SYNONYM_CREDIT * self._idf.get(t, unseen) # said in the corpus' own words return min(1.0, covered / total) def _passage(self, i: int, terms: dict[str, float]) -> tuple[float, str]: sents = self._sentences[i] best = (0.0, sents[0][0] if sents else "") for start in range(len(sents)): window = sents[start : start + self._window] present = set().union(*(toks for _, toks in window)) score = sum(self._idf.get(t, 0.0) * w for t, w in terms.items() if t in present) if score > best[0]: best = (score, " ".join(text for text, _ in window)) return best def search( self, query: str, *, audience: Audience | None = None, docs: set[str] | None = None, top_k: int | None = None, min_score: float | None = None, ) -> list[Hit]: """Best chunks for `query`, filtered by audience/doc ids; empty if nothing passes `min_score` (BM25).""" terms = self.query_terms(query) if not terms: return [] threshold = self._min_score if min_score is None else min_score scored = [ (s, i) for i, c in enumerate(self.chunks) if (audience is None or audience in c.doc.audiences) and (docs is None or c.doc.id in docs) and (s := self._score(i, terms)) >= threshold ] scored.sort(key=lambda x: (-x[0], x[1])) cued = self.cued_docs(query) hits = [] for bm25, i in scored[: max(self._candidates, top_k or 0)]: bonus, passage = self._passage(i, terms) chunk = self.chunks[i] hits.append( Hit(chunk, (bm25 + self._window_weight * bonus) * self.boost(chunk, audience, cued), bm25, passage) ) hits.sort(key=lambda h: -h.score) # stable: BM25 order breaks ties return hits[: top_k or self._top_k] @lru_cache def get_index() -> BM25Index: return BM25Index(load_chunks(), load_config()) def retrieve(query: str, **kwargs) -> list[Hit]: """Module-level shortcut: `retrieve("когда начинает действовать антидепрессант", audience="patient")`.""" return get_index().search(query, **kwargs)