scripts/huggingface/upload_state_splits_to_hf.py

#!/usr/bin/env python3
"""
Upload state-split parquet files to Hugging Face Datasets.

This script uploads the state-split files created by split_gold_by_state.py
to HuggingFace, organized by state for easier access and distribution.

Each state gets its own dataset with all relevant data:
- CommunityOne/one-data-AL (Alabama data)
- CommunityOne/one-data-AK (Alaska data)
- etc.

Benefits:
- Users can download only the state(s) they need
- Smaller file sizes (faster downloads)
- Better organization for state-specific analysis

Usage:
    # Upload all states
    python scripts/upload_state_splits_to_hf.py --all
    
    # Upload specific state
    python scripts/upload_state_splits_to_hf.py --state AL
    
    # Upload multiple states
    python scripts/upload_state_splits_to_hf.py --states AL AK AZ
    
    # Dry run (see what would be uploaded)
    python scripts/upload_state_splits_to_hf.py --all --dry-run
"""

import argparse
import os
from pathlib import Path
from typing import List, Dict
import pandas as pd
from datasets import Dataset
from huggingface_hub import login, create_repo, HfApi
from loguru import logger

# Load environment variables
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    logger.warning("python-dotenv not installed")


class StateDataUploader:
    """Upload state-split data files to HuggingFace."""
    
    # US States and territories
    ALL_STATES = [
        'AA', 'AE', 'AK', 'AL', 'AP', 'AR', 'AS', 'AZ', 'CA', 'CO', 'CT', 
        'DC', 'DE', 'FL', 'FM', 'GA', 'GU', 'HI', 'IA', 'ID', 'IL', 'IN', 
        'KS', 'KY', 'LA', 'MA', 'MD', 'ME', 'MH', 'MI', 'MN', 'MO', 'MP', 
        'MS', 'MT', 'NC', 'ND', 'NE', 'NH', 'NJ', 'NM', 'NV', 'NY', 'OH', 
        'OK', 'OR', 'PA', 'PR', 'PW', 'RI', 'SC', 'SD', 'TN', 'TX', 'UT', 
        'VA', 'VI', 'VT', 'WA', 'WI', 'WV', 'WY'
    ]
    
    # State names for better descriptions
    STATE_NAMES = {
        'AL': 'Alabama', 'AK': 'Alaska', 'AZ': 'Arizona', 'AR': 'Arkansas',
        'CA': 'California', 'CO': 'Colorado', 'CT': 'Connecticut', 'DE': 'Delaware',
        'FL': 'Florida', 'GA': 'Georgia', 'HI': 'Hawaii', 'ID': 'Idaho',
        'IL': 'Illinois', 'IN': 'Indiana', 'IA': 'Iowa', 'KS': 'Kansas',
        'KY': 'Kentucky', 'LA': 'Louisiana', 'ME': 'Maine', 'MD': 'Maryland',
        'MA': 'Massachusetts', 'MI': 'Michigan', 'MN': 'Minnesota', 'MS': 'Mississippi',
        'MO': 'Missouri', 'MT': 'Montana', 'NE': 'Nebraska', 'NV': 'Nevada',
        'NH': 'New Hampshire', 'NJ': 'New Jersey', 'NM': 'New Mexico', 'NY': 'New York',
        'NC': 'North Carolina', 'ND': 'North Dakota', 'OH': 'Ohio', 'OK': 'Oklahoma',
        'OR': 'Oregon', 'PA': 'Pennsylvania', 'RI': 'Rhode Island', 'SC': 'South Carolina',
        'SD': 'South Dakota', 'TN': 'Tennessee', 'TX': 'Texas', 'UT': 'Utah',
        'VT': 'Vermont', 'VA': 'Virginia', 'WA': 'Washington', 'WV': 'West Virginia',
        'WI': 'Wisconsin', 'WY': 'Wyoming', 'DC': 'District of Columbia',
        'PR': 'Puerto Rico', 'VI': 'Virgin Islands', 'GU': 'Guam',
        'AS': 'American Samoa', 'MP': 'Northern Mariana Islands',
        'AA': 'Armed Forces Americas', 'AE': 'Armed Forces Europe', 'AP': 'Armed Forces Pacific',
        'FM': 'Federated States of Micronesia', 'MH': 'Marshall Islands', 'PW': 'Palau'
    }
    
    def __init__(self, repo_prefix: str = None, token: str = None, splits_dir: str = "data/gold/by_state"):
        """
        Initialize uploader.
        
        Args:
            repo_prefix: HF repo prefix (e.g., "CommunityOne/one-data")
            token: HF token (or set HUGGINGFACE_TOKEN environment variable)
            splits_dir: Directory containing state-split parquet files
        """
        self.repo_prefix = repo_prefix or "CommunityOne/one-data"
        self.token = token or os.getenv("HUGGINGFACE_TOKEN")
        self.splits_dir = Path(splits_dir)
        
        if not self.token:
            raise ValueError(
                "Hugging Face token required! "
                "Get it from https://huggingface.co/settings/tokens "
                "and set HUGGINGFACE_TOKEN environment variable"
            )
        
        # Login
        login(token=self.token)
        logger.info(f"✅ Logged in to Hugging Face")
        logger.info(f"📦 Dataset prefix: {self.repo_prefix}")
    
    def get_state_files(self, state: str) -> List[Path]:
        """Get all parquet files for a specific state."""
        pattern = f"*_{state}.parquet"
        files = list(self.splits_dir.glob(pattern))
        return sorted(files)
    
    def upload_state(self, state: str, dry_run: bool = False) -> bool:
        """
        Upload all files for a specific state to HuggingFace.
        
        Args:
            state: State abbreviation (e.g., 'AL', 'AK')
            dry_run: If True, only report what would be done
            
        Returns:
            True if successful, False otherwise
        """
        files = self.get_state_files(state)
        
        if not files:
            logger.warning(f"No files found for state: {state}")
            return False
        
        repo_name = f"{self.repo_prefix}-{state}"
        state_name = self.STATE_NAMES.get(state, state)
        
        logger.info(f"📂 Processing {state_name} ({state})")
        logger.info(f"  Files: {len(files)}")
        
        if dry_run:
            logger.info(f"  [DRY RUN] Would create repo: {repo_name}")
            for f in files:
                logger.info(f"    [DRY RUN] Would upload: {f.name}")
            return True
        
        # Create repo
        try:
            create_repo(
                repo_id=repo_name,
                repo_type="dataset",
                private=False,
                exist_ok=True
            )
            logger.info(f"✅ Repository ready: https://huggingface.co/datasets/{repo_name}")
        except Exception as e:
            logger.error(f"❌ Failed to create repo: {e}")
            return False
        
        # Upload each file as a separate split
        api = HfApi()
        total_records = 0
        total_size = 0
        
        for file_path in files:
            try:
                # Determine split name from filename
                # e.g., nonprofits_organizations_AL.parquet -> organizations
                base_name = file_path.stem.replace(f"_{state}", "")
                split_name = base_name.replace("nonprofits_", "").replace("jurisdictions_", "").replace("domains_", "")
                
                # Read and upload
                df = pd.read_parquet(file_path)
                records = len(df)
                size_mb = file_path.stat().st_size / 1024 / 1024
                
                total_records += records
                total_size += size_mb
                
                logger.info(f"  📤 Uploading {split_name}: {records:,} records ({size_mb:.2f} MB)")
                
                # Convert to HF Dataset and push
                dataset = Dataset.from_pandas(df)
                dataset.push_to_hub(
                    repo_id=repo_name,
                    split=split_name,
                    token=self.token
                )
                
                logger.success(f"    ✅ Uploaded {split_name}")
                
            except Exception as e:
                logger.error(f"    ❌ Failed to upload {file_path.name}: {e}")
        
        logger.success(f"✅ {state_name} complete: {total_records:,} records, {total_size:.2f} MB")
        logger.success(f"   View at: https://huggingface.co/datasets/{repo_name}")
        
        return True
    
    def upload_all_states(self, dry_run: bool = False) -> Dict[str, bool]:
        """
        Upload all available states.
        
        Args:
            dry_run: If True, only report what would be done
            
        Returns:
            Dict mapping state to success status
        """
        logger.info("🚀 Uploading all states to HuggingFace...")
        logger.info(f"  Input directory: {self.splits_dir}")
        logger.info("")
        
        # Find which states have data
        available_states = set()
        for f in self.splits_dir.glob("*_??.parquet"):
            state = f.stem[-2:]
            if state in self.ALL_STATES:
                available_states.add(state)
        
        available_states = sorted(available_states)
        logger.info(f"Found data for {len(available_states)} states/territories")
        logger.info("")
        
        results = {}
        for state in available_states:
            try:
                success = self.upload_state(state, dry_run=dry_run)
                results[state] = success
                logger.info("")
            except Exception as e:
                logger.error(f"❌ Error processing {state}: {e}")
                results[state] = False
                logger.info("")
        
        # Summary
        successful = sum(1 for v in results.values() if v)
        logger.success("=" * 60)
        logger.success(f"✅ Uploaded {successful}/{len(results)} states")
        logger.success("=" * 60)
        
        return results


