File size: 2,424 Bytes
85a1f3d | 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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """Export a verified training checkpoint as inference-only Trackformer 1.2 weights."""
from __future__ import annotations
import argparse
import hashlib
import json
from pathlib import Path
import torch
EXPECTED_SHA256 = "f194a23d3f91ea76ad776dfad942fabd669eeae8b3fd665815463095367e9ee0"
EXPECTED_ARCH = "v173-moving-core-multiscale-attention"
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as stream:
for block in iter(lambda: stream.read(8 * 1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("checkpoint", type=Path)
parser.add_argument("output_dir", type=Path)
args = parser.parse_args()
if sha256(args.checkpoint) != EXPECTED_SHA256:
raise ValueError("Checkpoint SHA-256 does not match the selected release source")
payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
if (payload.get("version"), payload.get("architecture"), payload.get("epoch")) != (
"1.2.73", EXPECTED_ARCH, 4
):
raise ValueError("Checkpoint identity does not match the selected release source")
if len(payload.get("history", [])) != 4:
raise ValueError("Expected four completed epochs in checkpoint history")
args.output_dir.mkdir(parents=True, exist_ok=True)
weights = args.output_dir / "weights.pt"
torch.save(payload["model"], weights)
source = args.output_dir
manifest = {
"public_version": "1.2",
"release_status": "research_candidate_not_operational",
"source_checkpoint_version": payload["version"],
"source_checkpoint_epoch": payload["epoch"],
"architecture": payload["architecture"],
"source_checkpoint_sha256": EXPECTED_SHA256,
"inference_weights_sha256": sha256(weights),
"dataset_sha256": payload["dataset_sha256"],
"source_module_sha256": {name: sha256(source / name) for name in (
"model.py", "baseline_model.py", "v165_base.py")},
"data_contract": payload["data_contract"],
}
(args.output_dir / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
print(json.dumps({key: manifest[key] for key in (
"public_version", "architecture", "inference_weights_sha256")}, indent=2))
if __name__ == "__main__":
main()
|