File size: 1,529 Bytes
69928d6 | 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 | """Validated backbone handoff between the two Sudoku training stages."""
from pathlib import Path
import torch
from softasp.sudoku_reporting import file_sha256
def extract_backbone(checkpoint: Path, destination: Path) -> dict:
"""Extract a newly trained perception checkpoint, retaining source identity."""
saved = torch.load(checkpoint, map_location="cpu", weights_only=False)
if saved.get("config", {}).get("model", {}).get("name") != "perception_only_sudoku":
raise ValueError("Expected a perception_only_sudoku training checkpoint")
if saved.get("epoch", -1) < 0:
raise ValueError("Backbone checkpoint must follow at least one training epoch")
state = {
key: value
for key, value in saved["model_state_dict"].items()
if key.startswith("cell_cnn.")
}
if not state or any(not torch.isfinite(value).all() for value in state.values()):
raise ValueError("Expected finite trained cell_cnn tensors")
provenance = {
"source_checkpoint": str(checkpoint.resolve()),
"source_sha256": file_sha256(checkpoint),
"source_epoch": saved["epoch"],
"source_metric": saved.get("best_metric"),
"dataset_identity": saved["config"].get("resolved_dataset"),
}
# Refuse to overwrite any pre-existing model artifact.
with destination.open("xb") as stream:
torch.save({"model_state_dict": state, **provenance}, stream)
return {**provenance, "path": str(destination), "sha256": file_sha256(destination)}
|