Agentic-Service-Data-Eyond-Catalog

Running

Agentic-Service-Data-Eyond-Catalog / src /knowledge_parsing /parse.py

[NOTICKET] refactor(knowledge-parsing): English throughout the module

7f153d6 2 months ago

6.71 kB

"""Stage 2 — parse a document with MinerU, cached by content.
Why the cache is keyed to document CONTENT rather than to a run id: parsing is
by far the most expensive stage (19 s/page on CPU). A per-run cache would
re-parse everything each time the normaliser changes and the pipeline is run
again, for nothing. Keyed by content + settings + version, the same document is
parsed once — however many times the normaliser is revised.
Cache key = sha256(file bytes) + settings fingerprint + MinerU version. The
version is in there so figures from an older MinerU never quietly mix with
figures from a newer one — which matters, because those figures are used to
justify spend.
"""

from __future__ import annotations

import hashlib
import json
import shutil
import time
from dataclasses import dataclass
from pathlib import Path

from .config import PipelineConfig


@dataclass
class ParseResult:
    doc_id: str
    source: Path
    cache_dir: Path          # folder holding MinerU's output untouched
    content_list: Path       # *_content_list.json
    middle_json: Path | None
    pages: int
    seconds: float
    from_cache: bool
    backend_recorded: str | None    # read from _middle.json, NOT from the config
    mineru_version: str | None


def mineru_version() -> str:
    try:
        from mineru.version import __version__  # type: ignore
        return str(__version__)
    except Exception:
        try:
            from importlib.metadata import version
            return version("mineru")
        except Exception:
            return "unknown"


