import json, glob
from loadlib import load, sha
out = {}
for n in sorted(glob.glob("ckpt/confirm-*")):
    n = n.split("/")[-1]
    md = json.load(open(glob.glob(f"ckpt/{n}/*/metadata.json")[0]))
    out[n] = {"adapter_bf16_sha256": sha(load(n, prefix="language_model.")), "fp32_param_sha256": sha(load(n)),
              "metadata": {k: md.get(k) for k in ("base_model", "optimizer", "optimizer_config", "lora_config", "has_optimizer", "checkpoint_format")}}
json.dump(out, open("restore_check.json", "w"), indent=1)
for k, v in out.items(): print(k, v["adapter_bf16_sha256"][:16], v["fp32_param_sha256"][:16], v["metadata"]["optimizer"])