| """Export models with verified weights and encoder assets.""" | |
| import argparse | |
| import json | |
| from pathlib import Path | |
| import shutil | |
| import yaml | |
| from .checkpoint import load_checkpoint, save_checkpoint | |
| from .model import PayaModel | |
| from .provenance import RUNTIME_CONTRACT, sha256, verify_assets | |
| def export_model(model, assets, output): | |
| root = Path(output).resolve() | |
| if root.exists() and any(root.iterdir()): | |
| raise FileExistsError("Export destination must be empty") | |
| if any(root.is_relative_to(Path(path).resolve()) for path in assets.values()): | |
| raise ValueError("Export cannot be nested inside encoder assets") | |
| if verify_assets(assets) != model.asset_identity: | |
| raise ValueError("Export assets do not match the trained model") | |
| root.mkdir(parents=True, exist_ok=True) | |
| for name, path in assets.items(): | |
| source = Path(path).resolve() | |
| for entry in source.rglob("*"): | |
| if not entry.resolve().is_relative_to(source): | |
| raise ValueError("Asset tree contains an escaping link") | |
| shutil.copytree(source, root/name, ignore=shutil.ignore_patterns(".cache", "__pycache__")) | |
| save_checkpoint(root/"decision.pt", model) | |
| config = {"preparation": model.identity()["preparation"]["config"], "head": model.head.config} | |
| (root/"config.json").write_text(json.dumps(config, indent=2), encoding="utf-8") | |
| manifest = {"contract": RUNTIME_CONTRACT, "identity": model.identity(), | |
| "calibrated": model.calibration is not None, | |
| "files": {path.relative_to(root).as_posix(): sha256(path) | |
| for path in root.rglob("*") if path.is_file()}} | |
| (root/"manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8") | |
| return manifest | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--config", type=Path, default=Path("config/paya.yaml")) | |
| parser.add_argument("--checkpoint", type=Path, required=True) | |
| parser.add_argument("--output", type=Path, required=True) | |
| args = parser.parse_args() | |
| config = yaml.safe_load(args.config.read_text(encoding="utf-8")) | |
| model = PayaModel.from_assets(config["assets"], config.get("preparation"), config.get("head")) | |
| load_checkpoint(args.checkpoint, model) | |
| print(json.dumps(export_model(model, config["assets"], args.output), indent=2)[:1000]) | |