def main():
    parser = argparse.ArgumentParser(
        description="Upload state-split parquet files to HuggingFace Datasets",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Upload all states
  python scripts/upload_state_splits_to_hf.py --all
  
  # Upload Alabama data
  python scripts/upload_state_splits_to_hf.py --state AL
  
  # Upload multiple states
  python scripts/upload_state_splits_to_hf.py --states AL AK AZ CA
  
  # Dry run (see what would be uploaded)
  python scripts/upload_state_splits_to_hf.py --all --dry-run
        """
    )
    
    parser.add_argument('--all', action='store_true',
                       help='Upload all available states')
    parser.add_argument('--state', type=str,
                       help='Upload a specific state (e.g., AL)')
    parser.add_argument('--states', nargs='+',
                       help='Upload multiple states (e.g., AL AK AZ)')
    parser.add_argument('--dry-run', action='store_true',
                       help='Show what would be done without actually uploading')
    parser.add_argument('--repo-prefix', type=str, default='CommunityOne/one-data',
                       help='HuggingFace repo prefix (default: CommunityOne/one-data)')
    parser.add_argument('--splits-dir', type=str, default='data/gold/by_state',
                       help='Directory containing state-split files (default: data/gold/by_state)')
    
    args = parser.parse_args()
    
    # Initialize uploader
    uploader = StateDataUploader(
        repo_prefix=args.repo_prefix,
        splits_dir=args.splits_dir
    )
    
    # Handle commands
    if args.all:
        uploader.upload_all_states(dry_run=args.dry_run)
    
    elif args.state:
        uploader.upload_state(args.state.upper(), dry_run=args.dry_run)
    
    elif args.states:
        for state in args.states:
            uploader.upload_state(state.upper(), dry_run=args.dry_run)
            logger.info("")
    
    else:
        parser.print_help()


if __name__ == "__main__":
    main()