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:
@@ -0,0 +1,278 @@
|
||||
"""PDF -> TXT для RAG.
|
||||
|
||||
Запуск из корня репо (нужен poppler-utils для pdftotext):
|
||||
uv run --with pdfplumber python docs/source/rag/pdf_to_txt.py docs/source/rag/*.pdf --out config/rag/corpus
|
||||
Сверка результата с PDF:
|
||||
python docs/source/rag/check_txt.py docs/source/rag/*.pdf --out config/rag/corpus
|
||||
|
||||
- текст вне таблиц, строки склеены в абзацы (в том числе через границу страницы);
|
||||
- таблицы -> markdown; у таблицы-продолжения на новой странице повторяется шапка,
|
||||
строка, разорванная границей страницы, склеивается;
|
||||
- колонтитулы (повторяющиеся строки вверху/внизу страниц) вырезаются;
|
||||
- маркеры страниц `--- стр. N ---` для ссылок на источник;
|
||||
- текст с картинок (блок-схемы и т.п.) берётся из figures/<имя pdf>.p<N>.txt
|
||||
(ручная транскрипция) и вставляется в конец своей страницы;
|
||||
- если есть manual/<имя pdf>.txt, берётся он (документы, которые автоматом
|
||||
не разбираются, например вёрстка Chromium с выносками).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pdfplumber
|
||||
|
||||
FIGURES_DIR = Path(__file__).parent / "figures"
|
||||
MANUAL_DIR = Path(__file__).parent / "manual" # вручную выверенные txt: берутся как есть
|
||||
EDGE_LINES = 4 # сколько строк сверху/снизу страницы проверять на колонтитул
|
||||
SENTENCE_END = tuple(".:;!?»)")
|
||||
GRADE_TAIL = re.compile(r"^([ABCАВС]|\(уровень\b.*)$") # "Уровень убедительности рекомендаций\nA\n(уровень ...)"
|
||||
BULLET = re.compile(r"^([●•▪◦\-–—*]|\d+[.)]|[а-яa-z][.)]|[IVXА-Я]{1,4}[.)])\s")
|
||||
HYPHEN = "\x00" # временная метка переноса "сло-\nво", снимается в resolve_hyphens
|
||||
SOFT = "\x01" # временная метка разрыва строки внутри слова в узкой ячейке: "Соверше\nнно"
|
||||
DEFINITION = re.compile(r"^\S+ [–—] ") # "рТМС – ..." в списке сокращений
|
||||
ITEM = re.compile(r"^(\d+\s*[.)\-–]|[А-ЯA-Z]$)") # "1. СНИЖЕННОЕ НАСТРОЕНИЕ", "И" — пункт, не шапка
|
||||
OPTION = re.compile(r"^0 □") # варианты ответа пункта шкалы: "0 □ Отсутствует 1 □ ..."
|
||||
|
||||
|
||||
def norm_edge(line: str) -> str:
|
||||
return re.sub(r"\d+", "#", line.strip())
|
||||
|
||||
|
||||
def strip_junk(line: str, junk: set[str]) -> str:
|
||||
"""Убирает колонтитул: целую строку, удлинённый вариант (на альбомных
|
||||
страницах он шире) и прилипший к концу строки ("...на 7Страница 24 из 117")."""
|
||||
n = norm_edge(line)
|
||||
if n in junk or any(len(j) >= 30 and (n.startswith(j) or j.startswith(n) and len(n) >= 30) for j in junk):
|
||||
return ""
|
||||
for j in junk:
|
||||
if len(j) >= 10 and "#" in j:
|
||||
line = re.sub(re.escape(j).replace("\\#", r"\d+") + r"\s*$", "", line)
|
||||
return line
|
||||
|
||||
|
||||
def clean_cell(cell: str | None) -> str:
|
||||
if not cell:
|
||||
return ""
|
||||
cell = re.sub(r"(\w)-\s*\n\s*(?=[а-яёa-z])", r"\1" + HYPHEN, cell) # перенос в узкой ячейке
|
||||
cell = re.sub(r"([а-яё])\n(?=[а-яё])", r"\1" + SOFT, cell)
|
||||
return re.sub(r"\s*\n\s*", " ", cell).replace("|", "/").strip()
|
||||
|
||||
|
||||
def is_header(row: list[str]) -> bool:
|
||||
"""Шапка — короткие подписи колонок; у таблиц-опросников её нет и первая
|
||||
строка — уже пункт шкалы (его нельзя повторять на страницах-продолжениях)."""
|
||||
return not ITEM.match(row[0]) and all(len(c) <= 100 for c in row)
|
||||
|
||||
|
||||
def merge_split_rows(rows: list[list[str]]) -> list[list[str]]:
|
||||
"""Пункт шкалы, разрезанный на "N. НАЗВАНИЕ" и строку с вариантами "0 □ ...",
|
||||
склеивается обратно в одну строку."""
|
||||
out: list[list[str]] = []
|
||||
for row in rows:
|
||||
if out and OPTION.match(row[0]) and "□" not in out[-1][0]:
|
||||
out[-1] = [f"{a} {c}".strip() for a, c in zip(out[-1], row)]
|
||||
else:
|
||||
out.append(row)
|
||||
return out
|
||||
|
||||
|
||||
def table_to_md(header: list[str], body: list[list[str]]) -> str:
|
||||
out = ["| " + " | ".join(header) + " |", "|" + "---|" * len(header)]
|
||||
out += ["| " + " | ".join(r) + " |" for r in body]
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def reflow(lines: list[str]) -> str:
|
||||
"""Склеивает строки, разорванные вёрсткой, обратно в абзацы."""
|
||||
paras: list[str] = []
|
||||
for raw in lines:
|
||||
line = raw.strip()
|
||||
if not line:
|
||||
continue
|
||||
prev = paras[-1] if paras else ""
|
||||
if prev and GRADE_TAIL.match(line):
|
||||
paras[-1] = prev + " " + line
|
||||
continue
|
||||
joinable = (
|
||||
prev
|
||||
and line[0].islower()
|
||||
and not BULLET.match(line)
|
||||
and not DEFINITION.match(line)
|
||||
)
|
||||
if joinable and prev.endswith("-") and not prev.endswith(" -"):
|
||||
paras[-1] = prev[:-1] + HYPHEN + line # перенос или составное слово
|
||||
elif joinable:
|
||||
paras[-1] = prev + " " + line
|
||||
else:
|
||||
paras.append(line)
|
||||
return "\n".join(paras)
|
||||
|
||||
|
||||
def region_text(pdf_path: Path, num: int, page: Any, top: float, bottom: float) -> str:
|
||||
"""Текст горизонтальной полосы страницы. pdftotext надёжнее pdfplumber на
|
||||
буквах с диакритикой (у pdfplumber выходит "K o z a r i ć")."""
|
||||
y = math.floor(top)
|
||||
h = math.ceil(bottom) - y
|
||||
if h <= 0:
|
||||
return ""
|
||||
cmd = ["pdftotext", "-nopgbrk", "-f", str(num), "-l", str(num), "-x", "0", "-y", str(y),
|
||||
"-W", str(math.ceil(page.width)), "-H", str(h), str(pdf_path), "-"]
|
||||
return subprocess.run(cmd, capture_output=True, text=True, check=True).stdout
|
||||
|
||||
|
||||
def word_bag(text: str) -> Counter[str]:
|
||||
return Counter(w.lower() for w in re.findall(r"\w+", text.replace(HYPHEN, "")))
|
||||
|
||||
|
||||
def page_blocks(pdf_path: Path, num: int, page: Any) -> list[tuple[str, Any]]:
|
||||
"""Текстовые полосы и таблицы страницы сверху вниз."""
|
||||
# у таблицы, разорванной страницей, нет нижней (или верхней) рамки — без неё
|
||||
# pdfplumber теряет крайнюю строку; дорисовываем линии по концам вертикалей
|
||||
ends = {round(y) for e in page.vertical_edges if e["height"] > 5 for y in (e["top"], e["bottom"])}
|
||||
settings = {"explicit_horizontal_lines": sorted(ends)}
|
||||
tables = []
|
||||
for t in sorted(page.find_tables(settings), key=lambda t: t.bbox[1]):
|
||||
rows = [[clean_cell(c) for c in row] for row in t.extract()]
|
||||
# рамка, внутри которой текст лежит не в ячейках, — не таблица: иначе текст пропадёт
|
||||
band = word_bag(region_text(pdf_path, num, page, t.bbox[1], t.bbox[3]))
|
||||
if band - word_bag(" ".join(" ".join(r) for r in rows)):
|
||||
continue
|
||||
tables.append((t.bbox[1], t.bbox[3], rows))
|
||||
blocks: list[tuple[str, Any]] = []
|
||||
top = 0.0
|
||||
for t_top, t_bottom, rows in tables:
|
||||
if t_top < top: # вложенная/перекрывающаяся таблица
|
||||
continue
|
||||
blocks.append(("text", region_text(pdf_path, num, page, top, t_top)))
|
||||
blocks.append(("table", rows))
|
||||
top = t_bottom
|
||||
blocks.append(("text", region_text(pdf_path, num, page, top, page.height)))
|
||||
return blocks
|
||||
|
||||
|
||||
def resolve_hyphens(text: str) -> str:
|
||||
""""клинико-\nпсихопатологические" — составное слово, "антихо-\nлинерги-\nческие" —
|
||||
перенос. Склеиваем без дефиса, если склейка встречается в документе целиком или
|
||||
хотя бы одна половинка — не слово (обрывок слога); иначе оставляем дефис."""
|
||||
vocab = {w.lower() for w in re.findall(r"\w+", re.sub(rf"\w*[{HYPHEN}{SOFT}]\w*", " ", text))}
|
||||
|
||||
def fix(m: re.Match[str]) -> str:
|
||||
a, b = m.group(1), m.group(2)
|
||||
whole = (a + b).lower() in vocab
|
||||
compound = a.lower() in vocab and b.lower() in vocab
|
||||
return a + b if whole or not compound else f"{a}-{b}"
|
||||
|
||||
pattern = re.compile(rf"(\w+){HYPHEN}(\w+)")
|
||||
while HYPHEN in text:
|
||||
new = pattern.sub(fix, text, count=0)
|
||||
if new == text:
|
||||
break
|
||||
text = new
|
||||
text = text.replace(HYPHEN, "-")
|
||||
|
||||
def fix_soft(m: re.Match[str]) -> str:
|
||||
a, b = m.group(1), m.group(2)
|
||||
broken = (a + b).lower() in vocab and not (a.lower() in vocab and b.lower() in vocab)
|
||||
return a + b if broken else f"{a} {b}"
|
||||
|
||||
return re.sub(rf"(\w+){SOFT}(\w+)", fix_soft, text).replace(SOFT, " ")
|
||||
|
||||
|
||||
def convert(pdf_path: Path) -> str:
|
||||
manual = MANUAL_DIR / f"{pdf_path.stem}.txt"
|
||||
if manual.exists():
|
||||
return manual.read_text()
|
||||
with pdfplumber.open(pdf_path) as pdf:
|
||||
pages = [page_blocks(pdf_path, n, p) for n, p in enumerate(pdf.pages, 1)]
|
||||
|
||||
# колонтитулы: строки у края страницы, повторяющиеся на >= 40% страниц
|
||||
edge: Counter[str] = Counter()
|
||||
for page in pages:
|
||||
lines = [l for kind, b in page if kind == "text" for l in b.splitlines() if l.strip()]
|
||||
edge.update({norm_edge(l) for l in lines[:EDGE_LINES] + lines[-EDGE_LINES:]})
|
||||
junk = {k for k, n in edge.items() if n >= max(3, 0.4 * len(pages))}
|
||||
|
||||
# blocks: ["page", N] | ["text", str] | ["table", (header, body)] | ["figure", str]
|
||||
blocks: list[list[Any]] = []
|
||||
for num, page in enumerate(pages, 1):
|
||||
blocks.append(["page", num])
|
||||
first_on_page = True
|
||||
for kind, b in page:
|
||||
prev = next((x for x in reversed(blocks) if x[0] != "page"), None)
|
||||
if kind == "text":
|
||||
text = reflow([strip_junk(l, junk) for l in b.splitlines()])
|
||||
if not text:
|
||||
continue
|
||||
if first_on_page and prev and prev[0] == "text" and text[0].islower():
|
||||
# абзац перетёк на новую страницу: дописываем его хвост к началу
|
||||
head, _, rest = text.partition("\n")
|
||||
prev[1] = prev[1] + " " + head
|
||||
if rest:
|
||||
blocks.append(["text", rest])
|
||||
else:
|
||||
blocks.append(["text", text])
|
||||
else:
|
||||
rows = [r for r in b if any(r)]
|
||||
if not rows:
|
||||
continue
|
||||
continues = (
|
||||
first_on_page
|
||||
and prev is not None
|
||||
and prev[0] == "table"
|
||||
and len(prev[1][0]) == len(rows[0])
|
||||
and rows[0] != prev[1][0]
|
||||
)
|
||||
if continues:
|
||||
header, body = prev[1]
|
||||
tail = body if body else [header]
|
||||
if not rows[0][0] or rows[0][0][:1].islower() or (
|
||||
OPTION.match(rows[0][0]) and "□" not in tail[-1][0]
|
||||
):
|
||||
# строка таблицы разорвана границей страницы
|
||||
tail[-1] = [f"{a} {c}".strip() for a, c in zip(tail[-1], rows[0])]
|
||||
if not body:
|
||||
prev[1] = (tail[0], body)
|
||||
rows = rows[1:]
|
||||
if rows and is_header(header):
|
||||
blocks.append(["table", (header, rows)])
|
||||
elif rows: # шапки нет — не дублируем первый пункт
|
||||
blocks.append(["table", (rows[0], rows[1:])])
|
||||
else:
|
||||
blocks.append(["table", (rows[0], rows[1:])])
|
||||
first_on_page = False
|
||||
fig = FIGURES_DIR / f"{pdf_path.stem}.p{num}.txt"
|
||||
if fig.exists():
|
||||
blocks.append(["figure", fig.read_text().strip()])
|
||||
|
||||
out = [f"# {pdf_path.stem}", f"Источник: {pdf_path.name}, страниц: {len(pages)}"]
|
||||
for kind, b in blocks:
|
||||
if kind == "page":
|
||||
out.append(f"--- стр. {b} ---")
|
||||
elif kind == "text":
|
||||
out.append(b)
|
||||
elif kind == "table":
|
||||
header, *body = merge_split_rows([b[0], *b[1]])
|
||||
out.append(table_to_md(header, body))
|
||||
else:
|
||||
out.append("[Изображение: текст перенесён вручную]\n" + b)
|
||||
return resolve_hyphens("\n\n".join(out)) + "\n"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("pdfs", nargs="+", type=Path)
|
||||
parser.add_argument("--out", type=Path, help="папка для txt (по умолчанию рядом с pdf)")
|
||||
args = parser.parse_args()
|
||||
for src in args.pdfs:
|
||||
dst = (args.out / src.name if args.out else src).with_suffix(".txt")
|
||||
dst.write_text(convert(src))
|
||||
print(f"{src.name} -> {dst}")
|
||||
Reference in New Issue
Block a user