48 lines
1.7 KiB
Python
48 lines
1.7 KiB
Python
"""Explicitly add a reviewed PyTorch model hash to the local trust manifest."""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from app.utils.ModelTrust import MANIFEST_PATH, sha256_file
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Trust a reviewed local PyTorch model by SHA-256")
|
|
parser.add_argument("model", help="path to the reviewed .pt file")
|
|
parser.add_argument("--yes", action="store_true", help="confirm that the model source was reviewed")
|
|
args = parser.parse_args()
|
|
model = Path(args.model).resolve()
|
|
if model.suffix.lower() != ".pt" or not model.is_file():
|
|
parser.error("model must be an existing .pt file")
|
|
digest = sha256_file(model)
|
|
print("model:", model)
|
|
print("sha256:", digest)
|
|
if not args.yes:
|
|
print("No change made. Re-run with --yes after verifying the source and hash.")
|
|
return 2
|
|
data = {"trusted_sha256": []}
|
|
if MANIFEST_PATH.is_file():
|
|
with open(MANIFEST_PATH, "r", encoding="utf-8") as handle:
|
|
loaded = json.load(handle)
|
|
if isinstance(loaded, dict):
|
|
data = loaded
|
|
hashes = {str(item).lower() for item in data.get("trusted_sha256", [])}
|
|
hashes.add(digest)
|
|
data["trusted_sha256"] = sorted(hashes)
|
|
temp = MANIFEST_PATH.with_suffix(MANIFEST_PATH.suffix + ".tmp")
|
|
with open(temp, "w", encoding="utf-8") as handle:
|
|
json.dump(data, handle, indent=2)
|
|
handle.write("\n")
|
|
os.replace(temp, MANIFEST_PATH)
|
|
print("updated:", MANIFEST_PATH)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|