feat(llm): RAG knowledge base over clinical guidelines with cited extractive answers
CI / ruff · pytest (pull_request) Successful in 1m30s
CI / ruff · pytest (pull_request) Successful in 1m30s
- docs/source/rag: PDF sources (KR PTSD 2023, KR depression 2024, meds module spec), pdf->txt converter (tables as markdown, page markers, hand-transcribed figures), word-level check against pdftotext; fixes HAM-D/MADRS fake headers and split rows - config/rag: corpus txt + sources.yaml (documents, audiences, retrieval params, synonyms) - app/llm/rag: Snowball stemmer, section-aware chunker with page citations, BM25 + synonym expansion + densest-passage reranking, extractive answer() with coverage gate, patient dose/prescribing filter, optional LLM rephrasing rejected when it adds content; CLI python -m app.llm.rag - tests: gold retrieval set (hit@3 >= 15/18), off-topic, safety, faithfulness - ADR-011 Refs #60 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Claude Opus 5.5
@@ -0,0 +1,24 @@
|
||||
"""Knowledge base (RAG, #60): clinical guidelines -> citable chunks -> BM25 retrieval -> grounded answers.
|
||||
|
||||
from app.llm.rag import answer, retrieve
|
||||
hits = retrieve("синдром отмены антидепрессантов", audience="doctor")
|
||||
reply = await answer(scrub(text), audience="patient") # None -> rag.fallback_text()
|
||||
|
||||
Corpus and parameters: `config/rag/sources.yaml`; PDF sources and the converter: `docs/source/rag/`.
|
||||
"""
|
||||
|
||||
from app.llm.rag.corpus import Chunk, load_chunks
|
||||
from app.llm.rag.index import Hit, get_index, retrieve
|
||||
from app.llm.rag.qa import Citation, RagAnswer, answer, fallback_text
|
||||
|
||||
__all__ = [
|
||||
"Chunk",
|
||||
"Citation",
|
||||
"Hit",
|
||||
"RagAnswer",
|
||||
"answer",
|
||||
"fallback_text",
|
||||
"get_index",
|
||||
"load_chunks",
|
||||
"retrieve",
|
||||
]
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Manual check of the knowledge base on synthetic questions.
|
||||
|
||||
python -m app.llm.rag "через сколько начинают действовать антидепрессанты"
|
||||
python -m app.llm.rag --audience doctor "синдром отмены антидепрессантов"
|
||||
python -m app.llm.rag --search "ДПДГ при ПТСР" # raw retrieval: chunks, scores, passages
|
||||
python -m app.llm.rag --stats # corpus summary
|
||||
|
||||
The question always goes through `privacy.scrub()` first; the CLI never bypasses it.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
from collections import Counter
|
||||
|
||||
from app.core import privacy
|
||||
from app.llm.gateway import get_gateway
|
||||
from app.llm.rag.index import get_index
|
||||
from app.llm.rag.qa import answer, fallback_text
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(prog="python -m app.llm.rag")
|
||||
parser.add_argument("question", nargs="?")
|
||||
parser.add_argument("--audience", choices=["patient", "doctor", "team"], default="patient")
|
||||
parser.add_argument("--search", action="store_true", help="show retrieved chunks instead of an answer")
|
||||
parser.add_argument("--top", type=int, default=5)
|
||||
parser.add_argument("--stats", action="store_true")
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
index = get_index()
|
||||
if args.stats:
|
||||
by_doc = Counter(c.doc.short for c in index.chunks)
|
||||
print(f"chunks={len(index.chunks)}", ", ".join(f"{d}: {n}" for d, n in by_doc.items()))
|
||||
return 0
|
||||
if not args.question:
|
||||
parser.error("question is required")
|
||||
question = privacy.scrub(args.question)
|
||||
|
||||
if args.search:
|
||||
for hit in index.search(question, audience=args.audience, top_k=args.top, min_score=0):
|
||||
print(f"{hit.score:6.1f} (bm25 {hit.bm25:5.1f}) {hit.chunk.citation} {hit.chunk.section}")
|
||||
print(f" > {hit.passage[:300]}")
|
||||
return 0
|
||||
|
||||
print(f"provider={get_gateway().provider_name} audience={args.audience}\n")
|
||||
reply = asyncio.run(answer(question, audience=args.audience, index=index))
|
||||
print(reply.text if reply else fallback_text(args.audience))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,297 @@
|
||||
"""Knowledge-base corpus: `config/rag/sources.yaml` + txt files -> citable chunks.
|
||||
|
||||
The txt files come from `docs/source/rag/pdf_to_txt.py`: paragraphs are separated by blank lines, pages by
|
||||
`--- стр. N ---` markers, tables are markdown, text transcribed from figures follows an
|
||||
`[Изображение: ...]` line. Chunking rules:
|
||||
|
||||
* a chunk never crosses a section boundary (numbered headings "2.1 ...", "Приложение Г5", ...), so its
|
||||
section title is reliable context; the title is stored separately and indexed with extra weight;
|
||||
* "Комментарии: ..." after a recommendation stays with the recommendation (the comment explains it);
|
||||
* a line starting in lowercase continues the previous paragraph (cut by a page break the converter missed);
|
||||
* tables are split by rows, every table piece repeats the table header, so a chunk is readable alone;
|
||||
* each chunk keeps the page range it came from, for the citation "[КР ПТСР 2023, стр. 45]".
|
||||
"""
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.config import get_settings
|
||||
|
||||
Audience = str # "doctor" | "patient" | "team"
|
||||
|
||||
PAGE_MARKER = re.compile(r"^--- стр\. (\d+) ---$")
|
||||
TABLE_DIVIDER = re.compile(r"^\|(-+\|)+$")
|
||||
# "2.1 Жалобы и анамнез", "3.1.2. Купирующая терапия ..." — a numbered subsection title.
|
||||
SUBSECTION = re.compile(r"^(\d{1,2}(?:\.\d{1,2})+)\.?\s+([А-ЯЁA-Z].*)$") # not dates "28.07.2022 N ..."
|
||||
# Top-level KR sections: "1. Краткая информация ...", "3. Лечение, включая ...". Plain "1. ..." lines are
|
||||
# list items, so a top-level heading must start with one of the standard KR section names.
|
||||
TOPLEVEL = re.compile(
|
||||
r"^([1-9])\.\s+(Краткая информация|Диагностика|Лечение|Медицинская реабилитация|Профилактика"
|
||||
r"|Организация оказания|Дополнительная информация)"
|
||||
)
|
||||
NAMED = re.compile(
|
||||
r"^(Приложение [А-ЯЁ]\d*(?:-[А-ЯЁ]+\d*)?\b.*|Критерии оценки качества медицинской помощи|Список литературы"
|
||||
r"|Список сокращений|Термины и определения)$"
|
||||
)
|
||||
# Unnumbered short title lines (team specs): capitalised, short, no final punctuation, not a list item.
|
||||
SHORT_TITLE = re.compile(r"^[А-ЯЁ][^.:;!?|]{2,60}$")
|
||||
COMMENT = re.compile(r"^Комментари[йи]\s*:")
|
||||
SENTENCE_SPLIT = re.compile(r"(?<=[.!?…])\s+(?=[А-ЯЁA-Z0-9«\"(])")
|
||||
# Abbreviations after which a capital letter does not start a new sentence: "(см. Приложение Б5)".
|
||||
ABBREVIATION_END = re.compile(
|
||||
r"(?:\b(?:см|т\.\s?е|т\.\s?д|т\.\s?п|др|напр|рис|табл|стр|г|гг|им|ср|англ|лат|греч)|\b[А-ЯЁA-Z])\.$"
|
||||
)
|
||||
MAX_TITLE_CHARS = 300
|
||||
|
||||
|
||||
class DocumentMeta(BaseModel):
|
||||
id: str
|
||||
file: str
|
||||
title: str
|
||||
short: str
|
||||
issuer: str
|
||||
year: int
|
||||
icd: list[str] = Field(default_factory=list)
|
||||
audiences: list[Audience]
|
||||
|
||||
|
||||
class ChunkingConfig(BaseModel):
|
||||
max_chars: int = 1200
|
||||
min_chars: int = 200
|
||||
skip_sections: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class RetrievalConfig(BaseModel):
|
||||
top_k: int = 5
|
||||
bm25_k1: float = 1.5
|
||||
bm25_b: float = 0.75
|
||||
title_weight: int = 2
|
||||
min_score: float = 3.0
|
||||
rerank_candidates: int = 30
|
||||
window_sentences: int = 2
|
||||
window_weight: float = 0.5
|
||||
# section title prefix -> score multiplier, per audience ({"patient": {"Приложение В": 1.5}})
|
||||
section_boost: dict[str, dict[str, float]] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class AnswerConfig(BaseModel):
|
||||
max_chunks: int = 3
|
||||
max_passages: int = 2
|
||||
max_sentences: int = 3
|
||||
max_chars: int = 600
|
||||
min_coverage: float = 0.5
|
||||
second_passage_ratio: float = 0.8
|
||||
neighbor_ratio: float = 0.5
|
||||
|
||||
|
||||
class RagConfig(BaseModel):
|
||||
documents: list[DocumentMeta]
|
||||
chunking: ChunkingConfig = ChunkingConfig()
|
||||
retrieval: RetrievalConfig = RetrievalConfig()
|
||||
answer: AnswerConfig = AnswerConfig()
|
||||
synonyms: list[list[str]] = Field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Chunk:
|
||||
id: str # "<doc>#<n>"
|
||||
doc: DocumentMeta
|
||||
section: str # nearest heading ("" before the first one)
|
||||
page_start: int
|
||||
page_end: int
|
||||
text: str # display text: paragraphs and markdown table rows
|
||||
|
||||
@property
|
||||
def pages(self) -> str:
|
||||
return str(self.page_start) if self.page_start == self.page_end else f"{self.page_start}–{self.page_end}"
|
||||
|
||||
@property
|
||||
def citation(self) -> str:
|
||||
return f"[{self.doc.short}, стр. {self.pages}]"
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Unit:
|
||||
"""Indivisible piece of a chunk: a paragraph (with its comment) or one table row."""
|
||||
|
||||
text: str
|
||||
page: int
|
||||
table_header: str | None = None # "| a | b |\n|---|---|" for table rows
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Section:
|
||||
title: str
|
||||
units: list[_Unit] = field(default_factory=list)
|
||||
|
||||
|
||||
def rag_dir() -> Path:
|
||||
return get_settings().config_dir / "rag"
|
||||
|
||||
|
||||
@lru_cache
|
||||
def load_config(path: Path | None = None) -> RagConfig:
|
||||
path = path or rag_dir() / "sources.yaml"
|
||||
return RagConfig.model_validate(yaml.safe_load(path.read_text(encoding="utf-8")))
|
||||
|
||||
|
||||
def heading(line: str) -> str | None:
|
||||
"""The line's text if it is a section heading, else `None`."""
|
||||
if len(line) > MAX_TITLE_CHARS or line.startswith("|"):
|
||||
return None
|
||||
if SUBSECTION.match(line) or TOPLEVEL.match(line) or NAMED.match(line):
|
||||
return line
|
||||
return None
|
||||
|
||||
|
||||
def _is_short_title(block: list[str], next_block: list[str] | None) -> bool:
|
||||
"""A lone short line followed by a body paragraph (unnumbered headings in team docs)."""
|
||||
return (
|
||||
len(block) == 1
|
||||
and bool(SHORT_TITLE.match(block[0]))
|
||||
and next_block is not None
|
||||
and not PAGE_MARKER.match(next_block[0])
|
||||
and (len(next_block) > 1 or len(next_block[0]) > 80 or next_block[0].startswith("|"))
|
||||
)
|
||||
|
||||
|
||||
def _blocks(text: str) -> list[list[str]]:
|
||||
"""Blank-line separated blocks of non-empty lines."""
|
||||
return [[ln for ln in b.splitlines() if ln.strip()] for b in re.split(r"\n\s*\n", text) if b.strip()]
|
||||
|
||||
|
||||
def _sections(text: str, skip: list[str]) -> list[_Section]:
|
||||
blocks = _blocks(text)
|
||||
sections = [_Section(title="")]
|
||||
page = 1
|
||||
skipping = False
|
||||
for i, block in enumerate(blocks):
|
||||
first = block[0]
|
||||
if first.startswith("# ") or first.startswith("Источник:"):
|
||||
continue # file header written by the converter
|
||||
m = PAGE_MARKER.match(first)
|
||||
if m:
|
||||
page = int(m.group(1))
|
||||
continue
|
||||
nxt = next((b for b in blocks[i + 1 :] if not PAGE_MARKER.match(b[0])), None)
|
||||
if first.startswith("|"):
|
||||
if not skipping:
|
||||
sections[-1].units.extend(_table_units(block, page))
|
||||
continue
|
||||
if first.startswith("[Изображение"):
|
||||
if not skipping: # transcribed figure: one unit, keeps its own line breaks
|
||||
sections[-1].units.append(_Unit("\n".join(block[1:]) if len(block) > 1 else first, page))
|
||||
continue
|
||||
for j, line in enumerate(block):
|
||||
title = heading(line)
|
||||
if title is None and j == 0 and _is_short_title(block, nxt):
|
||||
title = line
|
||||
if title is not None:
|
||||
skipping = any(title.startswith(s) for s in skip)
|
||||
sections.append(_Section(title=title))
|
||||
continue
|
||||
if skipping:
|
||||
continue
|
||||
units = sections[-1].units
|
||||
if COMMENT.match(line) and units and units[-1].table_header is None:
|
||||
units[-1].text += "\n" + line
|
||||
elif line[0].islower() and units and units[-1].table_header is None:
|
||||
units[-1].text += " " + line # paragraph continued after a page break or a figure
|
||||
else:
|
||||
units.append(_Unit(line, page))
|
||||
return [s for s in sections if s.units]
|
||||
|
||||
|
||||
def _table_units(block: list[str], page: int) -> list[_Unit]:
|
||||
if len(block) >= 2 and TABLE_DIVIDER.match(block[1]):
|
||||
header = f"{block[0]}\n{block[1]}"
|
||||
rows = block[2:]
|
||||
else: # continuation without a header line (should not happen with the converter, but be safe)
|
||||
header, rows = "", block
|
||||
return [_Unit(row, page, table_header=header) for row in rows]
|
||||
|
||||
|
||||
def split_sentences(text: str) -> list[str]:
|
||||
"""Sentences of one line/paragraph; does not split after "см.", "т.е.", initials."""
|
||||
out: list[str] = []
|
||||
for part in SENTENCE_SPLIT.split(text):
|
||||
if out and ABBREVIATION_END.search(out[-1]):
|
||||
out[-1] += " " + part
|
||||
else:
|
||||
out.append(part)
|
||||
return [p.strip() for p in out if p.strip()]
|
||||
|
||||
|
||||
def _split_long(unit: _Unit, max_chars: int) -> list[_Unit]:
|
||||
if len(unit.text) <= max_chars or unit.table_header is not None:
|
||||
return [unit]
|
||||
parts: list[str] = []
|
||||
for sentence in split_sentences(unit.text):
|
||||
if parts and len(parts[-1]) + 1 + len(sentence) <= max_chars:
|
||||
parts[-1] += " " + sentence
|
||||
else:
|
||||
parts.append(sentence)
|
||||
return [_Unit(p, unit.page) for p in parts]
|
||||
|
||||
|
||||
def _render(units: list[_Unit]) -> str:
|
||||
out: list[str] = []
|
||||
header: str | None = None
|
||||
for u in units:
|
||||
if u.table_header is not None:
|
||||
if header != u.table_header:
|
||||
header = u.table_header
|
||||
if header:
|
||||
out.append(header)
|
||||
out.append(u.text)
|
||||
else:
|
||||
header = None
|
||||
out.append(u.text)
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def _pack(section: _Section, max_chars: int, min_chars: int) -> list[list[_Unit]]:
|
||||
units = [p for u in section.units for p in _split_long(u, max_chars)]
|
||||
groups: list[list[_Unit]] = [[]]
|
||||
size = 0
|
||||
for u in units:
|
||||
if groups[-1] and size + len(u.text) > max_chars:
|
||||
groups.append([])
|
||||
size = 0
|
||||
groups[-1].append(u)
|
||||
size += len(u.text) + 1
|
||||
if len(groups) > 1 and sum(len(u.text) for u in groups[-1]) < min_chars:
|
||||
groups[-2].extend(groups.pop())
|
||||
return [g for g in groups if g]
|
||||
|
||||
|
||||
def chunk_document(doc: DocumentMeta, text: str, cfg: ChunkingConfig) -> list[Chunk]:
|
||||
chunks: list[Chunk] = []
|
||||
for section in _sections(text, cfg.skip_sections):
|
||||
for group in _pack(section, cfg.max_chars, cfg.min_chars):
|
||||
chunks.append(
|
||||
Chunk(
|
||||
id=f"{doc.id}#{len(chunks)}",
|
||||
doc=doc,
|
||||
section=section.title,
|
||||
page_start=group[0].page,
|
||||
page_end=group[-1].page,
|
||||
text=_render(group),
|
||||
)
|
||||
)
|
||||
return chunks
|
||||
|
||||
|
||||
@lru_cache
|
||||
def load_chunks(config_path: Path | None = None) -> tuple[Chunk, ...]:
|
||||
cfg = load_config(config_path)
|
||||
base = (config_path.parent if config_path else rag_dir()).resolve()
|
||||
out: list[Chunk] = []
|
||||
for doc in cfg.documents:
|
||||
out.extend(chunk_document(doc, (base / doc.file).read_text(encoding="utf-8"), cfg.chunking))
|
||||
return tuple(out)
|
||||
@@ -0,0 +1,212 @@
|
||||
"""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)
|
||||
|
||||
# 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(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.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) -> float:
|
||||
"""Section multiplier for the audience (`retrieval.section_boost`), 1.0 by default."""
|
||||
boosts = self._section_boost.get(audience or "", {})
|
||||
return next((b for prefix, b in boosts.items() if chunk.section.startswith(prefix)), 1.0)
|
||||
|
||||
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]))
|
||||
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), 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)
|
||||
@@ -0,0 +1,317 @@
|
||||
"""Grounded answers from the knowledge base: every sentence shown is quoted from a cited chunk.
|
||||
|
||||
Pipeline (deterministic except for the optional patient rephrasing):
|
||||
|
||||
1. **Dose guard (patient).** Dose questions are never answered — the dose is the doctor's decision.
|
||||
2. **Retrieve.** Top `answer.max_chunks` chunks for the audience (`team` documents never reach patients
|
||||
or doctors, see `config/rag/sources.yaml`).
|
||||
3. **Sentences.** Chunks are split into sentences; KR boilerplate is stripped (literature refs "[1, 15]").
|
||||
Doctors also get table rows as "column: value; ..." lines and evidence grades. For patients, sentences
|
||||
about prescribing/dosing or with forbidden content (`QUOTE_FILTER`) are dropped.
|
||||
4. **Passage.** In each chunk the best sentence (idf-weighted overlap with the query) is extended with
|
||||
adjacent relevant sentences into one contiguous passage — the answer reads as written in the source,
|
||||
not as a collage. The best passage, plus a second one if it scores close, is kept.
|
||||
5. **No-answer gate.** If the passages cover less than `answer.min_coverage` of the question's idf mass
|
||||
(e.g. "рецепт борща"), there is no answer -> `None`.
|
||||
6. **Rephrase (patient only).** The excerpt may be simplified via `gateway.phrase("rag_answer")`; a
|
||||
rephrasing that adds content (`med_qa.adds_content`) is rejected and the excerpt is shown verbatim.
|
||||
Doctors always get the verbatim text.
|
||||
7. **Framing.** Intro, citations "[КР ПТСР 2023, стр. 28]" and the doctor note from `config/texts/rag.yaml`
|
||||
are added outside the LLM, so they are always present.
|
||||
|
||||
Safety screening of the incoming message is the caller's job (guardrail #2): `answer()` expects text that
|
||||
already passed `app.core.safety` and `privacy.scrub()`. Logs carry only codes and counts.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.config import get_settings
|
||||
from app.core.types import ScrubbedText
|
||||
from app.llm.gateway import FORBIDDEN_OUTPUT, MAX_REPLY_CHARS, get_gateway
|
||||
from app.llm.med_qa import DOSE_QUESTION, DOSING_SENTENCE, adds_content
|
||||
from app.llm.rag.corpus import TABLE_DIVIDER, AnswerConfig, Audience, Chunk, load_config, split_sentences
|
||||
from app.llm.rag.index import BM25Index, get_index, tokens
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
PHRASE_INTENT = "rag_answer"
|
||||
|
||||
LIT_REFS = re.compile(r"\s*\[\d+(?:\s*[-–,]\s*\d+)*\s*,?\]") # "[1, 15, 16]", "[48]", "[1-3]"
|
||||
GRADE = re.compile(r"\s*Уровень убедительности рекомендаций\s+[ABCАВС]\b.*$", re.IGNORECASE)
|
||||
BULLET = re.compile(r"^\s*(?:[-–—●•▪◦*]|\d+[.)])\s+")
|
||||
COMMENT_LABEL = re.compile(r"^Комментари[йи]\s*[:.]\s*")
|
||||
# KR sentences about prescribing are for the doctor: a patient reading "рекомендуется назначение
|
||||
# пароксетина" or "дозировка титруется" out of context takes it as advice (SPEC §7: no drug/dose advice).
|
||||
PRESCRIBING = re.compile(
|
||||
r"назнач|дозиров|титр|мг\b|мг/|доз[аеуыо]\b|дозы\b|/сут|в сутки|раза? в день|разово"
|
||||
r"|\d+\s*(?:[-–]\s*\d+\s*)?таблет",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
# "что такое флэшбек" -> the term; its definition "Флэшбек (flashback – англ.) – ..." is preferred.
|
||||
DEFINITION_QUESTION = re.compile(r"^\s*что (?:такое|значит|означает|это за)\s+(.+?)[?.!\s]*$", re.IGNORECASE)
|
||||
DEFINITION_HEAD = re.compile(r"\s[–—-]\s(?![^(]*\))|\s[–—-]?\s?это\s") # dash outside "(… – англ.)", or "это"
|
||||
DEFINITION_HEAD_CHARS = 60
|
||||
DEFINITION_BONUS = 2.0 # multiplier for a sentence that starts with the asked term (picks it within a chunk)
|
||||
# The gateway's output filter targets the bot's own words. Two of its patterns also hit neutral quotes from
|
||||
# the KR — the definition "ПТСР – это расстройство, которое..." and "диагноз ставится на основании..." —
|
||||
# and would cut exactly the patient-information text. Quotes are checked against the rest of the filter
|
||||
# (dose/drug imperatives, "у вас депрессия", confidentiality promises).
|
||||
QUOTE_EXEMPT_PREFIXES = ("диагноз", "(это|похоже на|признаки)")
|
||||
QUOTE_FILTER = tuple(p for p in FORBIDDEN_OUTPUT if not p.pattern.startswith(QUOTE_EXEMPT_PREFIXES))
|
||||
MIN_SENTENCE_CHARS = 25
|
||||
LEAD_IN_CHARS = 80
|
||||
# A best sentence opening with a reference back ("Это помогает...") needs the sentence it refers to.
|
||||
ANAPHORA = re.compile(
|
||||
r"^(Это|Эти|Этот|Эта|Он|Она|Оно|Они|Такие|Такой|Данн\w+|При этом|Кроме того|Однако|Поэтому)\b"
|
||||
) # a best sentence this short ("Что такое ПТСР?") pulls in the next one as its answer
|
||||
|
||||
|
||||
class Citation(BaseModel):
|
||||
doc: str # document id from sources.yaml
|
||||
short: str # "КР ПТСР 2023"
|
||||
title: str
|
||||
section: str
|
||||
pages: str # "28" | "28–29"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return f"{self.short}, стр. {self.pages}"
|
||||
|
||||
|
||||
class RagAnswer(BaseModel):
|
||||
text: str # framed answer: intro + body + sources (+ doctor note for patients)
|
||||
excerpt: str # the quoted passages, verbatim
|
||||
citations: list[Citation]
|
||||
audience: Audience
|
||||
coverage: float # share of the question covered by the passages (0..1)
|
||||
phrased_by_llm: bool
|
||||
|
||||
|
||||
class AudienceTexts(BaseModel):
|
||||
fallback: str
|
||||
intro: str
|
||||
sources_label: str
|
||||
doctor_note: str = ""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Passage:
|
||||
chunk: Chunk
|
||||
text: str
|
||||
score: float
|
||||
definition: bool = False # starts with the term asked in "что такое X?"
|
||||
|
||||
|
||||
@lru_cache
|
||||
def load_rag_texts(path: Path | None = None) -> dict[str, AudienceTexts]:
|
||||
path = path or get_settings().config_dir / "texts" / "rag.yaml"
|
||||
raw = yaml.safe_load(path.read_text(encoding="utf-8"))
|
||||
return {k: AudienceTexts.model_validate(v) for k, v in raw.items()}
|
||||
|
||||
|
||||
def fallback_text(audience: Audience = "patient") -> str:
|
||||
"""What the caller sends when `answer()` returns `None`."""
|
||||
return load_rag_texts()[audience].fallback
|
||||
|
||||
|
||||
def _clean(sentence: str, *, keep_grade: bool) -> str:
|
||||
s = BULLET.sub("", sentence)
|
||||
s = COMMENT_LABEL.sub("", s)
|
||||
if not keep_grade: # evidence grades matter to doctors, to a patient they are noise
|
||||
s = GRADE.sub("", s)
|
||||
s = LIT_REFS.sub("", s)
|
||||
return " ".join(s.split()).strip(" ;")
|
||||
|
||||
|
||||
def _cells(line: str) -> list[str]:
|
||||
return [c.strip() for c in line.strip().strip("|").split("|")]
|
||||
|
||||
|
||||
def _row_text(header: list[str], row: list[str]) -> str:
|
||||
pairs = [f"{h}: {v}" if h else v for h, v in zip(header, row, strict=False) if v]
|
||||
return LIT_REFS.sub("", "; ".join(pairs))
|
||||
|
||||
|
||||
def chunk_sentences(chunk: Chunk, *, clinical: bool) -> list[str | None]:
|
||||
"""Readable sentences of a chunk in reading order. `clinical` (doctor/team) keeps table rows as
|
||||
"column: value" lines and evidence grades. `None` marks a dropped row: passages do not span it."""
|
||||
out: list[str | None] = []
|
||||
lines = chunk.text.splitlines()
|
||||
header: list[str] | None = None
|
||||
for i, line in enumerate(lines):
|
||||
if line.startswith("|"):
|
||||
if TABLE_DIVIDER.match(line):
|
||||
continue
|
||||
if i + 1 < len(lines) and TABLE_DIVIDER.match(lines[i + 1]):
|
||||
header = _cells(line)
|
||||
continue
|
||||
out.append(_row_text(header or [], _cells(line)) if clinical else None)
|
||||
continue
|
||||
header = None
|
||||
out.extend(s for s in (_clean(p, keep_grade=clinical) for p in split_sentences(line)) if s)
|
||||
return out
|
||||
|
||||
|
||||
def _allowed(sentence: str | None, audience: Audience) -> bool:
|
||||
if sentence is None or len(sentence) < MIN_SENTENCE_CHARS and not sentence.endswith(("?", ":")):
|
||||
return False
|
||||
if sentence[0].islower():
|
||||
return False # tail of a paragraph cut by the chunk boundary: unreadable without its start
|
||||
if audience == "patient":
|
||||
return not (
|
||||
DOSING_SENTENCE.search(sentence)
|
||||
or PRESCRIBING.search(sentence)
|
||||
or any(p.search(sentence) for p in QUOTE_FILTER)
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
def _defined_term(question: str) -> set[str] | None:
|
||||
m = DEFINITION_QUESTION.match(question)
|
||||
return set(tokens(m.group(1))) or None if m else None
|
||||
|
||||
|
||||
def _is_definition(sentence: str, defined: set[str]) -> bool:
|
||||
""" "Флэшбек (flashback – англ.) – возврат..." / "ПТСР – это расстройство...": the term, then a dash."""
|
||||
head = DEFINITION_HEAD.split(sentence, maxsplit=1)
|
||||
return len(head) == 2 and len(head[0]) <= DEFINITION_HEAD_CHARS and bool(set(tokens(head[0])) & defined)
|
||||
|
||||
|
||||
def _sentence_score(index: BM25Index, sentence: str, terms: dict[str, float], defined: set[str] | None) -> float:
|
||||
score = index.overlap(sentence, terms)
|
||||
if defined and score > 0 and _is_definition(sentence, defined):
|
||||
score *= DEFINITION_BONUS
|
||||
return score
|
||||
|
||||
|
||||
def best_passage(
|
||||
index: BM25Index,
|
||||
chunk: Chunk,
|
||||
terms: dict[str, float],
|
||||
audience: Audience,
|
||||
cfg: AnswerConfig,
|
||||
defined: set[str] | None = None,
|
||||
) -> Passage | None:
|
||||
"""Best sentence of the chunk extended by adjacent relevant sentences (contiguous, within limits)."""
|
||||
sents = [s if _allowed(s, audience) else None for s in chunk_sentences(chunk, clinical=audience != "patient")]
|
||||
scores = [_sentence_score(index, s, terms, defined) if s else 0.0 for s in sents]
|
||||
if not scores or max(scores) <= 0:
|
||||
return None
|
||||
best = max(range(len(sents)), key=lambda i: (scores[i], -i))
|
||||
head = sents[best]
|
||||
assert head is not None # score > 0 implies an allowed sentence
|
||||
lo = hi = best
|
||||
size = len(head)
|
||||
|
||||
def fits(j: int) -> bool:
|
||||
if not 0 <= j < len(sents):
|
||||
return False
|
||||
s = sents[j]
|
||||
return s is not None and size + len(s) + 1 <= cfg.max_chars
|
||||
|
||||
# A short lead-in (question-like heading, "Комментарии:"-style intro) needs the sentence after it.
|
||||
def length(j: int) -> int:
|
||||
return len(sents[j] or "") + 1
|
||||
|
||||
if (size < LEAD_IN_CHARS or head.endswith(("?", ":"))) and fits(best + 1):
|
||||
hi += 1
|
||||
size += length(hi)
|
||||
if ANAPHORA.match(head) and fits(best - 1):
|
||||
lo -= 1
|
||||
size += length(lo)
|
||||
floor = cfg.neighbor_ratio * scores[best]
|
||||
while hi - lo + 1 < cfg.max_sentences:
|
||||
nxt = next((j for j in (hi + 1, lo - 1) if fits(j) and scores[j] > 0 and scores[j] >= floor), None)
|
||||
if nxt is None:
|
||||
break
|
||||
lo, hi = min(lo, nxt), max(hi, nxt)
|
||||
size += length(nxt)
|
||||
text = " ".join(s for s in sents[lo : hi + 1] if s)
|
||||
definition = scores[best] > index.overlap(head, terms) # got DEFINITION_BONUS
|
||||
return Passage(chunk, text, index.overlap(text, terms) * index.boost(chunk, audience), definition)
|
||||
|
||||
|
||||
def select(index: BM25Index, question: str, audience: Audience, docs: set[str] | None = None) -> list[Passage]:
|
||||
"""Best passage, plus a second one from another chunk if it scores within `second_passage_ratio`."""
|
||||
cfg = load_config().answer
|
||||
hits = index.search(question, audience=audience, docs=docs, top_k=cfg.max_chunks)
|
||||
terms = index.query_terms(question)
|
||||
defined = _defined_term(question)
|
||||
passages = [p for h in hits if (p := best_passage(index, h.chunk, terms, audience, cfg, defined))]
|
||||
if not passages:
|
||||
return []
|
||||
# A definition answers "что такое X?" better than any passage that merely mentions X often.
|
||||
passages.sort(key=lambda p: (not p.definition, -p.score)) # stable: retrieval rank breaks ties
|
||||
top = passages[0]
|
||||
chosen = [top]
|
||||
for p in passages[1:]:
|
||||
if len(chosen) >= cfg.max_passages:
|
||||
break
|
||||
if p.score >= cfg.second_passage_ratio * top.score and p.text != top.text:
|
||||
chosen.append(p)
|
||||
return chosen
|
||||
|
||||
|
||||
def _citation(chunk: Chunk) -> Citation:
|
||||
return Citation(
|
||||
doc=chunk.doc.id, short=chunk.doc.short, title=chunk.doc.title, section=chunk.section, pages=chunk.pages
|
||||
)
|
||||
|
||||
|
||||
async def _rephrase(excerpt: str, question: str) -> tuple[str, bool]:
|
||||
if len(excerpt) > MAX_REPLY_CHARS or any(p.search(excerpt) for p in FORBIDDEN_OUTPUT):
|
||||
return excerpt, False # the gateway would reject any reply this long / with these words anyway
|
||||
text = await get_gateway().phrase(PHRASE_INTENT, {"excerpt": excerpt, "question": question})
|
||||
if " ".join(text.split()) == " ".join(excerpt.split()):
|
||||
return excerpt, False
|
||||
if adds_content(text, excerpt, question):
|
||||
log.warning("rag rephrasing rejected: adds content")
|
||||
return excerpt, False
|
||||
return text, True
|
||||
|
||||
|
||||
async def answer(
|
||||
question: ScrubbedText,
|
||||
*,
|
||||
audience: Audience = "patient",
|
||||
docs: set[str] | None = None,
|
||||
index: BM25Index | None = None,
|
||||
) -> RagAnswer | None:
|
||||
"""Answer strictly from the knowledge base, with citations; `None` -> caller sends `fallback_text()`."""
|
||||
if audience == "patient" and DOSE_QUESTION.search(question.lower()):
|
||||
log.info("rag: dose question redirected to doctor")
|
||||
return None
|
||||
index = index or get_index()
|
||||
passages = select(index, question, audience, docs)
|
||||
if not passages:
|
||||
log.info("rag: no relevant fragment audience=%s", audience)
|
||||
return None
|
||||
coverage = index.coverage(" ".join(p.text for p in passages), question)
|
||||
if coverage < load_config().answer.min_coverage:
|
||||
log.info("rag: low coverage audience=%s coverage=%.2f", audience, coverage)
|
||||
return None
|
||||
citations = [_citation(p.chunk) for p in passages]
|
||||
texts = load_rag_texts()[audience]
|
||||
|
||||
if audience == "patient":
|
||||
excerpt = " ".join(p.text for p in passages)
|
||||
body, by_llm = await _rephrase(excerpt, question)
|
||||
sources = "; ".join(dict.fromkeys(c.label for c in citations))
|
||||
text = f"{texts.intro} {body}\n\n{texts.sources_label}: {sources}."
|
||||
else: # doctor/team: verbatim quotes, each with its own citation
|
||||
excerpt = "\n".join(f"{p.text} [{c.label}]" for p, c in zip(passages, citations, strict=True))
|
||||
by_llm = False
|
||||
text = f"{texts.intro}\n{excerpt}"
|
||||
if texts.doctor_note:
|
||||
text += f"\n\n{texts.doctor_note}"
|
||||
log.info("rag: answered audience=%s passages=%d coverage=%.2f llm=%s", audience, len(passages), coverage, by_llm)
|
||||
return RagAnswer(
|
||||
text=text, excerpt=excerpt, citations=citations, audience=audience, coverage=coverage, phrased_by_llm=by_llm
|
||||
)
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Russian Snowball stemmer (https://snowballstem.org/algorithms/russian/stemmer.html).
|
||||
|
||||
Self-contained because `pyproject.toml` is a frozen contract (#34): no new dependencies for the RAG.
|
||||
Latin words and digits pass through lowercased; "ё" is folded to "е".
|
||||
"""
|
||||
|
||||
import re
|
||||
from functools import lru_cache
|
||||
|
||||
_VOWELS = set("аеиоуыэюя")
|
||||
|
||||
# (suffixes, needs_a_ya): group-1 suffixes count only when preceded by "а"/"я" (which stays in the stem).
|
||||
_PERFECTIVE_GERUND = ((("вшись", "вши", "в"), True), (("ившись", "ывшись", "ивши", "ывши", "ив", "ыв"), False))
|
||||
_REFLEXIVE = ((("ся", "сь"), False),)
|
||||
_ADJECTIVE = (
|
||||
(
|
||||
(
|
||||
"ими", "ыми", "его", "ого", "ему", "ому", "ее", "ие", "ые", "ое", "ей", "ий", "ый", "ой", "ем",
|
||||
"им", "ым", "ом", "их", "ых", "ую", "юю", "ая", "яя", "ою", "ею",
|
||||
),
|
||||
False,
|
||||
),
|
||||
) # fmt: skip
|
||||
_PARTICIPLE = ((("ем", "нн", "вш", "ющ", "щ"), True), (("ивш", "ывш", "ующ"), False))
|
||||
_VERB = (
|
||||
(("ла", "на", "ете", "йте", "ли", "й", "л", "ем", "н", "ло", "но", "ет", "ют", "ны", "ть", "ешь", "нно"), True),
|
||||
(
|
||||
(
|
||||
"ила", "ыла", "ена", "ейте", "уйте", "ите", "или", "ыли", "ей", "уй", "ил", "ыл", "им", "ым", "ен",
|
||||
"ило", "ыло", "ено", "ят", "ует", "уют", "ит", "ыт", "ены", "ить", "ыть", "ишь", "ую", "ю",
|
||||
),
|
||||
False,
|
||||
),
|
||||
) # fmt: skip
|
||||
_NOUN = (
|
||||
(
|
||||
(
|
||||
"иями", "ями", "ами", "ией", "иям", "ием", "иях", "ев", "ов", "ие", "ье", "еи", "ии", "ей", "ой",
|
||||
"ий", "ям", "ем", "ам", "ом", "ах", "ях", "ию", "ью", "ия", "ья", "а", "е", "и", "й", "о", "у",
|
||||
"ы", "ь", "ю", "я",
|
||||
),
|
||||
False,
|
||||
),
|
||||
) # fmt: skip
|
||||
_DERIVATIONAL = ((("ость", "ост"), False),)
|
||||
_SUPERLATIVE = ((("ейше", "ейш"), False),)
|
||||
|
||||
_CYRILLIC_WORD = re.compile(r"^[а-я]+$")
|
||||
|
||||
|
||||
def _strip(word: str, start: int, groups: tuple[tuple[tuple[str, ...], bool], ...]) -> str | None:
|
||||
"""Remove the longest matching suffix lying in word[start:], or `None` if none matches."""
|
||||
best: tuple[int, str] | None = None
|
||||
for suffixes, needs_a_ya in groups:
|
||||
for suffix in suffixes:
|
||||
if not word.endswith(suffix) or len(word) - len(suffix) < start:
|
||||
continue
|
||||
if needs_a_ya and (len(word) - len(suffix) - 1 < start or word[-len(suffix) - 1] not in "ая"):
|
||||
continue
|
||||
if best is None or len(suffix) > best[0]:
|
||||
best = (len(suffix), word[: -len(suffix)])
|
||||
return None if best is None else best[1]
|
||||
|
||||
|
||||
def _regions(word: str) -> tuple[int, int]:
|
||||
"""RV (after the first vowel) and R2 start indices."""
|
||||
rv = next((i + 1 for i, ch in enumerate(word) if ch in _VOWELS), len(word))
|
||||
|
||||
def after_vc(start: int) -> int:
|
||||
for i in range(start + 1, len(word)):
|
||||
if word[i] not in _VOWELS and word[i - 1] in _VOWELS:
|
||||
return i + 1
|
||||
return len(word)
|
||||
|
||||
return rv, after_vc(after_vc(0))
|
||||
|
||||
|
||||
@lru_cache(maxsize=65536)
|
||||
def stem(word: str) -> str:
|
||||
word = word.lower().replace("ё", "е")
|
||||
if not _CYRILLIC_WORD.match(word):
|
||||
return word
|
||||
rv, r2 = _regions(word)
|
||||
|
||||
# Step 1
|
||||
out = _strip(word, rv, _PERFECTIVE_GERUND)
|
||||
if out is None:
|
||||
word = _strip(word, rv, _REFLEXIVE) or word
|
||||
out = _strip(word, rv, _ADJECTIVE)
|
||||
if out is not None:
|
||||
out = _strip(out, rv, _PARTICIPLE) or out
|
||||
else:
|
||||
out = _strip(word, rv, _VERB)
|
||||
if out is None:
|
||||
out = _strip(word, rv, _NOUN)
|
||||
word = out if out is not None else word
|
||||
|
||||
# Step 2
|
||||
if word.endswith("и") and len(word) - 1 >= rv:
|
||||
word = word[:-1]
|
||||
# Step 3
|
||||
word = _strip(word, r2, _DERIVATIONAL) or word
|
||||
# Step 4
|
||||
if word.endswith("нн") and len(word) - 1 >= rv:
|
||||
return word[:-1]
|
||||
sup = _strip(word, rv, _SUPERLATIVE)
|
||||
if sup is not None:
|
||||
word = sup
|
||||
return word[:-1] if word.endswith("нн") else word
|
||||
if word.endswith("ь") and len(word) - 1 >= rv:
|
||||
word = word[:-1]
|
||||
return word
|
||||
Reference in New Issue
Block a user