Upload upload_to_hf.py with huggingface_hub

upload_to_hf.py

@@ -0,0 +1,462 @@
+#!/usr/bin/env python3
+"""
+Upload three datasets to Hugging Face Hub.
+
+Datasets:
+  DATA1: Domain-Specific Code Dataset (115GB, 178 CSV files)
+  DATA2: Code-Documentation Alignment Dataset (2.9GB, 1 JSONL file)
+  DATA3: Programming Problems Generation Dataset (496MB, 1 JSONL file)
+
+Usage:
+  # First, login to Hugging Face:
+  huggingface-cli login
+
+  # Upload all three datasets:
+  python upload_to_hf.py --hf_user YOUR_USERNAME
+
+  # Upload a specific dataset:
+  python upload_to_hf.py --hf_user YOUR_USERNAME --dataset data1
+  python upload_to_hf.py --hf_user YOUR_USERNAME --dataset data2
+  python upload_to_hf.py --hf_user YOUR_USERNAME --dataset data3
+
+  # Use a Hugging Face organization instead of user:
+  python upload_to_hf.py --hf_user YOUR_ORG --dataset all
+
+  # Custom repo names:
+  python upload_to_hf.py --hf_user YOUR_USERNAME \\
+      --repo_name_data1 my-code-dataset \\
+      --repo_name_data2 my-alignment-dataset \\
+      --repo_name_data3 my-problems-dataset
+
+  # Dry run (only create repos and README, no file upload):
+  python upload_to_hf.py --hf_user YOUR_USERNAME --dry_run
+"""
+
+import os
+import argparse
+import logging
+import time
+from pathlib import Path
+from huggingface_hub import HfApi, create_repo
+from huggingface_hub.utils import HfHubHTTPError
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger = logging.getLogger(__name__)
+
+# ============================================================
+# Dataset paths
+# ============================================================
+DATA1_DIR = Path("/home/weifengsun/tangou1/domain_code/src/workdir/dataset_csv")
+DATA2_FILE = Path("/home/weifengsun/tangou1/step2/step22/output/alignment.jsonl")
+DATA3_FILE = Path("/home/weifengsun/tangou1/domain_code/src/datasets/instruct_data/programming_problems.jsonl")
+
+# README files
+DATA1_README = Path("/home/weifengsun/tangou1/DATA1_README.md")
+DATA2_README = Path("/home/weifengsun/tangou1/DATA2_README.md")
+DATA3_README = Path("/home/weifengsun/tangou1/DATA3_README.md")
+
+
+# ============================================================
+# Dataset Card templates (prepended to README content)
+# ============================================================
+
+DATA1_CARD = """---
+license: apache-2.0
+task_categories:
+  - text-generation
+language:
+  - code
+tags:
+  - code
+  - scientific-computing
+  - domain-specific
+  - chemistry
+  - biology
+  - physics
+size_categories:
+  - 1M<n<10M
+---
+
+"""
+
+DATA2_CARD = """---
+license: apache-2.0
+task_categories:
+  - text-generation
+  - text2text-generation
+language:
+  - code
+tags:
+  - code
+  - documentation
+  - docstring-generation
+  - code-documentation-alignment
+  - scientific-computing
+size_categories:
+  - 100K<n<1M
+---
+
+"""
+
+DATA3_CARD = """---
+license: apache-2.0
+task_categories:
+  - text-generation
+  - question-answering
+language:
+  - code
+  - en
+tags:
+  - code
+  - programming-problems
+  - scientific-computing
+  - problem-generation
+size_categories:
+  - 10K<n<100K
+---
+
+"""
+
+
+def build_readme(card_header: str, readme_path: Path) -> str:
+    """Combine YAML front-matter with existing README content."""
+    readme_content = ""
+    if readme_path.exists():
+        readme_content = readme_path.read_text(encoding="utf-8")
+    return card_header + readme_content
+
+
+# ============================================================
+# Upload functions
+# ============================================================
+
+def upload_file_with_retry(
+    api: HfApi,
+    file_path: Path,
+    path_in_repo: str,
+    repo_id: str,
+    max_retries: int = 5,
+    base_delay: float = 2.0,
+    max_delay: float = 300.0,
+    check_existing: bool = True,
+):
+    """
+    Upload a file with retry logic and rate limiting.
+    
+    Args:
+        api: HfApi instance
+        file_path: Local file path
+        path_in_repo: Path in repository
+        repo_id: Repository ID
+        max_retries: Maximum number of retries
+        base_delay: Base delay in seconds for exponential backoff
+        max_delay: Maximum delay in seconds
+        check_existing: Check if file already exists before uploading
+    """
+    # Check if file already exists
+    if check_existing:
+        try:
+            repo_info = api.repo_info(repo_id, repo_type="dataset", files_metadata=True)
+            existing_files = {f.rfilename for f in repo_info.siblings if hasattr(f, 'rfilename')}
+            if path_in_repo in existing_files:
+                logger.info(f"  File {path_in_repo} already exists, skipping.")
+                return True
+        except Exception as e:
+            logger.debug(f"Could not check existing files: {e}")
+    
+    for attempt in range(1, max_retries + 1):
+        try:
+            api.upload_file(
+                path_or_fileobj=str(file_path),
+                path_in_repo=path_in_repo,
+                repo_id=repo_id,
+                repo_type="dataset",
+            )
+            return True
+        except HfHubHTTPError as e:
+            status_code = getattr(e, 'status_code', None) or (e.response.status_code if hasattr(e, 'response') and e.response else None)
+            if status_code == 429:  # Too Many Requests
+                # Extract retry-after header if available
+                retry_after = None
+                if hasattr(e, 'response') and e.response:
+                    retry_after = e.response.headers.get("Retry-After")
+                if retry_after:
+                    wait_time = min(float(retry_after), max_delay)
+                    logger.warning(
+                        f"  Rate limited (429). Waiting {wait_time:.1f}s (Retry-After header)..."
+                    )
+                else:
+                    # Exponential backoff: 2^attempt seconds, capped at max_delay
+                    wait_time = min(base_delay * (2 ** (attempt - 1)), max_delay)
+                    logger.warning(
+                        f"  Rate limited (429). Waiting {wait_time:.1f}s (attempt {attempt}/{max_retries})..."
+                    )
+                time.sleep(wait_time)
+                continue
+            elif status_code and status_code >= 500:  # Server errors
+                wait_time = min(base_delay * (2 ** (attempt - 1)), max_delay)
+                logger.warning(
+                    f"  Server error ({status_code}). Waiting {wait_time:.1f}s (attempt {attempt}/{max_retries})..."
+                )
+                time.sleep(wait_time)
+                continue
+            else:
+                # Other HTTP errors (4xx except 429) - don't retry
+                logger.error(f"  HTTP error {status_code}: {e}")
+                raise
+        except Exception as e:
+            if attempt == max_retries:
+                logger.error(f"  Failed after {max_retries} attempts: {e}")
+                raise
+            wait_time = min(base_delay * (2 ** (attempt - 1)), max_delay)
+            logger.warning(
+                f"  Error: {e}. Waiting {wait_time:.1f}s (attempt {attempt}/{max_retries})..."
+            )
+            time.sleep(wait_time)
+    
+    return False
+
+
+def upload_data1(api: HfApi, repo_id: str, dry_run: bool = False, delay_between_files: float = 3.0):
+    """Upload DATA1: Domain-Specific Code Dataset (178 CSV files, ~115GB)."""
+    logger.info(f"{'[DRY RUN] ' if dry_run else ''}Uploading DATA1 to {repo_id}")
+
+    # Create repo
+    create_repo(repo_id, repo_type="dataset", exist_ok=True, private=False)
+    logger.info(f"Repository {repo_id} created/verified.")
+
+    # Upload README
+    readme = build_readme(DATA1_CARD, DATA1_README)
+    try:
+        api.upload_file(
+            path_or_fileobj=readme.encode("utf-8"),
+            path_in_repo="README.md",
+            repo_id=repo_id,
+            repo_type="dataset",
+        )
+        logger.info("README.md uploaded.")
+    except Exception as e:
+        logger.warning(f"README upload failed (may already exist): {e}")
+
+    if dry_run:
+        logger.info("[DRY RUN] Skipping file uploads.")
+        return
+
+    # Upload CSV files one by one (some files are very large)
+    csv_files = sorted(DATA1_DIR.glob("*.csv"))
+    total = len(csv_files)
+    logger.info(f"Found {total} CSV files to upload.")
+    logger.info(f"Using {delay_between_files}s delay between files to avoid rate limiting.")
+
+    successful = 0
+    failed = 0
+
+    for idx, csv_file in enumerate(csv_files, 1):
+        size_mb = csv_file.stat().st_size / (1024 * 1024)
+        logger.info(f"[{idx}/{total}] Uploading {csv_file.name} ({size_mb:.1f} MB)...")
+        
+        try:
+            success = upload_file_with_retry(
+                api=api,
+                file_path=csv_file,
+                path_in_repo=f"data/{csv_file.name}",
+                repo_id=repo_id,
+                max_retries=5,
+                base_delay=5.0,  # Start with 5s delay for 429 errors
+                max_delay=300.0,  # Max 5 minutes wait
+                check_existing=True,
+            )
+            
+            if success:
+                successful += 1
+                logger.info(f"[{idx}/{total}] ✓ {csv_file.name} uploaded. ({successful} successful, {failed} failed)")
+            else:
+                failed += 1
+                logger.error(f"[{idx}/{total}] ✗ {csv_file.name} failed after retries.")
+        except Exception as e:
+            failed += 1
+            logger.error(f"[{idx}/{total}] ✗ Failed to upload {csv_file.name}: {e}")
+        
+        # Add delay between files to avoid rate limiting (except for last file)
+        if idx < total:
+            logger.debug(f"Waiting {delay_between_files}s before next file...")
+            time.sleep(delay_between_files)
+    
+    logger.info(f"Upload complete: {successful} successful, {failed} failed out of {total} files.")
+
+
+def upload_data2(api: HfApi, repo_id: str, dry_run: bool = False):
+    """Upload DATA2: Code-Documentation Alignment Dataset (~2.9GB)."""
+    logger.info(f"{'[DRY RUN] ' if dry_run else ''}Uploading DATA2 to {repo_id}")
+
+    # Create repo
+    create_repo(repo_id, repo_type="dataset", exist_ok=True, private=False)
+    logger.info(f"Repository {repo_id} created/verified.")
+
+    # Upload README
+    readme = build_readme(DATA2_CARD, DATA2_README)
+    api.upload_file(
+        path_or_fileobj=readme.encode("utf-8"),
+        path_in_repo="README.md",
+        repo_id=repo_id,
+        repo_type="dataset",
+    )
+    logger.info("README.md uploaded.")
+
+    if dry_run:
+        logger.info("[DRY RUN] Skipping file uploads.")
+        return
+
+    size_mb = DATA2_FILE.stat().st_size / (1024 * 1024)
+    logger.info(f"Uploading {DATA2_FILE.name} ({size_mb:.1f} MB)...")
+    api.upload_file(
+        path_or_fileobj=str(DATA2_FILE),
+        path_in_repo=f"data/{DATA2_FILE.name}",
+        repo_id=repo_id,
+        repo_type="dataset",
+    )
+    logger.info(f"✓ {DATA2_FILE.name} uploaded.")
+
+
+def upload_data3(api: HfApi, repo_id: str, dry_run: bool = False):
+    """Upload DATA3: Programming Problems Generation Dataset (~496MB)."""
+    logger.info(f"{'[DRY RUN] ' if dry_run else ''}Uploading DATA3 to {repo_id}")
+
+    # Create repo
+    create_repo(repo_id, repo_type="dataset", exist_ok=True, private=False)
+    logger.info(f"Repository {repo_id} created/verified.")
+
+    # Upload README
+    readme = build_readme(DATA3_CARD, DATA3_README)
+    api.upload_file(
+        path_or_fileobj=readme.encode("utf-8"),
+        path_in_repo="README.md",
+        repo_id=repo_id,
+        repo_type="dataset",
+    )
+    logger.info("README.md uploaded.")
+
+    if dry_run:
+        logger.info("[DRY RUN] Skipping file uploads.")
+        return
+
+    size_mb = DATA3_FILE.stat().st_size / (1024 * 1024)
+    logger.info(f"Uploading {DATA3_FILE.name} ({size_mb:.1f} MB)...")
+    api.upload_file(
+        path_or_fileobj=str(DATA3_FILE),
+        path_in_repo=f"data/{DATA3_FILE.name}",
+        repo_id=repo_id,
+        repo_type="dataset",
+    )
+    logger.info(f"✓ {DATA3_FILE.name} uploaded.")
+
+
+# ============================================================
+# Main
+# ============================================================
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Upload datasets to Hugging Face Hub",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+    parser.add_argument(
+        "--hf_user", type=str, required=True,
+        help="Hugging Face username or organization name",
+    )
+    parser.add_argument(
+        "--dataset", type=str, default="all", choices=["all", "data1", "data2", "data3"],
+        help="Which dataset to upload (default: all)",
+    )
+    parser.add_argument(
+        "--repo_name_data1", type=str, default="SciCode-Domain-Code",
+        help="Repository name for DATA1 (default: SciCode-Domain-Code)",
+    )
+    parser.add_argument(
+        "--repo_name_data2", type=str, default="SciCode-Doc-Alignment",
+        help="Repository name for DATA2 (default: SciCode-Doc-Alignment)",
+    )
+    parser.add_argument(
+        "--repo_name_data3", type=str, default="SciCode-Programming-Problems",
+        help="Repository name for DATA3 (default: SciCode-Programming-Problems)",
+    )
+    parser.add_argument(
+        "--dry_run", action="store_true",
+        help="Only create repos and upload READMEs, skip data files",
+    )
+    parser.add_argument(
+        "--private", action="store_true",
+        help="Create private repositories (default: public)",
+    )
+    parser.add_argument(
+        "--delay", type=float, default=3.0,
+        help="Delay in seconds between file uploads (default: 3.0, increase if getting 429 errors)",
+    )
+    args = parser.parse_args()
+
+    # Verify data paths exist
+    checks = {
+        "data1": DATA1_DIR,
+        "data2": DATA2_FILE,
+        "data3": DATA3_FILE,
+    }
+    for name, path in checks.items():
+        if args.dataset in ("all", name) and not path.exists():
+            logger.error(f"Data path not found: {path}")
+            return
+
+    api = HfApi()
+
+    # Check authentication
+    try:
+        user_info = api.whoami()
+        logger.info(f"Logged in as: {user_info.get('name', user_info.get('fullname', 'unknown'))}")
+    except Exception:
+        logger.error(
+            "Not logged in to Hugging Face. Please run:\n"
+            "  huggingface-cli login\n"
+            "or set the HF_TOKEN environment variable."
+        )
+        return
+
+    repo_ids = {
+        "data1": f"{args.hf_user}/{args.repo_name_data1}",
+        "data2": f"{args.hf_user}/{args.repo_name_data2}",
+        "data3": f"{args.hf_user}/{args.repo_name_data3}",
+    }
+
+    upload_fns = {
+        "data1": upload_data1,
+        "data2": upload_data2,
+        "data3": upload_data3,
+    }
+
+    targets = ["data1", "data2", "data3"] if args.dataset == "all" else [args.dataset]
+
+    logger.info("=" * 60)
+    logger.info("Upload Plan:")
+    for t in targets:
+        logger.info(f"  {t.upper()} -> {repo_ids[t]}")
+    logger.info("=" * 60)
+
+    for t in targets:
+        try:
+            if t == "data1":
+                upload_fns[t](api, repo_ids[t], dry_run=args.dry_run, delay_between_files=args.delay)
+            else:
+                upload_fns[t](api, repo_ids[t], dry_run=args.dry_run)
+            logger.info(f"✓ {t.upper()} upload completed.\n")
+        except Exception as e:
+            logger.error(f"✗ {t.upper()} upload failed: {e}\n")
+
+    logger.info("All done!")
+
+
+if __name__ == "__main__":
+    main()
+