| |
| |
| |
| """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 |
|
|
