Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| #!/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() | |