Денис Кривоченко
836 lines
30 KiB
Python
836 lines
30 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Unified Document Parser — объединённый парсер DOCX и PDF.
|
||
|
||
Поддерживает два режима:
|
||
1. API mode (по умолчанию) — Unstructured API server (smart_indexer / parse_pdf)
|
||
2. Local mode — локальная библиотека unstructured (docx_pipeline)
|
||
|
||
Возвращает единый список чанков с метаданными, очищенный от шума
|
||
(шапки, футеры, листы согласования, подписи, страницы).
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import re
|
||
import json
|
||
import hashlib
|
||
import uuid
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import List, Dict, Optional, Any
|
||
|
||
|
||
# ============================================================
|
||
# CONFIG
|
||
# ============================================================
|
||
|
||
DEFAULT_API_URL = "http://192.168.1.103:8005/general/v0/general"
|
||
DEFAULT_LOCAL_MODE = False
|
||
|
||
|
||
def get_config() -> dict:
|
||
"""Собирает конфиг из переменных окружения или дефолтов."""
|
||
return {
|
||
"api_url": os.environ.get("UNSTRUCTURED_API_URL", DEFAULT_API_URL),
|
||
"local_mode": os.environ.get("LOCAL_MODE", "0").lower() in ("1", "true", "yes"),
|
||
"languages": os.environ.get("LANGUAGES", "rus").split(","),
|
||
"strategy": os.environ.get("STRATEGY", "hi_res"),
|
||
"pdf_infer_table_structure": os.environ.get("PDF_INFER_TABLE", "true").lower() == "true",
|
||
"min_text_length": int(os.environ.get("MIN_TEXT_LENGTH", "15")),
|
||
"min_chunk_length": int(os.environ.get("MIN_CHUNK_LENGTH", "40")),
|
||
"chunk_max_chars": int(os.environ.get("CHUNK_MAX_CHARS", "1800")),
|
||
"chunk_combine_under": int(os.environ.get("CHUNK_COMBINE_UNDER", "300")),
|
||
"chunk_new_after": int(os.environ.get("CHUNK_NEW_AFTER", "6000")),
|
||
"batch_size": int(os.environ.get("BATCH_SIZE", "32")),
|
||
}
|
||
|
||
|
||
# ============================================================
|
||
# CLEAN / FILTER HELPERS
|
||
# ============================================================
|
||
|
||
def clean_text(text: str) -> str:
|
||
"""Убирает множественные пробелы/переносы."""
|
||
if not text:
|
||
return ""
|
||
text = re.sub(r"\s+", " ", text)
|
||
return text.strip()
|
||
|
||
|
||
|
||
|
||
# ============================================================
|
||
# SECTION DETECTION & ID GENERATION
|
||
# ============================================================
|
||
|
||
# Regex for section titles like "1. Нормативные ссылки", "2. Термины и определения"
|
||
_SECTION_RE = re.compile(r"^\s*(\d+)[\.\s]\s*(.+)$")
|
||
|
||
# Keywords that indicate section titles from Unstructured API
|
||
SECTION_KEYWORDS = [
|
||
"Нормативные ссылки",
|
||
"Термины",
|
||
"Термины и определения",
|
||
"Области применения",
|
||
"Общие положения",
|
||
"Порядок организации",
|
||
"Обязанности",
|
||
"Контроль",
|
||
"Ответственность",
|
||
"Заключительные положения",
|
||
"Приложения",
|
||
]
|
||
|
||
def _detect_section_number(text: str) -> tuple:
|
||
"""Extract section number and title from text.
|
||
|
||
Returns (section_number, section_title) or (None, None).
|
||
"""
|
||
if not text:
|
||
return None, None
|
||
|
||
# Try numbered section pattern: "3. Нормативные ссылки"
|
||
m = _SECTION_RE.match(text.strip())
|
||
if m:
|
||
num = int(m.group(1))
|
||
title = m.group(2).strip()
|
||
return num, title
|
||
|
||
# Try unnumbered section pattern (fallback)
|
||
for kw in SECTION_KEYWORDS:
|
||
if kw in text:
|
||
# Extract just the keyword part as title
|
||
idx = text.index(kw)
|
||
title = text[idx:].strip()
|
||
return None, title
|
||
|
||
return None, None
|
||
|
||
|
||
def _generate_section_id(section_number: int, section_title: str) -> str:
|
||
"""Generate a stable section_id based on section number and title."""
|
||
if section_number is not None:
|
||
return f"section_{section_number}"
|
||
# Fallback: hash the title
|
||
h = hashlib.md5(section_title.encode("utf-8")).hexdigest()[:8]
|
||
return f"section_{h}"
|
||
|
||
|
||
def is_page_number(text: str) -> bool:
|
||
return bool(re.match(r"^Стр\.\s*\d+\s*из\s*\d+$", text))
|
||
|
||
|
||
def is_header_noise(text: str) -> bool:
|
||
patterns = [
|
||
"АО «ХК «Сибцем»",
|
||
"Тип документа:",
|
||
"Ведущее подразделение:",
|
||
"Дата утверждения:",
|
||
"Редакция 1",
|
||
"Оглавление",
|
||
]
|
||
return any(p in text for p in patterns)
|
||
|
||
|
||
def is_header_chunk(text: str) -> bool:
|
||
"""DOCX-специфичные фильтры шапки/обложки."""
|
||
text_lower = text.lower()
|
||
if any(kw in text_lower for kw in [
|
||
"утверждаю", "президент", "генеральный директор",
|
||
"согласован", "в.в. шаповалов",
|
||
]):
|
||
return True
|
||
if any(kw in text_lower for kw in ["кемерово", "2008г.", "рег-20-1", "регистр"]):
|
||
if len(text) < 300:
|
||
return True
|
||
if "лист соглас" in text_lower:
|
||
return True
|
||
if "содержание" in text_lower and len(text) < 500:
|
||
return True
|
||
if any(kw in text_lower for kw in [
|
||
"разработчик:", "дата разработки:", "стр. из:",
|
||
]):
|
||
return True
|
||
if any(kw in text_lower for kw in [
|
||
"список лиц", "разработавших", "финансовый директор",
|
||
"главный инженер",
|
||
]):
|
||
return True
|
||
return False
|
||
|
||
|
||
def is_footer_chunk(text: str) -> bool:
|
||
"""DOCX-специфичные фильтры футера."""
|
||
text_lower = text.lower()
|
||
if any(kw in text_lower for kw in [
|
||
"документ:", "стр. из", "дата разработки:",
|
||
"разработчик:", "унифицированная",
|
||
]):
|
||
return True
|
||
return False
|
||
|
||
|
||
def is_junk(item: dict, text: str, cfg: dict) -> bool:
|
||
"""Общий фильтр мусора для обоих режимов."""
|
||
if not text:
|
||
return True
|
||
if len(text) < cfg["min_text_length"]:
|
||
return True
|
||
if item.get("type") in ("Header", "Footer", "Image"):
|
||
return True
|
||
if is_page_number(text):
|
||
return True
|
||
if is_header_noise(text):
|
||
return True
|
||
# DOCX-специфичные фильтры
|
||
if is_header_chunk(text):
|
||
return True
|
||
if is_footer_chunk(text):
|
||
return True
|
||
return False
|
||
|
||
|
||
# ============================================================
|
||
# HTML TABLE → TEXT
|
||
# ============================================================
|
||
|
||
def _ocr_table_fixes(text: str) -> str:
|
||
"""Применяет патч-исправления для OCR-ошибок в таблицах."""
|
||
if not text:
|
||
return text
|
||
|
||
# Порядок важен — длинные патчи первыми, более специфичные до общих
|
||
fixes = [
|
||
# OCR-исправления заголовков (самые важные) — длинные сначала
|
||
("рредакции ы", "редакции"),
|
||
("рредакции :", "редакции :"),
|
||
("рредакции:", "редакции:"),
|
||
("рредакции", "редакции"),
|
||
("едакции ы", "редакции"),
|
||
("едакции :", "редакции :"),
|
||
("едакции:", "редакции:"),
|
||
("едакции", "редакции"),
|
||
("едаты ", "даты "),
|
||
("едаты:", "даты:"),
|
||
("едаты", "даты"),
|
||
# OCR "р" перед словами (массовое)
|
||
("р внесенных :", "внесенных :"),
|
||
("р внесенных:", "внесенных:"),
|
||
("р внесенных", "внесенных"),
|
||
("р системное", "системное"),
|
||
("р системного", "системного"),
|
||
("р системном", "системном"),
|
||
("р Департамента", "Департамента"),
|
||
("р технологий", "технологий"),
|
||
("р администрирование", "администрирование"),
|
||
("р администрирования", "администрирования"),
|
||
("р информационных", "информационных"),
|
||
("р Группы", "Группы"),
|
||
("р технолог", "технолог"),
|
||
("р технологическим", "технологическим"),
|
||
("р персонализированный", "персонализированный"),
|
||
# OCR "р" перед другими словами
|
||
("р ", " "),
|
||
# OCR "Ошесв" / "Обесв" / "Отесв"
|
||
("Ошесвииа", "Общества"),
|
||
("Отесват", "Общества"),
|
||
("Обесват", "Общества"),
|
||
("Ошесват", "Общества"),
|
||
("Ошесв", "Общества"),
|
||
# Дублирование слов (администрирование администрирования)
|
||
(" администрирования администрирования", " администрирования"),
|
||
("администрирования администрирования", " администрирования"),
|
||
(" администрирование администрирование", " администрирование"),
|
||
("администрирование администрирование", " администрирование"),
|
||
# Убираем вертикальные разделители
|
||
("| Создано", "Создано"),
|
||
("| ", " "),
|
||
("|", ""),
|
||
("Список лиц:", "Список лиц"),
|
||
("Список лиц,", "Список лиц"),
|
||
# Бауэр
|
||
("Бауэ ", "Бауэр "),
|
||
("Бауэ", "Бауэр"),
|
||
("Бауэ С", "Бауэр С"),
|
||
# Общие OCR-исправления
|
||
("пк", "ПК"),
|
||
("ПК пк", "ПК"),
|
||
("РроюСаг", "PhotoCar"),
|
||
("р р", " "),
|
||
]
|
||
for bad, good in fixes:
|
||
text = text.replace(bad, good)
|
||
|
||
# Убираем "р" как отдельный OCR-мусор — ТОЛЬКО в конце слов (не в середине)
|
||
# НЕ используем \br\s — он ломает русские слова (Номер→Номе)
|
||
text = re.sub(r'\sр\b', '', text)
|
||
|
||
# Убираем двойные/тройные двоеточия
|
||
text = re.sub(r':\s*:\s*', ':', text)
|
||
text = re.sub(r':$', '', text)
|
||
|
||
# Убираем "№:" из заголовков
|
||
text = re.sub(r'№:\s*', '', text)
|
||
|
||
# Убираем двойные буквы (Бауэрр → Бауэр и т.д.)
|
||
text = re.sub(r'(.)\1\1+', r'\1\1', text)
|
||
|
||
# Убираем тройные пробелы
|
||
text = re.sub(r' +', ' ', text)
|
||
|
||
# Чистим лишние пробелы
|
||
text = re.sub(r'\s+', ' ', text)
|
||
return text.strip()
|
||
|
||
|
||
def _is_ocr_table(html: str) -> bool:
|
||
"""Проверяет, является ли HTML таблицы сильно повреждённым OCR."""
|
||
if not html:
|
||
return False
|
||
# Если HTML содержит много OCR-мусора
|
||
ocr_markers = ["едации", "едаты", "Ошесв", "Обесв", "рредакции", "р систем", "р Департ"]
|
||
ocr_count = sum(1 for m in ocr_markers if m in html)
|
||
return ocr_count >= 2
|
||
|
||
|
||
def _extract_table_text_fallback(table_text: str) -> str:
|
||
"""Извлекает структуру таблицы из уже обработанного text поля API."""
|
||
if not table_text:
|
||
return ""
|
||
lines = []
|
||
current_record = None
|
||
for line in table_text.split('\n'):
|
||
line = clean_text(line)
|
||
if not line:
|
||
continue
|
||
# Проверяем "Запись N:"
|
||
m = re.match(r'^Запись\s+(\d+):\s*(.*)', line)
|
||
if m:
|
||
current_record = m.group(1)
|
||
rest = m.group(2)
|
||
if rest:
|
||
lines.append(f"Запись {current_record}: {clean_text(rest)}")
|
||
continue
|
||
lines.append(line)
|
||
|
||
return "\n".join(lines)
|
||
|
||
|
||
def html_table_to_text(html: str, table_text: str = "") -> str:
|
||
"""Конвертирует HTML таблицу в читаемый текст.
|
||
|
||
Сначала чистит HTML от OCR-ошибок, затем парсит.
|
||
Если OCR слишком сильный — fallback из table_text.
|
||
"""
|
||
if not html and not table_text:
|
||
return ""
|
||
|
||
# Определяем, стоит ли использовать HTML-парсинг
|
||
# Если в HTML слишком много OCR-мусора — используем fallback
|
||
html_clean = _ocr_table_fixes(html)
|
||
|
||
# Проверяем качество HTML после очистки
|
||
ocr_markers = ["едации", "едаты", "Ошесв", "Обесв", "рредакции", "р систем", "р Департ"]
|
||
ocr_count = sum(1 for m in ocr_markers if m in html_clean)
|
||
|
||
# Если после очистки всё ещё есть OCR-ошибки — используем fallback
|
||
if ocr_count >= 2:
|
||
return _extract_table_text_fallback(table_text) if table_text else ""
|
||
|
||
# Парсим HTML
|
||
from html.parser import HTMLParser
|
||
|
||
class TableParser(HTMLParser):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.rows = []
|
||
self.current_row = []
|
||
self.in_table = False
|
||
|
||
def handle_starttag(self, tag, attrs):
|
||
if tag == "table":
|
||
self.in_table = True
|
||
elif tag == "tr" and self.in_table:
|
||
self.current_row = []
|
||
elif tag in ("td", "th") and self.in_table:
|
||
attrs_dict = dict(attrs)
|
||
rowspan = int(attrs_dict.get("rowspan", 1))
|
||
colspan = int(attrs_dict.get("colspan", 1))
|
||
self.current_row.append({
|
||
"text": "",
|
||
"rowspan": rowspan,
|
||
"colspan": colspan,
|
||
})
|
||
|
||
def handle_data(self, data):
|
||
if self.current_row:
|
||
self.current_row[-1]["text"] += data
|
||
|
||
def handle_endtag(self, tag):
|
||
if tag == "tr" and self.in_table and self.current_row:
|
||
self.rows.append(self.current_row)
|
||
self.current_row = []
|
||
|
||
parser = TableParser()
|
||
parser.feed(html_clean)
|
||
|
||
rows = parser.rows
|
||
if not rows:
|
||
return _extract_table_text_fallback(table_text) if table_text else ""
|
||
|
||
# Убираем пустые строки
|
||
rows = [r for r in rows if any(cell["text"].strip() for cell in r)]
|
||
|
||
if not rows:
|
||
return _extract_table_text_fallback(table_text) if table_text else ""
|
||
|
||
# Определяем ширину таблицы с учётом rowspan
|
||
max_cols = max(len(row) for row in rows)
|
||
table_width = max(max_cols, 1)
|
||
|
||
# Строим виртуальную сетку с rowspan
|
||
grid = [[None] * table_width for _ in range(len(rows))]
|
||
for row_idx, row in enumerate(rows):
|
||
col_idx = 0
|
||
for cell in row:
|
||
# Пропускаем занятые ячейки (rowspan)
|
||
while col_idx < table_width and grid[row_idx][col_idx] is not None:
|
||
col_idx += 1
|
||
if col_idx >= table_width:
|
||
break
|
||
grid[row_idx][col_idx] = {
|
||
"text": cell["text"],
|
||
"rowspan": cell["rowspan"],
|
||
"colspan": cell["colspan"],
|
||
}
|
||
# Занимаем ячейки для colspan
|
||
for c in range(1, cell["colspan"]):
|
||
next_col = col_idx + c
|
||
if next_col < table_width:
|
||
grid[row_idx][next_col] = {"text": "", "rowspan": 1, "colspan": 1}
|
||
|
||
# Определяем, является ли первая строка заголовком
|
||
def is_header_row(row_strings):
|
||
"""row_strings = list of str. Headers are short, no numbers."""
|
||
texts = [s.strip() for s in row_strings if s.strip()]
|
||
if not texts:
|
||
return False
|
||
short = all(len(t) < 50 for t in texts)
|
||
has_number = any(re.match(r"^\d+\.?\s*$", t) for t in texts)
|
||
return short and not has_number
|
||
|
||
# Парсим ячейки в плоские строки
|
||
def parse_row_cells(row_cells):
|
||
result = []
|
||
for cell in row_cells:
|
||
if cell is None:
|
||
result.append("")
|
||
else:
|
||
result.append(_ocr_table_fixes(cell["text"].strip()))
|
||
return result
|
||
|
||
parsed_rows = []
|
||
for row in grid:
|
||
parsed_rows.append(parse_row_cells(row))
|
||
|
||
# Фильтруем пустые строки
|
||
parsed_rows = [r for r in parsed_rows if any(c.strip() for c in r)]
|
||
|
||
if not parsed_rows:
|
||
return _extract_table_text_fallback(table_text) if table_text else ""
|
||
|
||
# Определяем заголовок
|
||
header_idx = 0
|
||
if is_header_row(parsed_rows[0]):
|
||
header_idx = 0
|
||
else:
|
||
for i, row in enumerate(parsed_rows):
|
||
if is_header_row(row):
|
||
header_idx = i
|
||
break
|
||
|
||
headers = parsed_rows[header_idx]
|
||
|
||
# Формируем вывод
|
||
lines = []
|
||
for row_idx, row in enumerate(parsed_rows[header_idx + 1:], header_idx + 1):
|
||
lines.append(f"Запись {row_idx}:")
|
||
for col_idx, cell in enumerate(row):
|
||
if not cell:
|
||
continue
|
||
header = headers[col_idx] if col_idx < len(headers) else f"Колонка {col_idx+1}"
|
||
header = _ocr_table_fixes(header)
|
||
lines.append(f"{header}: {cell}")
|
||
lines.append("")
|
||
|
||
result = "\n".join(lines).strip()
|
||
return result if result else _extract_table_text_fallback(table_text) if table_text else ""
|
||
|
||
|
||
# ============================================================
|
||
# MERGE FRAGMENTS
|
||
# ============================================================
|
||
|
||
def merge_fragments(elements: list) -> list:
|
||
"""Объединяет короткие чанки в буфер, таблицы — отдельно.
|
||
|
||
Улучшена логика слияния:
|
||
- Учитывает номер страницы (объединяет только с соседних)
|
||
- Учитывает section_title (не объединяет через разделы)
|
||
- Более высокий порог для разделения (300 вместо 80)
|
||
"""
|
||
merged = []
|
||
buffer = []
|
||
buffer_page = None
|
||
buffer_section = None
|
||
|
||
def flush_buffer():
|
||
nonlocal buffer
|
||
if not buffer:
|
||
return
|
||
text = clean_text(" ".join(buffer))
|
||
if text:
|
||
merged.append({
|
||
"type": "MergedText",
|
||
"text": text,
|
||
"page_number": buffer_page,
|
||
"section_title": buffer_section,
|
||
})
|
||
buffer = []
|
||
|
||
for el in elements:
|
||
text = clean_text(el.get("text", ""))
|
||
if not text:
|
||
continue
|
||
el_type = el.get("type")
|
||
page = el.get("metadata", {}).get("page_number", 0)
|
||
section = el.get("section_title", "")
|
||
|
||
if el_type == "Table":
|
||
flush_buffer()
|
||
merged.append({
|
||
"type": "Table",
|
||
"text": text,
|
||
"html": el.get("metadata", {}).get("text_as_html", ""),
|
||
"page_number": page,
|
||
"section_title": section,
|
||
})
|
||
continue
|
||
|
||
if el_type in ("NarrativeText", "ListItem", "UncategorizedText"):
|
||
# Check if we should flush the buffer (different section or page jump)
|
||
if buffer_section and section and buffer_section != section:
|
||
flush_buffer()
|
||
buffer = []
|
||
buffer_section = section
|
||
buffer_page = page
|
||
|
||
if len(text) < 300:
|
||
buffer.append(text)
|
||
buffer_page = page
|
||
buffer_section = section
|
||
else:
|
||
flush_buffer()
|
||
merged.append({
|
||
"type": "Text",
|
||
"text": text,
|
||
"page_number": page,
|
||
"section_title": section,
|
||
})
|
||
else:
|
||
flush_buffer()
|
||
|
||
flush_buffer()
|
||
return merged
|
||
|
||
|
||
# ============================================================
|
||
# API MODE — parse via Unstructured API
|
||
# ============================================================
|
||
|
||
def parse_document_api(file_path: str, cfg: dict) -> list:
|
||
"""Парсинг через Unstructured API (работает и с PDF, и с DOCX)."""
|
||
import requests
|
||
|
||
filename = os.path.basename(file_path)
|
||
print(f"\n📄 API Парсинг: {filename}")
|
||
|
||
payload_data = {
|
||
"strategy": cfg["strategy"],
|
||
"languages": cfg["languages"],
|
||
"pdf_infer_table_structure": str(cfg["pdf_infer_table_structure"]),
|
||
}
|
||
|
||
ext = os.path.splitext(filename)[1].lower()
|
||
mime_type = "application/pdf" if ext == ".pdf" else (
|
||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||
)
|
||
|
||
with open(file_path, "rb") as f:
|
||
files = {"files": (filename, f, mime_type)}
|
||
timeout = cfg.get("timeout", 900)
|
||
response = requests.post(
|
||
cfg["api_url"],
|
||
files=files,
|
||
data=payload_data,
|
||
timeout=timeout,
|
||
)
|
||
response.raise_for_status()
|
||
return response.json()
|
||
|
||
|
||
# ============================================================
|
||
# LOCAL MODE — parse DOCX via python unstructured
|
||
# ============================================================
|
||
|
||
def parse_document_local(file_path: str, cfg: dict) -> list:
|
||
"""Локальный парсинг DOCX через unstructured.partition.docx."""
|
||
from unstructured.partition.docx import partition_docx
|
||
from unstructured.chunking.title import chunk_by_title
|
||
|
||
print(f"\n📄 Local DOCX Парсинг: {os.path.basename(file_path)}")
|
||
|
||
elements = partition_docx(
|
||
file_path,
|
||
strategy=cfg["strategy"],
|
||
infer_table_structure=True,
|
||
)
|
||
print(f" Partitioned: {len(elements)} elements")
|
||
|
||
chunks = chunk_by_title(
|
||
elements,
|
||
max_characters=cfg["chunk_max_chars"],
|
||
combine_text_under_n_chars=cfg["chunk_combine_under"],
|
||
new_after_n_chars=cfg["chunk_new_after"],
|
||
)
|
||
print(f" Chunked: {len(chunks)} chunks")
|
||
|
||
# Конвертируем в единый формат с API
|
||
result = []
|
||
for c in chunks:
|
||
item = {
|
||
"type": c.category,
|
||
"text": c.text,
|
||
"metadata": {
|
||
"page_number": getattr(c.metadata, "page_number", 0) if hasattr(c, "metadata") and c.metadata else 0,
|
||
"text_as_html": getattr(c.metadata, "text_as_html", "") if hasattr(c, "metadata") and c.metadata else "",
|
||
},
|
||
}
|
||
# section_title из metadata
|
||
if hasattr(c, "metadata") and c.metadata:
|
||
if hasattr(c.metadata, "page_label") and c.metadata.page_label:
|
||
item["metadata"]["page_number"] = c.metadata.page_label
|
||
result.append(item)
|
||
|
||
return result
|
||
|
||
|
||
# ============================================================
|
||
# CLEAN + MERGE (единый пайплайн)
|
||
# ============================================================
|
||
|
||
def clean_and_merge(data: list, cfg: dict) -> list:
|
||
"""Очистка + фильтрация + слияние — общий для обоих режимов."""
|
||
cleaned = []
|
||
current_title = "Общая информация"
|
||
current_section_num = None
|
||
|
||
for item in data:
|
||
raw_text = item.get("text", "")
|
||
text = clean_text(raw_text)
|
||
if is_junk(item, text, cfg):
|
||
continue
|
||
|
||
# Try to detect section titles from content (not just Title type)
|
||
section_num, section_title = _detect_section_number(text)
|
||
if section_num is not None or section_title != text:
|
||
if section_num is not None:
|
||
current_section_num = section_num
|
||
current_title = section_title if section_title else text
|
||
# Don't add the title itself as content
|
||
continue
|
||
|
||
item["section_title"] = current_title
|
||
item["section_number"] = current_section_num
|
||
cleaned.append(item)
|
||
|
||
merged = merge_fragments(cleaned)
|
||
return merged
|
||
|
||
|
||
# ============================================================
|
||
# BUILD CHUNKS + CONTENT
|
||
# ============================================================
|
||
|
||
def build_chunks(merged: list, filename: str, cfg: dict) -> list:
|
||
"""Формирует финальные чанки с content для индексации.
|
||
|
||
Добавлен section_id, улучшена фильтрация мусорных чанков.
|
||
"""
|
||
chunks = []
|
||
for idx, chunk in enumerate(merged):
|
||
chunk_type = chunk["type"]
|
||
section_title = chunk["section_title"]
|
||
page_number = chunk["page_number"]
|
||
section_num = chunk.get("section_number", None)
|
||
|
||
# Generate section_id
|
||
section_id = _generate_section_id(section_num, section_title)
|
||
|
||
if chunk_type == "Table":
|
||
table_text = html_table_to_text(chunk.get("html", ""), chunk.get("text", ""))
|
||
content = f"Раздел: {section_title}\nТип: Таблица\n\n{table_text}"
|
||
else:
|
||
content = f"Раздел: {section_title}\n\n{chunk['text']}"
|
||
|
||
content = clean_text(content)
|
||
if len(content) < cfg["min_chunk_length"]:
|
||
continue
|
||
|
||
# Filter out garbage chunks (too short content after header)
|
||
if chunk_type == "Table" and len(table_text) < 50 if "table_text" in locals() else len(content) < 100:
|
||
continue
|
||
|
||
# Filter out chunks that are just headers with no real content
|
||
if chunk["type"] in ("MergedText", "Text") and content.count("\n") <= 1 and len(content) < 100:
|
||
continue
|
||
|
||
chunk_hash = hashlib.sha256(
|
||
(filename + str(idx) + content).encode("utf-8")
|
||
).hexdigest()
|
||
|
||
chunks.append({
|
||
"id": str(uuid.uuid5(uuid.NAMESPACE_DNS, chunk_hash)),
|
||
"content": content,
|
||
"metadata": {
|
||
"filename": filename,
|
||
"chunk_index": idx,
|
||
"section_title": section_title,
|
||
"section_id": section_id,
|
||
"page_number": page_number,
|
||
"content_type": chunk_type,
|
||
"chunk_hash": chunk_hash,
|
||
},
|
||
})
|
||
|
||
return chunks
|
||
|
||
|
||
|
||
|
||
|
||
# ============================================================
|
||
# MAIN ENTRY POINT
|
||
# ============================================================
|
||
|
||
def parse_document(file_path: str, cfg: Optional[dict] = None) -> dict:
|
||
"""
|
||
Универсальный парсер документов (PDF + DOCX).
|
||
|
||
Возвращает:
|
||
{
|
||
"filename": str,
|
||
"chunks": list[dict], # каждый: {id, content, metadata}
|
||
"stats": {
|
||
"total_elements": int,
|
||
"total_chunks_before_filter": int,
|
||
"total_chunks_after_filter": int,
|
||
"total_final_chunks": int,
|
||
}
|
||
}
|
||
"""
|
||
if cfg is None:
|
||
cfg = get_config()
|
||
|
||
filename = os.path.basename(file_path)
|
||
|
||
# --- PARSE ---
|
||
if cfg["local_mode"]:
|
||
raw_data = parse_document_local(file_path, cfg)
|
||
else:
|
||
raw_data = parse_document_api(file_path, cfg)
|
||
|
||
total_elements = len(raw_data)
|
||
|
||
# --- CLEAN + MERGE ---
|
||
merged = clean_and_merge(raw_data, cfg)
|
||
total_after_filter = len(merged)
|
||
|
||
# --- BUILD CHUNKS ---
|
||
chunks = build_chunks(merged, filename, cfg)
|
||
total_final = len(chunks)
|
||
|
||
stats = {
|
||
"total_elements": total_elements,
|
||
"total_chunks_before_filter": total_elements,
|
||
"total_chunks_after_filter": total_after_filter,
|
||
"total_final_chunks": total_final,
|
||
}
|
||
|
||
return {
|
||
"filename": filename,
|
||
"chunks": chunks,
|
||
"stats": stats,
|
||
}
|
||
|
||
|
||
# ============================================================
|
||
# CLI
|
||
# ============================================================
|
||
|
||
def main():
|
||
"""CLI: python document_parser.py <file> [--api|--local] [--json]"""
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="Unified DOCX/PDF Document Parser")
|
||
parser.add_argument("file", help="Path to PDF or DOCX file")
|
||
parser.add_argument("--api", action="store_true", default=None,
|
||
help="Use API mode (default: auto-detect from env)")
|
||
parser.add_argument("--local", action="store_true", default=None,
|
||
help="Use local unstructured library")
|
||
parser.add_argument("--json", action="store_true", default=False,
|
||
help="Output parsed JSON to stdout")
|
||
parser.add_argument("--output", "-o", default=None,
|
||
help="Output JSON file path")
|
||
args = parser.parse_args()
|
||
|
||
file_path = os.path.abspath(args.file)
|
||
if not os.path.exists(file_path):
|
||
print(f"❌ File not found: {file_path}")
|
||
sys.exit(1)
|
||
|
||
# Override mode from CLI
|
||
if args.local:
|
||
cfg = get_config()
|
||
cfg["local_mode"] = True
|
||
elif args.api:
|
||
cfg = get_config()
|
||
cfg["local_mode"] = False
|
||
else:
|
||
cfg = get_config()
|
||
|
||
result = parse_document(file_path, cfg)
|
||
|
||
print(f"\n✅ {result['filename']}")
|
||
print(f" Elements: {result['stats']['total_elements']}")
|
||
print(f" After filter: {result['stats']['total_chunks_after_filter']}")
|
||
print(f" Final chunks: {result['stats']['total_final_chunks']}")
|
||
|
||
# Output
|
||
output_json = json.dumps(result, ensure_ascii=False, indent=2)
|
||
|
||
if args.json or args.output:
|
||
if args.output:
|
||
with open(args.output, "w", encoding="utf-8") as f:
|
||
f.write(output_json)
|
||
print(f"\n📄 Saved to: {args.output}")
|
||
else:
|
||
print(output_json)
|
||
|
||
return result
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|