File size: 1,119 Bytes
534cc76 | 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 | """Load an exported model for inference."""
import json
from pathlib import Path
from .checkpoint import load_checkpoint
from .model import PayaModel
from .provenance import RUNTIME_CONTRACT, contained_file, sha256
from .vendor import signal
__all__ = ["load_model"]
def load_model(directory, device="cpu"):
root = Path(directory).resolve()
manifest = json.loads((root/"manifest.json").read_text(encoding="utf-8"))
if manifest.get("contract") != RUNTIME_CONTRACT:
raise ValueError("Unsupported model architecture")
for name, expected in manifest["files"].items():
if sha256(contained_file(root, name)) != expected:
raise ValueError(f"Model file checksum changed: {name}")
config = json.loads((root/"config.json").read_text(encoding="utf-8"))
assets = {name: root/name for name in (signal.ASSET_KEY, "text")}
model = PayaModel.from_assets(assets, config["preparation"], config["head"], device)
if manifest["identity"] != model.identity():
raise ValueError("Model identities differ")
load_checkpoint(root/"decision.pt", model)
return model
|