"""Upload training checkpoints to HuggingFace Hub."""
import os
import glob
from huggingface_hub import HfApi

api = HfApi()
ckpts = glob.glob("lightning_logs/**/checkpoints/*.ckpt", recursive=True)
ckpts += glob.glob("checkpoints/**/*.ckpt", recursive=True)
print(f"Found {len(ckpts)} checkpoints")

for ckpt in ckpts:
    dest = f"checkpoints/{os.path.basename(ckpt)}"
    print(f"  Uploading {ckpt} -> {dest}")
    api.upload_file(
        path_or_fileobj=ckpt,
        path_in_repo=dest,
        repo_id="icarus112/sem-v6-training",
    )
    print(f"  Done")

print("All checkpoints uploaded")