repro-bundle/download_checkpoints.py

import os
from huggingface_hub import hf_hub_download
import shutil

def main():
    repo_id = "RL-MIND/OrthoReg_checkpoints"
    model = "ViT-B-32"
    base_dir = "checkpoints"
    
    datasets = ['Cars', 'DTD', 'EuroSAT', 'GTSRB', 'MNIST', 'RESISC45', 'SUN397', 'SVHN']
    
    # Download zeroshot.pt
    print("Downloading zeroshot.pt...")
    zeroshot_src = f"standard_1e-05_{model}/CarsVal/zeroshot.pt"
    zeroshot_dest = os.path.join(base_dir, model, "zeroshot.pt")
    os.makedirs(os.path.dirname(zeroshot_dest), exist_ok=True)
    
    path = hf_hub_download(repo_id=repo_id, filename=zeroshot_src)
    shutil.copy(path, zeroshot_dest)
    print(f"Saved zeroshot.pt to {zeroshot_dest}")
    
    # Download finetuned.pt and head.pt for each dataset
    for dataset in datasets:
        print(f"\nProcessing {dataset}...")
        
        # finetuned.pt
        ft_src = f"standard_1e-05_{model}/{dataset}Val/finetuned.pt"
        ft_dest = os.path.join(base_dir, model, dataset, "finetuned.pt")
        os.makedirs(os.path.dirname(ft_dest), exist_ok=True)
        
        print(f"Downloading finetuned.pt for {dataset}...")
        path_ft = hf_hub_download(repo_id=repo_id, filename=ft_src)
        shutil.copy(path_ft, ft_dest)
        print(f"Saved finetuned.pt to {ft_dest}")
        
        # head.pt
        head_src = f"standard_1e-05_{model}/head_{dataset}Val.pt"
        head_dest = os.path.join(base_dir, model, f"head_{dataset}.pt")
        
        print(f"Downloading classification head for {dataset}...")
        path_head = hf_hub_download(repo_id=repo_id, filename=head_src)
        shutil.copy(path_head, head_dest)
        print(f"Saved head to {head_dest}")

    print("\nAll downloads completed successfully!")

if __name__ == "__main__":
    main()