Spaces:
Running on Zero
Running on Zero
| """Strict verification and loading for the pinned official checkpoint.""" | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| from pathlib import Path | |
| from typing import Any, Mapping | |
| import torch | |
| from torch import nn | |
| from .architecture import model_parameter_count | |
| from .config import CHECKPOINT_SHA256 | |
| from .schemas import CheckpointVerificationReport | |
| class CheckpointError(RuntimeError): | |
| """Base class for safe checkpoint failures.""" | |
| class CheckpointVerificationError(CheckpointError): | |
| pass | |
| class CheckpointCompatibilityError(CheckpointError): | |
| pass | |
| def sha256_file(path: str | Path, chunk_size: int = 1024 * 1024) -> str: | |
| digest = hashlib.sha256() | |
| with Path(path).open("rb") as handle: | |
| while chunk := handle.read(chunk_size): | |
| digest.update(chunk) | |
| return digest.hexdigest() | |
| def _bounded_metadata(value: Any, *, depth: int = 0) -> Any: | |
| """Produce a JSON-safe, bounded metadata representation.""" | |
| if depth > 4: | |
| return "<maximum-depth>" | |
| if value is None or isinstance(value, (bool, int, float)): | |
| return value | |
| if isinstance(value, str): | |
| return value if len(value) <= 2000 else value[:2000] + "…" | |
| if isinstance(value, Mapping): | |
| items = list(value.items())[:100] | |
| result = {str(key): _bounded_metadata(item, depth=depth + 1) for key, item in items} | |
| if len(value) > 100: | |
| result["<truncated-items>"] = len(value) - 100 | |
| return result | |
| if isinstance(value, (list, tuple)): | |
| result = [_bounded_metadata(item, depth=depth + 1) for item in value[:100]] | |
| if len(value) > 100: | |
| result.append(f"<truncated-items:{len(value) - 100}>") | |
| return result | |
| if isinstance(value, torch.Tensor): | |
| return {"tensor_shape": list(value.shape), "dtype": str(value.dtype)} | |
| text = repr(value) | |
| return text if len(text) <= 500 else text[:500] + "…" | |
| def _load_verified_payload(path: Path, expected_sha256: str) -> tuple[dict[str, Any], int, str]: | |
| if not path.is_file(): | |
| raise CheckpointVerificationError(f"Checkpoint file does not exist: {path}") | |
| file_size = path.stat().st_size | |
| digest = sha256_file(path) | |
| if digest.lower() != expected_sha256.lower(): | |
| raise CheckpointVerificationError( | |
| f"Checkpoint SHA-256 mismatch: expected {expected_sha256}, received {digest}" | |
| ) | |
| try: | |
| payload = torch.load(path, map_location="cpu", weights_only=False) | |
| except Exception as error: | |
| raise CheckpointVerificationError(f"Verified checkpoint could not be decoded: {error}") from error | |
| if not isinstance(payload, dict): | |
| raise CheckpointVerificationError("Checkpoint top level must be a dictionary") | |
| return payload, file_size, digest | |
| def inspect_checkpoint( | |
| path: str | Path, *, expected_sha256: str = CHECKPOINT_SHA256 | |
| ) -> tuple[dict[str, torch.Tensor], CheckpointVerificationReport]: | |
| checkpoint_path = Path(path).expanduser().resolve() | |
| payload, file_size, digest = _load_verified_payload(checkpoint_path, expected_sha256) | |
| state_dict = payload.get("state_dict") | |
| if not isinstance(state_dict, Mapping) or not state_dict: | |
| raise CheckpointVerificationError("Checkpoint does not contain a non-empty 'state_dict'") | |
| if not all(isinstance(key, str) and isinstance(value, torch.Tensor) for key, value in state_dict.items()): | |
| raise CheckpointVerificationError("Checkpoint state_dict contains non-tensor entries") | |
| # The official file already uses exact model-relative keys. No prefix | |
| # normalization is performed; any mismatch is an incompatibility. | |
| normalized = dict(state_dict) | |
| metadata = _bounded_metadata(payload.get("meta", {})) | |
| try: | |
| json.dumps(metadata) | |
| except TypeError as error: # defensive invariant for diagnostics output | |
| raise CheckpointVerificationError("Checkpoint metadata is not serializable") from error | |
| report = CheckpointVerificationReport( | |
| path=str(checkpoint_path), | |
| file_size=file_size, | |
| sha256=digest, | |
| top_level_keys=tuple(sorted(str(key) for key in payload)), | |
| tensor_count=len(normalized), | |
| parameter_count=0, | |
| missing_keys=(), | |
| unexpected_keys=(), | |
| key_transformations=(), | |
| checkpoint_metadata=metadata, | |
| ) | |
| return normalized, report | |
| def load_checkpoint_strict( | |
| model: nn.Module, | |
| path: str | Path, | |
| *, | |
| expected_sha256: str = CHECKPOINT_SHA256, | |
| ) -> CheckpointVerificationReport: | |
| state_dict, report = inspect_checkpoint(path, expected_sha256=expected_sha256) | |
| try: | |
| incompatible = model.load_state_dict(state_dict, strict=True) | |
| except RuntimeError as error: | |
| raise CheckpointCompatibilityError(f"Strict state-dict loading failed: {error}") from error | |
| missing = tuple(incompatible.missing_keys) | |
| unexpected = tuple(incompatible.unexpected_keys) | |
| if missing or unexpected: | |
| raise CheckpointCompatibilityError( | |
| f"Strict state-dict loading reported missing={missing}, unexpected={unexpected}" | |
| ) | |
| return CheckpointVerificationReport( | |
| path=report.path, | |
| file_size=report.file_size, | |
| sha256=report.sha256, | |
| top_level_keys=report.top_level_keys, | |
| tensor_count=report.tensor_count, | |
| parameter_count=model_parameter_count(model), | |
| missing_keys=missing, | |
| unexpected_keys=unexpected, | |
| key_transformations=report.key_transformations, | |
| checkpoint_metadata=report.checkpoint_metadata, | |
| ) | |