| """Exact architecture and integrity-checked loader for YellowCab v0.""" |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from collections import Counter |
| from pathlib import Path |
| from typing import Any, Mapping |
|
|
| import torch |
| from safetensors.torch import load_file |
| from torch import nn |
| from torchvision.models import efficientnet_b0 |
|
|
|
|
| class CheckpointIntegrityError(RuntimeError): |
| """Raised when a release checkpoint does not match its signed-off contract.""" |
|
|
|
|
| class TemporalFusionHead(nn.Module): |
| """The trained temporal and telemetry fusion head.""" |
|
|
| def __init__( |
| self, |
| *, |
| image_feature_dim: int, |
| telemetry_dim: int, |
| num_classes: int, |
| hidden_dim: int = 256, |
| dropout: float = 0.25, |
| ) -> None: |
| super().__init__() |
| self.image_feature_dim = image_feature_dim |
| self.image_norm = nn.LayerNorm(image_feature_dim) |
| self.temporal = nn.GRU( |
| input_size=image_feature_dim, |
| hidden_size=hidden_dim, |
| batch_first=True, |
| ) |
| self.telemetry = nn.Sequential( |
| nn.LayerNorm(telemetry_dim), |
| nn.Linear(telemetry_dim, 96), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(96, 96), |
| nn.GELU(), |
| ) |
| self.classifier = nn.Sequential( |
| nn.LayerNorm(hidden_dim + 96), |
| nn.Linear(hidden_dim + 96, hidden_dim), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(hidden_dim, num_classes), |
| ) |
|
|
| def forward( |
| self, |
| image_features: torch.Tensor, |
| telemetry: torch.Tensor, |
| ) -> torch.Tensor: |
| temporal_output, _ = self.temporal(self.image_norm(image_features)) |
| telemetry_output = self.telemetry(telemetry) |
| fused = torch.cat((temporal_output[:, -1], telemetry_output), dim=-1) |
| return self.classifier(fused) |
|
|
|
|
| class TaxiManeuverModel(nn.Module): |
| """Frozen EfficientNet-B0 plus the trained v0 fusion head.""" |
|
|
| def __init__(self, head_config: Mapping[str, Any]) -> None: |
| super().__init__() |
| self.encoder = efficientnet_b0(weights=None) |
| self.encoder.classifier = nn.Identity() |
| self.head = TemporalFusionHead(**dict(head_config)) |
|
|
| def forward( |
| self, |
| images: torch.Tensor, |
| telemetry: torch.Tensor, |
| ) -> torch.Tensor: |
| if images.ndim != 5 or images.shape[1] != 3 or images.shape[2] != 3: |
| raise ValueError("images must have shape [batch, 3, 3, height, width]") |
| batch_size, frame_count, channels, height, width = images.shape |
| flat_images = images.reshape( |
| batch_size * frame_count, |
| channels, |
| height, |
| width, |
| ) |
| image_features = self.encoder(flat_images).reshape( |
| batch_size, |
| frame_count, |
| -1, |
| ) |
| return self.head(image_features, telemetry) |
|
|
|
|
| def load_release_config(repository_path: str | Path) -> dict[str, Any]: |
| """Load and minimally validate the one canonical release configuration.""" |
|
|
| root = Path(repository_path) |
| path = root / "config.json" |
| try: |
| config = json.loads(path.read_text(encoding="utf-8")) |
| except (OSError, json.JSONDecodeError) as exc: |
| raise CheckpointIntegrityError("release configuration is unavailable") from exc |
| if config.get("schema") != "gdc_taxi_maneuver_release_config_v1": |
| raise CheckpointIntegrityError("unsupported release configuration schema") |
| for identity_field in ("model_id", "model_name", "model_version"): |
| identity_value = config.get(identity_field) |
| if not isinstance(identity_value, str) or not identity_value.strip(): |
| raise CheckpointIntegrityError( |
| f"release configuration {identity_field} is invalid" |
| ) |
| labels = config.get("labels") |
| feature_order = config.get("telemetry", {}).get("feature_order") |
| if not isinstance(labels, list) or len(labels) != 5 or len(set(labels)) != 5: |
| raise CheckpointIntegrityError("canonical label configuration is invalid") |
| if ( |
| not isinstance(feature_order, list) |
| or len(feature_order) != 12 |
| or len(set(feature_order)) != 12 |
| ): |
| raise CheckpointIntegrityError("canonical telemetry configuration is invalid") |
| return config |
|
|
|
|
| def sha256_file(path: str | Path) -> str: |
| digest = hashlib.sha256() |
| with Path(path).open("rb") as handle: |
| for chunk in iter(lambda: handle.read(1024 * 1024), b""): |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def _verify_checkpoint_metadata( |
| *, |
| checkpoint_path: Path, |
| weights: Mapping[str, torch.Tensor], |
| model: nn.Module, |
| checkpoint_config: Mapping[str, Any], |
| ) -> None: |
| expected_bytes = int(checkpoint_config["bytes"]) |
| expected_sha256 = str(checkpoint_config["sha256"]) |
| if checkpoint_path.stat().st_size != expected_bytes: |
| raise CheckpointIntegrityError("checkpoint byte size does not match config.json") |
| if sha256_file(checkpoint_path) != expected_sha256: |
| raise CheckpointIntegrityError("checkpoint SHA-256 does not match config.json") |
|
|
| expected_state = model.state_dict() |
| actual_keys = set(weights) |
| expected_keys = set(expected_state) |
| missing = sorted(expected_keys - actual_keys) |
| unexpected = sorted(actual_keys - expected_keys) |
| if missing or unexpected: |
| raise CheckpointIntegrityError( |
| "checkpoint tensor keys do not exactly match the model architecture" |
| ) |
|
|
| for name, expected_tensor in expected_state.items(): |
| actual_tensor = weights[name] |
| if actual_tensor.shape != expected_tensor.shape: |
| raise CheckpointIntegrityError( |
| f"checkpoint tensor shape mismatch for {name}" |
| ) |
| if actual_tensor.dtype != expected_tensor.dtype: |
| raise CheckpointIntegrityError( |
| f"checkpoint tensor dtype mismatch for {name}" |
| ) |
|
|
| actual_tensor_count = len(weights) |
| actual_state_numel = sum(tensor.numel() for tensor in weights.values()) |
| actual_dtype_counts = Counter(str(tensor.dtype) for tensor in weights.values()) |
| expected_dtype_counts = { |
| str(name): int(count) |
| for name, count in checkpoint_config["state_dtypes"].items() |
| } |
| if actual_tensor_count != int(checkpoint_config["state_tensor_count"]): |
| raise CheckpointIntegrityError("checkpoint tensor count does not match config.json") |
| if actual_state_numel != int(checkpoint_config["state_numel"]): |
| raise CheckpointIntegrityError("checkpoint state size does not match config.json") |
| if dict(actual_dtype_counts) != expected_dtype_counts: |
| raise CheckpointIntegrityError("checkpoint dtype inventory does not match config.json") |
|
|
| parameter_count = sum(parameter.numel() for parameter in model.parameters()) |
| if parameter_count != int(checkpoint_config["parameter_count"]): |
| raise CheckpointIntegrityError( |
| "model parameter count does not match config.json" |
| ) |
|
|
|
|
| def load_model( |
| repository_path: str | Path, |
| *, |
| device: torch.device | str = "cpu", |
| ) -> tuple[TaxiManeuverModel, dict[str, Any]]: |
| """Load the checkpoint only after hash, keys, shapes and dtypes all match.""" |
|
|
| root = Path(repository_path) |
| config = load_release_config(root) |
| checkpoint_config = config["checkpoint"] |
| checkpoint_path = root / str(checkpoint_config["file"]) |
| try: |
| weights = load_file(str(checkpoint_path), device="cpu") |
| except (OSError, ValueError) as exc: |
| raise CheckpointIntegrityError("checkpoint cannot be read") from exc |
|
|
| model = TaxiManeuverModel(config["architecture"]["head_config"]) |
| _verify_checkpoint_metadata( |
| checkpoint_path=checkpoint_path, |
| weights=weights, |
| model=model, |
| checkpoint_config=checkpoint_config, |
| ) |
| try: |
| model.load_state_dict(weights, strict=True) |
| except RuntimeError as exc: |
| raise CheckpointIntegrityError("strict checkpoint loading failed") from exc |
| model.eval() |
| model.requires_grad_(False) |
| model.to(torch.device(device)) |
| return model, config |
|
|
|
|
| __all__ = [ |
| "CheckpointIntegrityError", |
| "TaxiManeuverModel", |
| "TemporalFusionHead", |
| "load_model", |
| "load_release_config", |
| "sha256_file", |
| ] |
|
|
