cesnimda
411 lines
17 KiB
Python
411 lines
17 KiB
Python
"""Resource-bounded CV document decoder invoked only as a child process."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import io
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import sys
|
|
import zipfile
|
|
from pathlib import PurePosixPath
|
|
|
|
|
|
MAX_EXTRACT_CHARS = 200_000
|
|
MAX_PDF_PAGES = 40
|
|
MAX_DOCX_ENTRIES = 256
|
|
MAX_DOCX_UNCOMPRESSED_BYTES = 32 * 1024 * 1024
|
|
MAX_DOCX_ENTRY_BYTES = 8 * 1024 * 1024
|
|
MAX_DOCX_COMPRESSION_RATIO = 100
|
|
MAX_IMAGE_DIMENSION = 12_000
|
|
MAX_IMAGE_PIXELS = 40_000_000
|
|
MAX_PDF_PAGE_PIXELS = 12_000_000
|
|
MAX_PDF_OCR_PIXELS = 120_000_000
|
|
OCR_LANGUAGES = "eng"
|
|
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
|
|
|
|
|
|
class ParserError(Exception):
|
|
def __init__(self, code: str):
|
|
super().__init__(code)
|
|
self.code = code
|
|
|
|
|
|
def _apply_resource_limits(cpu_seconds: int, memory_bytes: int) -> None:
|
|
if os.name != "posix":
|
|
return
|
|
import resource
|
|
|
|
resource.setrlimit(resource.RLIMIT_CPU, (cpu_seconds, cpu_seconds + 1))
|
|
resource.setrlimit(resource.RLIMIT_AS, (memory_bytes, memory_bytes))
|
|
resource.setrlimit(resource.RLIMIT_FSIZE, (1024 * 1024, 1024 * 1024))
|
|
resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64))
|
|
resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
|
|
|
|
|
|
def normalize_text(value: str) -> str:
|
|
value = value.replace("\x00", " ").replace("\r\n", "\n").replace("\r", "\n")
|
|
lines: list[str] = []
|
|
blank = False
|
|
for raw_line in value.split("\n"):
|
|
line = re.sub(r"[\t\f\v]+", " ", raw_line).strip()
|
|
if not line:
|
|
if lines and not blank:
|
|
lines.append("")
|
|
blank = True
|
|
continue
|
|
lines.append(line)
|
|
blank = False
|
|
return "\n".join(lines).strip()
|
|
|
|
|
|
def _check_extracted_size(text: str, limit: int = MAX_EXTRACT_CHARS) -> str:
|
|
if len(text) > limit:
|
|
raise ParserError("extracted_text_limit")
|
|
return text
|
|
|
|
|
|
def _check_image_size(image, *, pixel_limit: int = MAX_IMAGE_PIXELS) -> None:
|
|
width, height = image.size
|
|
if width <= 0 or height <= 0 or width > MAX_IMAGE_DIMENSION or height > MAX_IMAGE_DIMENSION:
|
|
raise ParserError("image_dimensions")
|
|
if width * height > pixel_limit:
|
|
raise ParserError("image_pixel_limit")
|
|
if getattr(image, "n_frames", 1) != 1:
|
|
raise ParserError("image_frames")
|
|
|
|
|
|
def validate_docx_container(
|
|
data: bytes,
|
|
*,
|
|
entry_limit: int = MAX_DOCX_ENTRIES,
|
|
total_limit: int = MAX_DOCX_UNCOMPRESSED_BYTES,
|
|
entry_size_limit: int = MAX_DOCX_ENTRY_BYTES,
|
|
ratio_limit: int = MAX_DOCX_COMPRESSION_RATIO,
|
|
) -> None:
|
|
try:
|
|
with zipfile.ZipFile(io.BytesIO(data)) as archive:
|
|
entries = archive.infolist()
|
|
names = {entry.filename for entry in entries}
|
|
if "[Content_Types].xml" not in names or "word/document.xml" not in names:
|
|
raise ParserError("docx_signature")
|
|
if len(entries) > entry_limit:
|
|
raise ParserError("docx_entry_limit")
|
|
total_size = 0
|
|
for entry in entries:
|
|
total_size += entry.file_size
|
|
if entry.file_size > entry_size_limit:
|
|
raise ParserError("docx_entry_size")
|
|
if entry.file_size > 0 and entry.compress_size == 0:
|
|
raise ParserError("docx_ratio")
|
|
if entry.compress_size > 0 and entry.file_size / entry.compress_size > ratio_limit:
|
|
raise ParserError("docx_ratio")
|
|
if total_size > total_limit:
|
|
raise ParserError("docx_expanded_size")
|
|
except ParserError:
|
|
raise
|
|
except (zipfile.BadZipFile, OSError) as exc:
|
|
raise ParserError("docx_signature") from exc
|
|
|
|
|
|
def extract_plain_text(data: bytes) -> str:
|
|
try:
|
|
decoded = data.decode("utf-8-sig")
|
|
except UnicodeDecodeError as exc:
|
|
raise ParserError("text_encoding") from exc
|
|
return _check_extracted_size(normalize_text(decoded))
|
|
|
|
|
|
def extract_docx_text(data: bytes) -> str:
|
|
validate_docx_container(data)
|
|
from docx import Document
|
|
|
|
document = Document(io.BytesIO(data))
|
|
parts: list[str] = []
|
|
blocks = document.iter_inner_content() if hasattr(document, "iter_inner_content") else document.paragraphs
|
|
for block in blocks:
|
|
if hasattr(block, "rows"):
|
|
for row in block.rows:
|
|
cells = [cell.text.strip() for cell in row.cells if cell.text and cell.text.strip()]
|
|
if cells:
|
|
parts.append(" | ".join(cells))
|
|
elif getattr(block, "text", "").strip():
|
|
text = block.text.strip()
|
|
style = (getattr(getattr(block, "style", None), "name", "") or "").lower()
|
|
if "role" in style and parts:
|
|
parts.append("")
|
|
parts.append(f"- {text}" if "bullet" in style else text)
|
|
return _check_extracted_size(normalize_text("\n".join(parts)))
|
|
|
|
|
|
def _ocr_image(image) -> str:
|
|
import pytesseract
|
|
|
|
if image.mode not in ("RGB", "L"):
|
|
image = image.convert("RGB")
|
|
return normalize_text(pytesseract.image_to_string(image, lang=OCR_LANGUAGES))
|
|
|
|
|
|
def extract_image_text(data: bytes, *, pixel_limit: int = MAX_IMAGE_PIXELS) -> str:
|
|
from PIL import Image
|
|
|
|
with Image.open(io.BytesIO(data)) as image:
|
|
_check_image_size(image, pixel_limit=pixel_limit)
|
|
return _check_extracted_size(_ocr_image(image))
|
|
|
|
|
|
def extract_pdf_text(data: bytes, *, page_limit: int = MAX_PDF_PAGES) -> tuple[str, bool, int]:
|
|
from pypdf import PdfReader
|
|
|
|
page_count = 0
|
|
extracted_pages: list[str] = []
|
|
try:
|
|
reader = PdfReader(io.BytesIO(data))
|
|
page_count = len(reader.pages)
|
|
if page_count > page_limit:
|
|
raise ParserError("pdf_page_limit")
|
|
for page in reader.pages:
|
|
try:
|
|
extracted_pages.append(page.extract_text(extraction_mode="layout") or "")
|
|
except TypeError:
|
|
extracted_pages.append(page.extract_text() or "")
|
|
except ParserError:
|
|
raise
|
|
except Exception:
|
|
extracted_pages = []
|
|
|
|
text = _check_extracted_size(normalize_text("\n".join(extracted_pages)))
|
|
if len(text) >= 80:
|
|
return text, False, page_count
|
|
|
|
import fitz
|
|
from PIL import Image
|
|
|
|
document = fitz.open(stream=data, filetype="pdf")
|
|
try:
|
|
page_count = document.page_count
|
|
if page_count > page_limit:
|
|
raise ParserError("pdf_page_limit")
|
|
ocr_pages: list[str] = []
|
|
total_pixels = 0
|
|
for page in document:
|
|
pixmap = page.get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False)
|
|
page_pixels = pixmap.width * pixmap.height
|
|
total_pixels += page_pixels
|
|
if page_pixels > MAX_PDF_PAGE_PIXELS or total_pixels > MAX_PDF_OCR_PIXELS:
|
|
raise ParserError("pdf_ocr_limit")
|
|
with Image.open(io.BytesIO(pixmap.tobytes("png"))) as image:
|
|
_check_image_size(image, pixel_limit=MAX_PDF_PAGE_PIXELS)
|
|
ocr_pages.append(_ocr_image(image))
|
|
_check_extracted_size("\n".join(ocr_pages))
|
|
return _check_extracted_size(normalize_text("\n".join(ocr_pages))), True, page_count
|
|
finally:
|
|
document.close()
|
|
|
|
|
|
def extract(extension: str, data: bytes) -> dict[str, object]:
|
|
if extension in {".txt", ".md"}:
|
|
text = extract_plain_text(data)
|
|
return {"text": text, "ocr_used": False, "page_count": None}
|
|
if extension == ".docx":
|
|
text = extract_docx_text(data)
|
|
return {"text": text, "ocr_used": False, "page_count": None}
|
|
if extension == ".pdf":
|
|
text, ocr_used, page_count = extract_pdf_text(data)
|
|
return {"text": text, "ocr_used": ocr_used, "page_count": page_count}
|
|
if extension in IMAGE_EXTENSIONS:
|
|
text = extract_image_text(data)
|
|
return {"text": text, "ocr_used": True, "page_count": 1}
|
|
raise ParserError("unsupported_type")
|
|
|
|
|
|
def _validate_import_docx(data: bytes) -> None:
|
|
validate_docx_container(data)
|
|
from defusedxml import ElementTree
|
|
|
|
try:
|
|
with zipfile.ZipFile(io.BytesIO(data)) as archive:
|
|
names: set[str] = set()
|
|
for entry in archive.infolist():
|
|
name = entry.filename
|
|
path = PurePosixPath(name)
|
|
if name in names or path.is_absolute() or ".." in path.parts or "\\" in name or ":" in name or entry.flag_bits & 1:
|
|
raise ParserError("docx_unsafe_archive")
|
|
names.add(name)
|
|
if name.lower().endswith((".bin", ".emf", ".wmf")) or name.startswith("word/embeddings/"):
|
|
raise ParserError("external_content")
|
|
if name.endswith((".xml", ".rels")):
|
|
root = ElementTree.fromstring(archive.read(entry))
|
|
for node in root.iter():
|
|
if node.tag.rsplit("}", 1)[-1] in {"altChunk", "OLEObject", "object"}:
|
|
raise ParserError("external_content")
|
|
if node.attrib.get("TargetMode", "").lower() == "external":
|
|
# Hyperlinks are inert source data; converters cannot follow them in the sandbox.
|
|
from urllib.parse import urlsplit
|
|
if node.attrib.get("Type", "").endswith("/hyperlink") and urlsplit(node.attrib.get("Target", "")).scheme.lower() in {"http", "https", "mailto"}:
|
|
continue
|
|
raise ParserError("external_content")
|
|
target = node.attrib.get("Target", "")
|
|
if ":" in target or "\\" in target or target.startswith("/"):
|
|
raise ParserError("external_content")
|
|
except ParserError:
|
|
raise
|
|
except Exception as exc:
|
|
raise ParserError("docx_signature") from exc
|
|
|
|
|
|
def _validate_import_pdf(data: bytes) -> int:
|
|
if not data.startswith(b"%PDF-"):
|
|
raise ParserError("pdf_signature")
|
|
from pypdf import PdfReader
|
|
from pypdf.generic import ArrayObject, DictionaryObject, IndirectObject
|
|
|
|
try:
|
|
reader = PdfReader(io.BytesIO(data), strict=True)
|
|
if reader.is_encrypted:
|
|
raise ParserError("pdf_encrypted")
|
|
count = len(reader.pages)
|
|
if count > MAX_PDF_PAGES:
|
|
raise ParserError("pdf_page_limit")
|
|
forbidden = {"/OpenAction", "/AA", "/JS", "/JavaScript", "/Launch", "/GoToR", "/SubmitForm", "/ImportData", "/EmbeddedFiles", "/Filespec", "/RichMedia", "/XFA"}
|
|
pending = [(reader.trailer, 0)]
|
|
seen: set[int] = set()
|
|
while pending:
|
|
value, depth = pending.pop()
|
|
if isinstance(value, IndirectObject):
|
|
value = value.get_object()
|
|
if not isinstance(value, (DictionaryObject, ArrayObject)) or id(value) in seen:
|
|
continue
|
|
seen.add(id(value))
|
|
if depth > 40 or len(seen) > 10_000:
|
|
raise ParserError("pdf_object_limit")
|
|
if isinstance(value, DictionaryObject):
|
|
from urllib.parse import urlsplit
|
|
if "/URI" in value and urlsplit(str(value["/URI"])).scheme.lower() not in {"http", "https", "mailto"}:
|
|
raise ParserError("external_content")
|
|
if forbidden.intersection(value) or str(value.get("/S", "")) in forbidden or str(value.get("/Type", "")) in forbidden:
|
|
raise ParserError("external_content")
|
|
pending.extend((item, depth + 1) for item in value.values())
|
|
else:
|
|
pending.extend((item, depth + 1) for item in value)
|
|
return count
|
|
except ParserError:
|
|
raise
|
|
except Exception as exc:
|
|
raise ParserError("pdf_invalid") from exc
|
|
|
|
|
|
def extract_document(extension: str, data: bytes) -> dict[str, object]:
|
|
"""Model-free Docling conversion of validated bytes, never URLs or local paths."""
|
|
if not data or len(data) > 5 * 1024 * 1024:
|
|
raise ParserError("file_size_limit")
|
|
if extension == ".docx":
|
|
_validate_import_docx(data)
|
|
page_count = None
|
|
warnings = ["docx_page_numbers_unavailable", "ocr_disabled"]
|
|
elif extension == ".pdf":
|
|
page_count = _validate_import_pdf(data)
|
|
warnings = ["pdf_native_order_no_layout_or_tables", "ocr_disabled"]
|
|
else:
|
|
raise ParserError("unsupported_type")
|
|
|
|
from docling.datamodel.base_models import DocumentStream, InputFormat, ConversionStatus
|
|
from docling.datamodel.backend_options import MsWordBackendOptions, ThreadedDoclingParseBackendOptions
|
|
from docling.datamodel.pipeline_options import ConvertPipelineOptions, NativePdfPipelineOptions
|
|
from docling.document_converter import DocumentConverter, WordFormatOption, NativePdfFormatOption
|
|
from docling_core.types.doc import TextItem, TableItem
|
|
|
|
common = dict(enable_remote_services=False, allow_external_plugins=False, do_picture_classification=False, do_picture_description=False, document_timeout=20)
|
|
if extension == ".pdf":
|
|
source_format = InputFormat.PDF
|
|
option = NativePdfFormatOption(
|
|
pipeline_options=NativePdfPipelineOptions(**common, generate_page_images=False, generate_picture_images=False, parser_threads=1),
|
|
backend_options=ThreadedDoclingParseBackendOptions(parser_threads=1, max_concurrent_results=1, render_pages=False, include_bitmap_images=False),
|
|
)
|
|
else:
|
|
source_format = InputFormat.DOCX
|
|
option = WordFormatOption(
|
|
pipeline_options=ConvertPipelineOptions(**common),
|
|
backend_options=MsWordBackendOptions(render_chart_images=False, max_total_bytes=MAX_DOCX_UNCOMPRESSED_BYTES, max_file_bytes=MAX_DOCX_ENTRY_BYTES, max_member_count=MAX_DOCX_ENTRIES),
|
|
)
|
|
converter = DocumentConverter(allowed_formats=[source_format], format_options={source_format: option})
|
|
conversion = converter.convert(DocumentStream(name="document" + extension, stream=io.BytesIO(data)), max_num_pages=MAX_PDF_PAGES, max_file_size=5 * 1024 * 1024)
|
|
if conversion.status != ConversionStatus.SUCCESS:
|
|
raise ParserError("parser_failed")
|
|
blocks: list[dict[str, object]] = []
|
|
characters = 0
|
|
for item, _ in conversion.document.iterate_items():
|
|
if isinstance(item, TableItem):
|
|
text = normalize_text(item.export_to_markdown(doc=conversion.document))
|
|
kind = "table"
|
|
elif isinstance(item, TextItem):
|
|
text = normalize_text(item.text)
|
|
kind = {"section_header": "section_header", "title": "title", "list_item": "list_item"}.get(item.label.value, "paragraph")
|
|
else:
|
|
continue
|
|
if not text:
|
|
continue
|
|
characters += len(text) + (1 if blocks else 0)
|
|
if characters > MAX_EXTRACT_CHARS:
|
|
raise ParserError("extracted_text_limit")
|
|
if len(blocks) >= 5000:
|
|
raise ParserError("block_limit")
|
|
page = item.prov[0].page_no if extension == ".pdf" and item.prov else None
|
|
blocks.append({"id": f"b{len(blocks) + 1:04d}", "text": text, "kind": kind, "page": page})
|
|
if not blocks:
|
|
raise ParserError("no_text")
|
|
return {"text": "\n".join(block["text"] for block in blocks), "ocr_used": False, "page_count": page_count, "blocks": blocks, "warnings": warnings, "method": "docling-native-pdf-v1" if extension == ".pdf" else "docling-docx-v1"}
|
|
|
|
|
|
def _deny_import_side_effects(event: str, args: tuple) -> None:
|
|
# Model-free conversion has no reason to open a socket or invoke an external tool.
|
|
if event in {"socket.connect", "socket.connect_ex", "socket.getaddrinfo", "subprocess.Popen", "os.system"}:
|
|
raise ParserError("external_content")
|
|
|
|
|
|
def _write_result(path: Path, result: dict[str, object]) -> None:
|
|
path.write_text(json.dumps(result, ensure_ascii=False, separators=(",", ":")), encoding="utf-8")
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(add_help=False)
|
|
parser.add_argument("--input", required=True)
|
|
parser.add_argument("--output", required=True)
|
|
parser.add_argument("--extension", required=True)
|
|
parser.add_argument("--cpu-seconds", type=int, default=20)
|
|
parser.add_argument("--memory-bytes", type=int, default=512 * 1024 * 1024)
|
|
parser.add_argument("--mode", choices=("legacy", "document"), default="legacy")
|
|
args = parser.parse_args()
|
|
output_path = Path(args.output)
|
|
try:
|
|
_apply_resource_limits(args.cpu_seconds, args.memory_bytes)
|
|
data = Path(args.input).read_bytes()
|
|
if args.mode == "document":
|
|
os.environ.update({"HF_HUB_OFFLINE": "1", "HF_HUB_DISABLE_TELEMETRY": "1", "TRANSFORMERS_OFFLINE": "1"})
|
|
# Cache trusted platform metadata before the document sandbox denies Linux uname.
|
|
import platform
|
|
platform.platform()
|
|
sys.addaudithook(_deny_import_side_effects)
|
|
result = extract_document(args.extension.lower(), data)
|
|
else:
|
|
result = extract(args.extension.lower(), data)
|
|
if not result.get("text"):
|
|
raise ParserError("no_text")
|
|
_write_result(output_path, {"ok": True, **result})
|
|
return 0
|
|
except ParserError as exc:
|
|
_write_result(output_path, {"ok": False, "code": exc.code})
|
|
return 2
|
|
except ImportError:
|
|
_write_result(output_path, {"ok": False, "code": "docling_unavailable" if args.mode == "document" else "parser_failed"})
|
|
return 3
|
|
except BaseException:
|
|
_write_result(output_path, {"ok": False, "code": "parser_failed"})
|
|
return 3
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|