import json
from pathlib import Path

base_dir = Path(
    "/work/jf381/output/fst_353M_bert_update_2_6b_test/checkpoints/mrpc/lr_2e-05"
)

best = {
    "score": float("-inf"),
    "restart": None,
    "path": None,
}

for meta_path in base_dir.glob("restart_*/run_meta.json"):
    with open(meta_path) as f:
        meta = json.load(f)

    score = meta.get("best_score")
    restart = meta.get("restart")

    if score is not None and score > best["score"]:
        best.update({
            "score": score,
            "restart": restart,
            "path": str(meta_path.parent),
        })

print("Highest best_score:")
print(f"  score   : {best['score']}")
print(f"  restart : {best['restart']}")
print(f"  path    : {best['path']}")