| |
| |
| """Content-bound loading of the published mixed MXFP8 checkpoint.""" |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| from safetensors import safe_open |
|
|
| from ax_doclayout_mlx._identity import EXPECTED_FILES, EXPECTED_TENSORS, MODEL_ID |
| from ax_doclayout_mlx.errors import ArtifactError |
|
|
|
|
| def validate_bundle(directory: Path) -> dict: |
| """Verify every required file, then the exact safetensors inventory.""" |
| try: |
| root = Path(directory).expanduser() |
| if not root.is_dir() or root.is_symlink(): |
| raise ArtifactError("regular_bundle_directory_required") |
| for name, expected in EXPECTED_FILES.items(): |
| path = root / name |
| if path.is_symlink() or not path.is_file(): |
| raise ArtifactError("regular_bundle_files_required") |
| with path.open("rb") as stream: |
| actual = hashlib.file_digest(stream, "sha256").hexdigest() |
| if actual != expected: |
| raise ArtifactError("bundle_checksum_mismatch") |
| manifest = json.loads((root / "manifest.json").read_text()) |
| if (manifest["model_id"] != MODEL_ID or manifest["mode"] != "mxfp8" |
| or manifest["group_size"] != 32 |
| or manifest["convolution_precision"] != "float16" |
| or manifest["convolutions_quantized"] is not False): |
| raise ArtifactError("bundle_precision_contract_mismatch") |
| with safe_open(root / "model.safetensors", framework="np") as weights: |
| if set(weights.keys()) != set(EXPECTED_TENSORS): |
| raise ArtifactError("checkpoint_inventory_mismatch") |
| for name, expected in EXPECTED_TENSORS.items(): |
| item = weights.get_slice(name) |
| if (list(item.get_shape()) != expected["shape"] |
| or item.get_dtype() != expected["dtype"]): |
| raise ArtifactError("checkpoint_tensor_contract_mismatch") |
| return manifest |
| except ArtifactError: |
| raise |
| except Exception: |
| raise ArtifactError("bundle_validation_failed") from None |
|
|
