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>
281 lines
14 KiB
Python
281 lines
14 KiB
Python
"""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 ---` для ссылок на источник; абзац, перетёкший на новую
|
||
страницу, склеивается, а место перехода помечается `⟦стр. 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():
|
||
# абзац перетёк на новую страницу: дописываем его хвост к началу, а место
|
||
# перехода помечаем ⟦стр. N⟧ — чанкер по метке ставит верную страницу в цитату
|
||
head, _, rest = text.partition("\n")
|
||
prev[1] = prev[1] + f" ⟦стр. {num}⟧ " + 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}")
|