from huggingface_hub import snapshot_download, hf_hub_download
import argparse

# set args for repo_id, local_dir, repo_type,

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Download a dataset or model from the Hugging Face Hub"
    )
    parser.add_argument(
        "--repo_id", type=str, help="The ID of the repository to download"
    )
    parser.add_argument(
        "--local_dir",
        type=str,
        help="The local directory to download the repository to",
    )
    parser.add_argument(
        "--repo_type",
        type=str,
        help="The type of repository to download (dataset or model)",
    )
    parser.add_argument("--file_name", type=str, help="The file name to download")
    args = parser.parse_args()
    if args.file_name:
        hf_hub_download(
            repo_id=args.repo_id,
            filename=args.file_name,
            repo_type=args.repo_type,
            local_dir=args.local_dir,
        )
    else:
        snapshot_download(
            repo_id=args.repo_id,
            local_dir=args.local_dir,
            repo_type=args.repo_type,
            local_dir_use_symlinks=False,
            resume_download=True,
        )