File size: 6,333 Bytes
1b4ed07 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | # Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
"""Load a verified immutable Hub bundle and detect layout on RGB pages."""
from __future__ import annotations
import importlib
import shutil
from collections.abc import Callable
from pathlib import Path
from tempfile import TemporaryDirectory
from ax_doclayout.artifact import MODEL_ID, ArtifactReceipt
from ax_doclayout.contract import (
HUB_COMMIT_RE,
Element,
ModelIdentity,
PageImage,
PageStructure,
)
from ax_doclayout.errors import (
ArtifactError,
ArtifactUnavailableError,
AXDocLayoutError,
BackendUnavailableError,
InferenceError,
InvalidPageImageError,
InvalidRevisionError,
)
from ax_doclayout.hub_release import PINNED_HUB_RECEIPT, qualify_published_bundle
from ax_doclayout.postprocess import page_structure_from_elements
BACKENDS = ("pytorch", "onnx")
_PredictElements = Callable[[PageImage], tuple[Element, ...]]
def _download_bundle(directory: Path, *, receipt: ArtifactReceipt) -> None:
"""Materialize regular files from the Hub cache before strict verification."""
try:
from huggingface_hub import hf_hub_download
except ImportError:
raise BackendUnavailableError("huggingface-hub is not installed") from None
try:
for item in receipt.artifact_files:
cached = hf_hub_download(
repo_id=receipt.model_id,
filename=item.name,
revision=receipt.model_revision,
token=False,
)
# Hub cache entries may be symlinks. Only the copied, hash-verified
# regular files are handed to artifact validation and the backend.
shutil.copyfile(cached, directory / item.name)
except Exception: # noqa: BLE001 - redact third-party download error details
raise ArtifactUnavailableError("hub_bundle_download_failed") from None
def _load_backend(directory: Path, *, backend: str) -> _PredictElements:
"""Load all backend state before the temporary verified bundle is removed."""
if backend == "pytorch":
from ax_doclayout._local_pytorch import LocalPytorchBundle
return LocalPytorchBundle(directory).predict_elements
import onnxruntime as ort
from ax_doclayout._cuda_onnx import create_cuda_session
from ax_doclayout._local_onnx import LocalOnnxBundle
session = create_cuda_session(directory / "model.onnx", ort=ort)
runner = LocalOnnxBundle(
directory, session_run=lambda feeds: session.run(None, dict(feeds)),
)
def predict_elements(page: PageImage) -> tuple[Element, ...]:
return runner.predict(page).elements
return predict_elements
class AXDocLayout:
"""A trusted Hub artifact with explicit CUDA ONNX or CPU PyTorch inference."""
def __init__(self, model_id: str, revision: str, backend: str) -> None:
self.model_id = model_id
self.revision = revision
self.backend = backend
self._predict_elements: _PredictElements | None = None
self._identity: ModelIdentity | None = None
@classmethod
def from_pretrained(
cls,
model_id: str,
*,
revision: str,
backend: str | None = None,
) -> AXDocLayout:
if not isinstance(revision, str) or HUB_COMMIT_RE.fullmatch(revision) is None:
raise InvalidRevisionError(
"revision must be a 40-character immutable Hub commit; "
"'main', a tag, or 'latest' is not accepted"
)
if not isinstance(backend, str) or backend not in BACKENDS:
raise BackendUnavailableError(
"backend must be given explicitly as one of pytorch, onnx"
)
receipt = PINNED_HUB_RECEIPT
if not isinstance(model_id, str) or model_id != MODEL_ID:
raise ArtifactError("artifact_model_id_invalid")
if receipt is None or revision != receipt.model_revision:
raise ArtifactUnavailableError("no trusted artifact for this revision")
try:
dependencies = ("torch", "safetensors") if backend == "pytorch" else ("onnxruntime",)
for dependency in dependencies:
try:
importlib.import_module(dependency)
except (ImportError, OSError):
raise BackendUnavailableError("requested backend is not installed") from None
with TemporaryDirectory(prefix="ax-doclayout-hub-") as temporary:
directory = Path(temporary)
_download_bundle(directory, receipt=receipt)
qualified = qualify_published_bundle(
model_id=model_id, revision=revision, backend=backend,
receipt=receipt, bundle=directory,
)
if qualified is None:
raise ArtifactUnavailableError("no trusted artifact for this revision")
predictor = _load_backend(qualified, backend=backend)
except AXDocLayoutError:
raise
except Exception: # noqa: BLE001 - keep public errors content-free
raise ArtifactError("hub_backend_load_failed") from None
model = cls(model_id, revision, backend)
model._predict_elements = predictor
model._identity = ModelIdentity(
id=receipt.model_id, revision=receipt.model_revision,
artifact_sha256=receipt.artifact_sha256.removeprefix("sha256:"),
)
return model
def predict(self, page: PageImage) -> PageStructure:
if self._predict_elements is None or self._identity is None:
raise ArtifactUnavailableError("no loaded artifact; construct through from_pretrained")
if not isinstance(page, PageImage):
raise InvalidPageImageError("page must be a PageImage")
try:
elements = self._predict_elements(page)
return page_structure_from_elements(
page=page, elements=elements, identity=self._identity,
)
except AXDocLayoutError:
raise
except Exception: # noqa: BLE001 - keep public errors content-free
raise InferenceError("page_inference_failed") from None
|