def file_hash(path: Path, length: int = 16) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for block in iter(lambda: f.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()[:length]


def cache_key(source: Path, cfg: PipelineConfig) -> str:
    return f"{file_hash(source)}-{cfg.fingerprint()}-{mineru_version()}"


def _find_output(root: Path) -> tuple[Path | None, Path | None]:
    """Locate content_list & middle json anywhere under root.
    Deliberately a glob rather than a fixed path: MinerU's subfolder layout
    differs per backend ('auto' for pipeline, something else for vlm). Globbing
    means this module needs no change when the backend does.
    """
    content = next(iter(sorted(root.rglob("*_content_list.json"))), None)
    middle = next(iter(sorted(root.rglob("*_middle.json"))), None)
    return content, middle


def _read_middle(middle: Path | None) -> tuple[str | None, str | None]:
    """The backend & version that ACTUALLY ran, taken from MinerU's own output."""
    if not middle or not middle.exists():
        return None, None
    try:
        d = json.loads(middle.read_text(encoding="utf-8"))
        return d.get("_backend"), d.get("_version_name")
    except Exception:
        return None, None


def parse_document(source: Path, cfg: PipelineConfig) -> ParseResult:
    """Parse one document. A cache hit is not re-run."""
    doc_id = source.stem
    target = cfg.cache_dir / "parse" / cache_key(source, cfg)
    marker = target / ".done"

    if marker.exists():
        content, middle = _find_output(target)
        if content:
            meta = json.loads(marker.read_text(encoding="utf-8"))
            backend_recorded, version = _read_middle(middle)
            return ParseResult(
                doc_id=doc_id, source=source, cache_dir=target,
                content_list=content, middle_json=middle,
                pages=meta.get("pages", 0), seconds=meta.get("seconds", 0.0),
                from_cache=True,
                backend_recorded=backend_recorded, mineru_version=version,
            )
        shutil.rmtree(target, ignore_errors=True)   # cache corrupt, redo it

    from mineru.cli.common import do_parse, read_fn

    # MinerU dispatches on prefixes: "pipeline", then `startswith("vlm-")` and
    # `startswith("hybrid-")`. The bare names "vlm"/"hybrid" match NONE of those
    # branches — MinerU simply does nothing, without error, and the failure only
    # surfaces here as "content_list not found", which points at the wrong place
    # entirely. The short names stay in the config because they are part of the
    # cache key and read better; the translation happens here, right before
    # MinerU is touched.
    mineru_backend = {
        "pipeline": "pipeline",
        "vlm": "vlm-engine",
        "hybrid": "hybrid-engine",
    }[cfg.backend]

    # The playground does this before parsing; aliases like "latin"/"en" are
    # resolved to a canonical code. Matched here so results are identical.
    lang = cfg.canonical_lang()

    # with_suffix() REPLACES the last suffix, and the folder name ends in the
    # MinerU version ("...-3.4.4") — so ".4" would be thrown away and two
    # different versions would collide on one temp directory.
    staging = target.with_name(target.name + ".partial")
    shutil.rmtree(staging, ignore_errors=True)
    staging.mkdir(parents=True, exist_ok=True)

    pdf_bytes = read_fn(source)
    started = time.perf_counter()
    do_parse(
        output_dir=str(staging),
        pdf_file_names=[doc_id],
        pdf_bytes_list=[pdf_bytes],
        p_lang_list=[lang],
        backend=mineru_backend,
        formula_enable=cfg.formula_enable,
        table_enable=cfg.table_enable,
        f_draw_layout_bbox=cfg.write_debug_pdf,
        f_draw_span_bbox=cfg.write_debug_pdf,
        start_page_id=cfg.start_page,
        end_page_id=cfg.end_page,
        **({"effort": cfg.effort} if cfg.backend == "hybrid" else {}),
    )
    seconds = time.perf_counter() - started

    content, middle = _find_output(staging)
    if content is None:
        raise RuntimeError(
            f"MinerU finished but no *_content_list.json was found in {staging}"
        )

    pages = _count_pages(source, cfg)
    staging.rename(target)
    content = target / content.relative_to(staging)
    middle = target / middle.relative_to(staging) if middle else None

    marker.write_text(
        json.dumps({"pages": pages, "seconds": seconds, "doc_id": doc_id}),
        encoding="utf-8",
    )
    backend_recorded, version = _read_middle(middle)

    return ParseResult(
        doc_id=doc_id, source=source, cache_dir=target,
        content_list=content, middle_json=middle,
        pages=pages, seconds=seconds, from_cache=False,
        backend_recorded=backend_recorded, mineru_version=version,
    )


def _count_pages(source: Path, cfg: PipelineConfig) -> int:
    if cfg.end_page is not None:
        return cfg.end_page - cfg.start_page + 1
    try:
        from pypdf import PdfReader
        return len(PdfReader(str(source)).pages)
    except Exception:
        return 0

	"""Stage 2 — parse a document with MinerU, cached by content.

	Why the cache is keyed to document CONTENT rather than to a run id: parsing is
	by far the most expensive stage (19 s/page on CPU). A per-run cache would
	re-parse everything each time the normaliser changes and the pipeline is run
	again, for nothing. Keyed by content + settings + version, the same document is
	parsed once — however many times the normaliser is revised.

	Cache key = sha256(file bytes) + settings fingerprint + MinerU version. The
	version is in there so figures from an older MinerU never quietly mix with
	figures from a newer one — which matters, because those figures are used to
	justify spend.
	"""

	from __future__ import annotations

	import hashlib
	import json
	import shutil
	import time
	from dataclasses import dataclass
	from pathlib import Path

	from .config import PipelineConfig


	@dataclass
	class ParseResult:
	doc_id: str
	source: Path
	cache_dir: Path # folder holding MinerU's output untouched
	content_list: Path # *_content_list.json
	middle_json: Path \| None
	pages: int
	seconds: float
	from_cache: bool
	backend_recorded: str \| None # read from _middle.json, NOT from the config
	mineru_version: str \| None


	def mineru_version() -> str:
	try:
	from mineru.version import __version__ # type: ignore
	return str(__version__)
	except Exception:
	try:
	from importlib.metadata import version
	return version("mineru")
	except Exception:
	return "unknown"


	def file_hash(path: Path, length: int = 16) -> str:
	h = hashlib.sha256()
	with path.open("rb") as f:
	for block in iter(lambda: f.read(1 << 20), b""):
	h.update(block)
	return h.hexdigest()[:length]


	def cache_key(source: Path, cfg: PipelineConfig) -> str:
	return f"{file_hash(source)}-{cfg.fingerprint()}-{mineru_version()}"


	def _find_output(root: Path) -> tuple[Path \| None, Path \| None]:
	"""Locate content_list & middle json anywhere under root.

	Deliberately a glob rather than a fixed path: MinerU's subfolder layout
	differs per backend ('auto' for pipeline, something else for vlm). Globbing
	means this module needs no change when the backend does.
	"""
	content = next(iter(sorted(root.rglob("*_content_list.json"))), None)
	middle = next(iter(sorted(root.rglob("*_middle.json"))), None)
	return content, middle


	def _read_middle(middle: Path \| None) -> tuple[str \| None, str \| None]:
	"""The backend & version that ACTUALLY ran, taken from MinerU's own output."""
	if not middle or not middle.exists():
	return None, None
	try:
	d = json.loads(middle.read_text(encoding="utf-8"))
	return d.get("_backend"), d.get("_version_name")
	except Exception:
	return None, None


	def parse_document(source: Path, cfg: PipelineConfig) -> ParseResult:
	"""Parse one document. A cache hit is not re-run."""
	doc_id = source.stem
	target = cfg.cache_dir / "parse" / cache_key(source, cfg)
	marker = target / ".done"

	if marker.exists():
	content, middle = _find_output(target)
	if content:
	meta = json.loads(marker.read_text(encoding="utf-8"))
	backend_recorded, version = _read_middle(middle)
	return ParseResult(
	doc_id=doc_id, source=source, cache_dir=target,
	content_list=content, middle_json=middle,
	pages=meta.get("pages", 0), seconds=meta.get("seconds", 0.0),
	from_cache=True,
	backend_recorded=backend_recorded, mineru_version=version,
	)
	shutil.rmtree(target, ignore_errors=True) # cache corrupt, redo it

	from mineru.cli.common import do_parse, read_fn

	# MinerU dispatches on prefixes: "pipeline", then `startswith("vlm-")` and
	# `startswith("hybrid-")`. The bare names "vlm"/"hybrid" match NONE of those
	# branches — MinerU simply does nothing, without error, and the failure only
	# surfaces here as "content_list not found", which points at the wrong place
	# entirely. The short names stay in the config because they are part of the
	# cache key and read better; the translation happens here, right before
	# MinerU is touched.
	mineru_backend = {
	"pipeline": "pipeline",
	"vlm": "vlm-engine",
	"hybrid": "hybrid-engine",
	}[cfg.backend]

	# The playground does this before parsing; aliases like "latin"/"en" are
	# resolved to a canonical code. Matched here so results are identical.
	lang = cfg.canonical_lang()

	# with_suffix() REPLACES the last suffix, and the folder name ends in the
	# MinerU version ("...-3.4.4") — so ".4" would be thrown away and two
	# different versions would collide on one temp directory.
	staging = target.with_name(target.name + ".partial")
	shutil.rmtree(staging, ignore_errors=True)
	staging.mkdir(parents=True, exist_ok=True)

	pdf_bytes = read_fn(source)
	started = time.perf_counter()
	do_parse(
	output_dir=str(staging),
	pdf_file_names=[doc_id],
	pdf_bytes_list=[pdf_bytes],
	p_lang_list=[lang],
	backend=mineru_backend,
	formula_enable=cfg.formula_enable,
	table_enable=cfg.table_enable,
	f_draw_layout_bbox=cfg.write_debug_pdf,
	f_draw_span_bbox=cfg.write_debug_pdf,
	start_page_id=cfg.start_page,
	end_page_id=cfg.end_page,
	**({"effort": cfg.effort} if cfg.backend == "hybrid" else {}),
	)
	seconds = time.perf_counter() - started

	content, middle = _find_output(staging)
	if content is None:
	raise RuntimeError(
	f"MinerU finished but no *_content_list.json was found in {staging}"
	)

	pages = _count_pages(source, cfg)
	staging.rename(target)
	content = target / content.relative_to(staging)
	middle = target / middle.relative_to(staging) if middle else None

	marker.write_text(
	json.dumps({"pages": pages, "seconds": seconds, "doc_id": doc_id}),
	encoding="utf-8",
	)
	backend_recorded, version = _read_middle(middle)

	return ParseResult(
	doc_id=doc_id, source=source, cache_dir=target,
	content_list=content, middle_json=middle,
	pages=pages, seconds=seconds, from_cache=False,
	backend_recorded=backend_recorded, mineru_version=version,
	)


	def _count_pages(source: Path, cfg: PipelineConfig) -> int:
	if cfg.end_page is not None:
	return cfg.end_page - cfg.start_page + 1
	try:
	from pypdf import PdfReader
	return len(PdfReader(str(source)).pages)
	except Exception:
	return 0