ElebobaandClaude Opus 5.5
fe29ebd629
CI / ruff · pytest (pull_request) Successful in 1m2s
- converter marks a paragraph glued across a page break with ⟦стр. N⟧; the chunker
strips the marker and extends the chunk's page range (wrong-page citations 24 -> 3,
the rest are scale-form table rows split by a page)
- [NAME]/[PHONE] placeholders from privacy.scrub() no longer count as unknown words
(a scrubbed question used to fall below the coverage gate)
- a question naming one disease boosts its guideline in retrieval and passage choice
("ТМС при депрессии" was answered from the PTSD guideline); "антидепрессант" is not a cue
Refs #62
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
230 lines
11 KiB
Python
230 lines
11 KiB
Python
"""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)
|