107 lines
4.4 KiB
Python
107 lines
4.4 KiB
Python
"""PDF/DOCX 원문을 위치 좌표가 보존된 세그먼트로 변환한다."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Iterator
|
|
|
|
from app.engine.provenance import DocumentRecord, SegmentRecord, stable_id
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExtractionResult:
|
|
document: DocumentRecord
|
|
segments: tuple[SegmentRecord, ...]
|
|
warnings: tuple[str, ...] = ()
|
|
|
|
|
|
def file_sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as stream:
|
|
for block in iter(lambda: stream.read(1024 * 1024), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def chunk_with_offsets(text: str, size: int = 1000, stride: int = 500) -> Iterator[tuple[int, int, str]]:
|
|
if size < 1 or stride < 1:
|
|
raise ValueError("size and stride must be positive")
|
|
text = text or ""
|
|
for start in range(0, len(text), stride):
|
|
end = min(len(text), start + size)
|
|
chunk = text[start:end]
|
|
if chunk.strip():
|
|
yield start, end, chunk
|
|
if end >= len(text):
|
|
break
|
|
|
|
|
|
def extract_pdf(path: Path, size: int = 1000, stride: int = 500) -> ExtractionResult:
|
|
try:
|
|
from pypdf import PdfReader
|
|
except ImportError as exc:
|
|
raise RuntimeError("PDF 추출에는 pypdf가 필요합니다") from exc
|
|
sha = file_sha256(path)
|
|
document_id = stable_id("doc", sha)
|
|
document = DocumentRecord(
|
|
document_id=document_id, title=path.stem, source_path=str(path), source_sha256=sha,
|
|
metadata={"format": "pdf", "coordinate_scope": "page"},
|
|
)
|
|
segments: list[SegmentRecord] = []
|
|
warnings: list[str] = []
|
|
reader = PdfReader(str(path))
|
|
for page_number, page in enumerate(reader.pages, 1):
|
|
text = page.extract_text() or ""
|
|
if not text.strip():
|
|
warnings.append(f"page {page_number}: text 없음(OCR 필요 가능)")
|
|
continue
|
|
for local_index, (start, end, chunk) in enumerate(chunk_with_offsets(text, size, stride), 1):
|
|
segments.append(SegmentRecord(
|
|
segment_id=stable_id("seg", document_id, str(page_number), str(start), chunk),
|
|
document_id=document_id, text=chunk,
|
|
ordinal=f"p{page_number:05d}-{local_index:04d}", coordinate_scope="page",
|
|
page_number=page_number, char_start=start, char_end=end,
|
|
source_locator=f"{path.name}#page={page_number}&chars={start}-{end}",
|
|
metadata={"extractor": "pypdf", "page_text_coordinates": True},
|
|
))
|
|
return ExtractionResult(document, tuple(segments), tuple(warnings))
|
|
|
|
|
|
def extract_docx(path: Path, size: int = 1000, stride: int = 500) -> ExtractionResult:
|
|
try:
|
|
from docx import Document
|
|
except ImportError as exc:
|
|
raise RuntimeError("DOCX 추출에는 python-docx가 필요합니다") from exc
|
|
sha = file_sha256(path)
|
|
document_id = stable_id("doc", sha)
|
|
document = DocumentRecord(
|
|
document_id=document_id, title=path.stem, source_path=str(path), source_sha256=sha,
|
|
metadata={"format": "docx", "coordinate_scope": "paragraph"},
|
|
)
|
|
segments: list[SegmentRecord] = []
|
|
source = Document(str(path))
|
|
for paragraph_number, paragraph in enumerate(source.paragraphs, 1):
|
|
text = paragraph.text or ""
|
|
for local_index, (start, end, chunk) in enumerate(chunk_with_offsets(text, size, stride), 1):
|
|
segments.append(SegmentRecord(
|
|
segment_id=stable_id("seg", document_id, str(paragraph_number), str(start), chunk),
|
|
document_id=document_id, text=chunk,
|
|
ordinal=f"para{paragraph_number:06d}-{local_index:04d}",
|
|
coordinate_scope="paragraph", paragraph_number=paragraph_number,
|
|
char_start=start, char_end=end,
|
|
source_locator=f"{path.name}#paragraph={paragraph_number}&chars={start}-{end}",
|
|
metadata={"extractor": "python-docx", "paragraph_text_coordinates": True},
|
|
))
|
|
return ExtractionResult(document, tuple(segments))
|
|
|
|
|
|
def extract_source(path: Path, size: int = 1000, stride: int = 500) -> ExtractionResult:
|
|
suffix = path.suffix.lower()
|
|
if suffix == ".pdf":
|
|
return extract_pdf(path, size, stride)
|
|
if suffix == ".docx":
|
|
return extract_docx(path, size, stride)
|
|
raise ValueError(f"지원하지 않는 원본 형식: {path.suffix}")
|