[content] slurm/utils/mihir/upload_hf_dataset.py (python)
#!/usr/bin/env python3
"""
Upload a folder to a Hugging Face repository.

This script efficiently uploads large folders to Hugging Face using
the huggingface_hub library with support for progress tracking and
resume capabilities.

Usage:
    python upload_hf_dataset.py <folder_path> <repo_name> [--repo-type TYPE] [--private]
    
Example:
    python upload_hf_dataset.py ./my_dataset p-doom/my-dataset --repo-type dataset
"""

import argparse
import os
import sys
from pathlib import Path

try:
    from huggingface_hub import HfApi, upload_folder
    from huggingface_hub.utils import HfHubHTTPError
except ImportError:
    print("Error: huggingface_hub is not installed.")
    print("Please install it with: pip install huggingface_hub")
    sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        description="Upload a folder to a Hugging Face repository",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Upload a dataset folder
  python upload_hf_dataset.py ./my_data p-doom/my-dataset --repo-type dataset
  
  # Upload a model folder (default)
  python upload_hf_dataset.py ./checkpoints p-doom/my-model
  
  # Upload as private repository
  python upload_hf_dataset.py ./my_data p-doom/my-dataset --repo-type dataset --private
        """
    )
    
    parser.add_argument(
        "folder",
        type=str,
        help="Path to the folder to upload"
    )
    
    parser.add_argument(
        "repo_name",
        type=str,
        help="Hugging Face repository name (e.g., 'username/repo-name')"
    )
    
    parser.add_argument(
        "--repo-type",
        type=str,
        default="model",
        choices=["model", "dataset", "space"],
        help="Type of repository (default: model)"
    )
    
    parser.add_argument(
        "--private",
        action="store_true",
        help="Make the repository private"
    )
    
    parser.add_argument(
        "--ignore-patterns",
        type=str,
        nargs="*",
        default=[".git", "__pycache__", "*.pyc", ".DS_Store", ".gitignore"],
        help="Patterns to ignore when uploading (default: .git, __pycache__, *.pyc, .DS_Store, .gitignore)"
    )
    
    parser.add_argument(
        "--commit-message",
        type=str,
        default="Upload folder",
        help="Commit message for the upload (default: 'Upload folder')"
    )
    
    args = parser.parse_args()
    
    # Validate folder path
    folder_path = Path(args.folder).expanduser().resolve()
    
    if not folder_path.exists():
        print(f"Error: Folder '{folder_path}' does not exist.")
        sys.exit(1)
    
    if not folder_path.is_dir():
        print(f"Error: '{folder_path}' is not a directory.")
        sys.exit(1)
    
    # Validate repo name format
    if "/" not in args.repo_name:
        print(f"Error: Repository name must be in format 'username/repo-name', got '{args.repo_name}'")
        sys.exit(1)
    
    print(f"Uploading folder: {folder_path}")
    print(f"Repository: {args.repo_name}")
    print(f"Repository type: {args.repo_type}")
    print(f"Private: {args.private}")
    print("-" * 60)
    
    # Check if HF token is available
    api = HfApi()
    try:
        # Test authentication
        whoami = api.whoami()
        print(f"Authenticated as: {whoami.get('name', 'unknown')}")
    except Exception as e:
        print("Error: Not authenticated with Hugging Face.")
        print("Please log in using: huggingface-cli login")
        print(f"Or set HUGGING_FACE_HUB_TOKEN environment variable.")
        sys.exit(1)
    
    # Create repository if it doesn't exist
    try:
        api.create_repo(
            repo_id=args.repo_name,
            repo_type=args.repo_type,
            private=args.private,
            exist_ok=True
        )
        print(f"Repository '{args.repo_name}' is ready.")
    except HfHubHTTPError as e:
        print(f"Error creating repository: {e}")
        sys.exit(1)
    
    # Upload the folder
    try:
        print("
Starting upload...")
        print("This may take a while for large folders...")
        
        result = upload_folder(
            folder_path=str(folder_path),
            repo_id=args.repo_name,
            repo_type=args.repo_type,
            ignore_patterns=args.ignore_patterns,
            commit_message=args.commit_message,
            multi_commits=True,  # Enable multi-commit for large uploads
            multi_commits_verbose=True,
        )
        
        print("
" + "=" * 60)
        print("Upload completed successfully!")
        print(f"Repository URL: https://huggingface.co/{args.repo_name}")
        if args.repo_type == "dataset":
            print(f"Dataset URL: https://huggingface.co/datasets/{args.repo_name}")
        print("=" * 60)
        
    except KeyboardInterrupt:
        print("

Upload interrupted by user.")
        print("You can resume the upload by running the same command again.")
        sys.exit(130)
    except Exception as e:
        print(f"
Error during upload: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()