File size: 4,336 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 | # Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
"""Digest-verified PyTorch inference over a local nine-file bundle.
This path loads the staged safetensors into the in-repo RT-DETRv2-M
architecture. It is not ``from_pretrained`` and it never marks a bundle
publishable.
"""
from __future__ import annotations
from pathlib import Path
import torch
from ax_doclayout._torch_rtdetrv2 import build_rtdetrv2_r50vd_m
from ax_doclayout.artifact import (
MODEL_ID,
STAGED_SAFETENSORS_SHA256,
_reject_json_constant,
_unique_object,
verify_artifact_manifest,
)
from ax_doclayout.contract import Element, PageImage
from ax_doclayout.errors import ArtifactError, BackendUnavailableError
from ax_doclayout.postprocess import (
DEFAULT_SCORE_THRESHOLD,
elements_from_letterbox_detections,
)
from ax_doclayout.preprocess import letterbox_nchw, letterbox_page, orig_target_sizes
ARCHITECTURE_ID = "rtdetrv2_r50vd_m"
class LocalPytorchBundle:
"""Run the staged safetensors after the manifest binds every payload digest."""
def __init__(
self,
directory: Path,
*,
expected_safetensors_sha256: str | None = STAGED_SAFETENSORS_SHA256,
) -> None:
try:
from safetensors import SafetensorError
from safetensors.torch import load_file
except ImportError as error:
raise BackendUnavailableError("safetensors is not installed") from error
manifest = verify_artifact_manifest(directory)
if manifest.model_id != MODEL_ID:
raise ArtifactError("artifact_model_id_invalid")
weights_digest = ""
for item in manifest.artifact_files:
if item.name == "model.safetensors":
weights_digest = item.sha256
break
if not weights_digest:
raise ArtifactError("artifact_file_unavailable")
if (
expected_safetensors_sha256 is not None
and weights_digest != expected_safetensors_sha256
):
raise ArtifactError("artifact_safetensors_digest_mismatch")
weights_path = directory / "model.safetensors"
if not weights_path.is_file():
raise ArtifactError("artifact_file_unavailable")
try:
state = load_file(str(weights_path))
except (SafetensorError, OSError, ValueError) as error:
raise ArtifactError("artifact_safetensors_unreadable") from error
model = build_rtdetrv2_r50vd_m()
try:
model.load_state_dict(state, strict=True)
except RuntimeError as error:
raise ArtifactError("artifact_state_dict_mismatch") from error
model.eval()
self._model = model
self._directory = directory
def predict_elements(
self,
page: PageImage,
*,
threshold: float = DEFAULT_SCORE_THRESHOLD,
) -> tuple[Element, ...]:
"""Detect on one RGB page and return original-raster elements."""
letterbox = letterbox_page(page)
images = torch.from_numpy(letterbox_nchw(letterbox))
sizes = torch.from_numpy(orig_target_sizes(letterbox))
with torch.no_grad():
labels, boxes, scores = self._model(images, sizes)
return elements_from_letterbox_detections(
labels=labels.cpu().numpy(),
boxes=boxes.cpu().numpy(),
scores=scores.cpu().numpy(),
page=page,
letterbox=letterbox,
threshold=threshold,
)
def read_publishable(path: Path) -> bool:
"""Read the staging receipt flag. Missing or malformed receipts fail closed."""
import json
try:
payload = json.loads(
path.read_text(encoding="utf-8"),
object_pairs_hook=_unique_object,
parse_constant=_reject_json_constant,
)
except (OSError, ValueError) as error:
raise ArtifactError("staging_receipt_unreadable") from error
if not isinstance(payload, dict) or "publishable" not in payload:
raise ArtifactError("staging_publishable_missing")
flag = payload["publishable"]
if not isinstance(flag, bool):
raise ArtifactError("staging_publishable_invalid")
return flag
|