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>
351 lines
13 KiB
Python
351 lines
13 KiB
Python
"""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]"; inline
|
||
"⟦стр. N⟧" markers (paragraph glued across a page break) are removed from the text but move the range.
|
||
"""
|
||
|
||
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+) ---$")
|
||
# A paragraph that flows onto the next page is glued by the converter; "⟦стр. N⟧" marks where page N begins.
|
||
INLINE_PAGE = re.compile(r"\s*⟦стр\. (\d+)⟧\s*")
|
||
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)
|
||
# doc id -> regexes naming its disease; a question naming one disease favours that document
|
||
doc_cues: dict[str, list[str]] = Field(default_factory=dict)
|
||
doc_cue_boost: float = 1.0
|
||
|
||
|
||
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 # page where the unit starts
|
||
table_header: str | None = None # "| a | b |\n|---|---|" for table rows
|
||
breaks: list[tuple[int, int]] = field(default_factory=list) # (offset in text, page starting there)
|
||
|
||
@property
|
||
def page_end(self) -> int:
|
||
return self.breaks[-1][1] if self.breaks else self.page
|
||
|
||
def page_at(self, offset: int) -> int:
|
||
page = self.page
|
||
for start, p in self.breaks:
|
||
if start <= offset:
|
||
page = p
|
||
return page
|
||
|
||
def append(self, line: str, page: int, sep: str) -> None:
|
||
"""Glue a continuation line; record a page break if it comes from a later page."""
|
||
base = len(self.text) + len(sep)
|
||
text, breaks = _strip_inline_pages(line)
|
||
if page > self.page_end:
|
||
self.breaks.append((base, page))
|
||
self.breaks.extend((base + off, p) for off, p in breaks)
|
||
self.text += sep + text
|
||
|
||
|
||
def _strip_inline_pages(line: str) -> tuple[str, list[tuple[int, int]]]:
|
||
"""Text without "⟦стр. N⟧" markers + (offset, page) of each marker in the cleaned text."""
|
||
out, breaks, pos = [], [], 0
|
||
for m in INLINE_PAGE.finditer(line):
|
||
out.append(line[pos : m.start()])
|
||
breaks.append((sum(len(x) for x in out) + 1, int(m.group(1))))
|
||
out.append(" ")
|
||
pos = m.end()
|
||
out.append(line[pos:])
|
||
return "".join(out), breaks
|
||
|
||
|
||
def _new_unit(line: str, page: int) -> _Unit:
|
||
text, breaks = _strip_inline_pages(line)
|
||
return _Unit(text, page, breaks=breaks)
|
||
|
||
|
||
@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(_new_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].append(line, page, "\n")
|
||
elif line[0].islower() and units and units[-1].table_header is None:
|
||
units[-1].append(line, page, " ") # paragraph continued after a page break or a figure
|
||
else:
|
||
units.append(_new_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)
|
||
out: list[_Unit] = []
|
||
pos = 0
|
||
for part in parts:
|
||
start = unit.text.find(part[:40], pos)
|
||
start = pos if start < 0 else start
|
||
end = start + len(part) - 1
|
||
first, last = unit.page_at(start), unit.page_at(end)
|
||
out.append(_Unit(part, first, breaks=[(len(part) - 1, last)] if last != first else []))
|
||
pos = start + len(part)
|
||
return out
|
||
|
||
|
||
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=max(u.page_end for u in group),
|
||
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)
|