upload_data_to_hf.py

#!/usr/bin/env python3
"""
Upload data folder to Hugging Face Hub.

This script:
1. Creates a compressed tar.gz archive of the data folder
2. Uploads it to Hugging Face Hub as a dataset

Usage:
    python upload_data_to_hf.py --repo-id <username/dataset-name> [--data-dir data] [--skip-archive]

Requirements:
    pip install huggingface_hub
"""

import os
import sys
import argparse
import tarfile
import tempfile
from pathlib import Path
from huggingface_hub import HfApi, login
from tqdm import tqdm


def get_git_root() -> Path:
    """Get the root directory of the git repository."""
    import subprocess
    try:
        return Path(
            subprocess.check_output(
                ["git", "rev-parse", "--show-toplevel"],
                stderr=subprocess.DEVNULL,
                text=True
            ).strip()
        )
    except (subprocess.CalledProcessError, FileNotFoundError):
        # Fallback: use current working directory
        return Path.cwd()


def create_data_archive(data_dir: Path, output_path: Path) -> None:
    """
    Create a compressed tar.gz archive of the data directory.
    
    Args:
        data_dir: Path to the data directory
        output_path: Path where the archive will be saved
    """
    print(f"Creating archive of {data_dir}...")
    print(f"Output: {output_path}")
    
    # Count files for progress bar
    file_count = sum(1 for _ in data_dir.rglob('*') if _.is_file())
    print(f"Found {file_count} files to archive")
    
    with tarfile.open(output_path, 'w:gz') as tar:
        # Add all files with progress bar
        files = list(data_dir.rglob('*'))
        for file_path in tqdm(files, desc="Archiving files"):
            if file_path.is_file():
                # Use relative path from data_dir parent
                arcname = file_path.relative_to(data_dir.parent)
                tar.add(file_path, arcname=arcname, recursive=False)
    
    archive_size = output_path.stat().st_size / (1024 * 1024)  # MB
    print(f"\nArchive created: {output_path}")
    print(f"Archive size: {archive_size:.2f} MB")
    print(f"Compression ratio: {archive_size / (data_dir.stat().st_size / (1024 * 1024)):.2%}")


def upload_to_hf(repo_id: str, archive_path: Path, token: str = None) -> None:
    """
    Upload the archive to Hugging Face Hub.
    
    Args:
        repo_id: Hugging Face repository ID (e.g., "username/dataset-name")
        archive_path: Path to the archive file
        token: Hugging Face token (if None, will use login or environment variable)
    """
    print(f"\nUploading to Hugging Face: {repo_id}")
    
    # Login if token provided
    if token:
        login(token=token)
    else:
        # Try to login (will use cached token or prompt)
        try:
            login()
        except Exception as e:
            print(f"Warning: Could not login automatically: {e}")
            print("Please run: huggingface-cli login")
            print("Or set HF_TOKEN environment variable")
    
    api = HfApi()
    
    # Create repository if it doesn't exist
    try:
        api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True)
        print(f"Repository {repo_id} is ready")
    except Exception as e:
        print(f"Warning: Could not create repository: {e}")
        print("Continuing with upload...")
    
    # Upload file with progress
    print(f"Uploading {archive_path.name}...")
    try:
        api.upload_file(
            path_or_fileobj=str(archive_path),
            path_in_repo=archive_path.name,
            repo_id=repo_id,
            repo_type="dataset",
        )
        print(f"\n✓ Successfully uploaded to {repo_id}")
        print(f"  File: {archive_path.name}")
        print(f"\nTo download, use:")
        print(f"  python download_data_from_hf.py --repo-id {repo_id}")
    except Exception as e:
        print(f"\n✗ Upload failed: {e}")
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        description="Upload data folder to Hugging Face Hub"
    )
    parser.add_argument(
        "--repo-id",
        type=str,
        required=True,
        help="Hugging Face repository ID (e.g., 'username/dataset-name')"
    )
    parser.add_argument(
        "--data-dir",
        type=str,
        default="data",
        help="Data directory to upload (default: 'data')"
    )
    parser.add_argument(
        "--output-archive",
        type=str,
        default=None,
        help="Output archive path (default: 'data.tar.gz' in git root)"
    )
    parser.add_argument(
        "--skip-archive",
        action="store_true",
        help="Skip creating archive (use existing archive)"
    )
    parser.add_argument(
        "--token",
        type=str,
        default=None,
        help="Hugging Face token (or set HF_TOKEN env var)"
    )
    parser.add_argument(
        "--keep-archive",
        action="store_true",
        help="Keep the archive file after upload"
    )
    
    args = parser.parse_args()
    
    # Get git root
    git_root = get_git_root()
    print(f"Git root: {git_root}")
    
    # Resolve data directory
    data_dir = git_root / args.data_dir
    if not data_dir.exists():
        print(f"Error: Data directory not found: {data_dir}")
        sys.exit(1)
    
    # Determine archive path
    if args.output_archive:
        archive_path = Path(args.output_archive).resolve()
    else:
        archive_path = git_root / "data.tar.gz"
    
    # Create archive if needed
    if not args.skip_archive:
        if archive_path.exists():
            response = input(f"Archive {archive_path} already exists. Overwrite? [y/N]: ")
            if response.lower() != 'y':
                print("Using existing archive...")
            else:
                create_data_archive(data_dir, archive_path)
        else:
            create_data_archive(data_dir, archive_path)
    else:
        if not archive_path.exists():
            print(f"Error: Archive not found: {archive_path}")
            sys.exit(1)
        print(f"Using existing archive: {archive_path}")
    
    # Upload to Hugging Face
    token = args.token or os.getenv("HF_TOKEN")
    upload_to_hf(args.repo_id, archive_path, token)
    
    # Clean up archive if requested
    if not args.keep_archive and archive_path.exists():
        response = input(f"\nDelete archive file {archive_path}? [y/N]: ")
        if response.lower() == 'y':
            archive_path.unlink()
            print(f"Deleted {archive_path}")


if __name__ == "__main__":
    main()