scripts/push_all_to_hub.py

#!/usr/bin/env python3
"""Push all BlitzKode artifacts to HuggingFace Hub in one command.

Uploads
-------
1. LoRA adapter (1.5B) -> neuralbroker/blitzkode-1.5b-lora
2. LoRA adapter (0.5B) -> neuralbroker/blitzkode-lora-0.5b
3. GGUF model file    -> neuralbroker/blitzkode  (into a GGUF-specific branch/folder)
4. Project docs       -> neuralbroker/blitzkode  (README, model card, runbook, eval results)

Usage
-----
    # Export HF_TOKEN first, then run:
    python scripts/push_all_to_hub.py

    # Or pass token directly:
    python scripts/push_all_to_hub.py --token hf_XXXX

    # Dry-run to validate without pushing:
    python scripts/push_all_to_hub.py --dry-run
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]

PUSH_SCRIPT = REPO_ROOT / "scripts" / "push_to_hub.py"

ARTIFACTS = [
    {
        "label": "1.5B LoRA adapter (primary)",
        "checkpoint": REPO_ROOT / "checkpoints" / "blitzkode-1.5b-lora" / "final",
        "repo_id": "neuralbroker/blitzkode-1.5b-lora",
        "commit_message": "Upload BlitzKode 1.5B LoRA adapter v2.1 (100-step SFT)",
    },
    {
        "label": "0.5B LoRA adapter (lightweight)",
        "checkpoint": REPO_ROOT / "checkpoints" / "available-lora-0.5b-full" / "final",
        "repo_id": "neuralbroker/blitzkode-lora-0.5b",
        "commit_message": "Upload BlitzKode 0.5B LoRA adapter v2.1 (50-step SFT)",
    },
]

PROJECT_DOC_FILES = [
    (REPO_ROOT / "README.md", "README.md"),
    (REPO_ROOT / "MODEL_CARD.md", "MODEL_CARD.md"),
    (REPO_ROOT / "docs" / "PROJECT_OVERVIEW.md", "docs/PROJECT_OVERVIEW.md"),
    (REPO_ROOT / "docs" / "PRODUCTION_RUNBOOK.md", "docs/PRODUCTION_RUNBOOK.md"),
    (REPO_ROOT / "docs" / "evaluation_results.json", "docs/evaluation_results.json"),
]


def resolve_token(args_token: str | None, dry_run: bool) -> str:
    token = (args_token or os.environ.get("HF_TOKEN", "")).strip()
    if token or dry_run:
        return token

    try:
        from huggingface_hub import get_token  # noqa: PLC0415

        token = (get_token() or "").strip()
    except ImportError:
        token = ""

    if token:
        return token

    print(
        "\n[ERROR] HuggingFace token required.\n"
        "  Option 1: set HF_TOKEN=hf_XXXX\n"
        "  Option 2: python scripts/push_all_to_hub.py --token hf_XXXX\n"
        "  Option 3: run huggingface-cli login\n"
        "  Option 4: run --dry-run to validate without pushing\n"
        "\nGet a write token at: https://huggingface.co/settings/tokens",
        file=sys.stderr,
    )
    sys.exit(1)


def push_gguf(token: str, gguf_path: Path, dry_run: bool) -> None:
    if not gguf_path.exists():
        print(f"  [SKIP] GGUF not found: {gguf_path}")
        return

    size_gb = gguf_path.stat().st_size / 1024**3
    print(f"\n  Uploading GGUF ({size_gb:.2f} GB) -> neuralbroker/blitzkode ...")
    if dry_run:
        print("  [DRY RUN] skipped.")
        return

    from huggingface_hub import HfApi  # noqa: PLC0415
    from huggingface_hub.errors import HfHubHTTPError  # noqa: PLC0415

    api = HfApi(token=token)
    try:
        api.create_repo("neuralbroker/blitzkode", repo_type="model", exist_ok=True)
        api.upload_file(
            path_or_fileobj=str(gguf_path),
            path_in_repo="blitzkode.gguf",
            repo_id="neuralbroker/blitzkode",
            repo_type="model",
            commit_message="Update GGUF model Q8_0 (1.5B merged + quantised)",
        )
        print("  [OK] GGUF uploaded -> https://huggingface.co/neuralbroker/blitzkode")
    except HfHubHTTPError as exc:
        print(f"  [ERROR] GGUF upload failed: {exc}", file=sys.stderr)


def push_project_docs(token: str, dry_run: bool) -> bool:
    print("\n  Uploading project docs -> neuralbroker/blitzkode ...")
    missing = [path for path, _ in PROJECT_DOC_FILES if not path.exists()]
    if missing:
        for path in missing:
            print(f"  [ERROR] Missing doc file: {path}", file=sys.stderr)
        return False

    for path, repo_path in PROJECT_DOC_FILES:
        size_kb = path.stat().st_size / 1024
        print(f"  - {repo_path} ({size_kb:.1f} KB)")

    if dry_run:
        print("  [DRY RUN] skipped.")
        return True

    from huggingface_hub import HfApi  # noqa: PLC0415
    from huggingface_hub.errors import HfHubHTTPError  # noqa: PLC0415

    api = HfApi(token=token)
    try:
        api.create_repo("neuralbroker/blitzkode", repo_type="model", exist_ok=True)
        for path, repo_path in PROJECT_DOC_FILES:
            api.upload_file(
                path_or_fileobj=str(path),
                path_in_repo=repo_path,
                repo_id="neuralbroker/blitzkode",
                repo_type="model",
                commit_message=f"Update project documentation: {repo_path}",
            )
        print("  [OK] Project docs uploaded -> https://huggingface.co/neuralbroker/blitzkode")
        return True
    except HfHubHTTPError as exc:
        print(f"  [ERROR] Project docs upload failed: {exc}", file=sys.stderr)
        return False


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--token", default=None, help="HuggingFace write token (or set HF_TOKEN env var / huggingface-cli login).")
    parser.add_argument("--dry-run", action="store_true", help="Validate only, do not push.")
    args = parser.parse_args()

    token = resolve_token(args.token, args.dry_run)

    print("=" * 72)
    print("BLITZKODE - PUSH ALL ARTIFACTS TO HUGGING FACE HUB")
    if args.dry_run:
        print("(DRY RUN - nothing will be pushed)")
    print("=" * 72)

    failures: list[str] = []

    for art in ARTIFACTS:
        print(f"\n{'-' * 60}")
        print(f"  Artifact : {art['label']}")
        print(f"  Repo     : {art['repo_id']}")

        checkpoint: Path = art["checkpoint"]
        if not checkpoint.exists():
            print(f"  [SKIP] Checkpoint not found: {checkpoint}")
            continue

        cmd = [
            sys.executable,
            str(PUSH_SCRIPT),
            "--checkpoint",
            str(checkpoint),
            "--repo-id",
            art["repo_id"],
            "--commit-message",
            art["commit_message"],
        ]
        if args.dry_run:
            cmd.append("--dry-run")
        if token:
            cmd += ["--token", token]

        result = subprocess.run(cmd)
        if result.returncode != 0:
            failures.append(art["repo_id"])
            print(f"  [FAIL] Push exited with code {result.returncode}", file=sys.stderr)

    # GGUF upload (direct via huggingface_hub)
    print(f"\n{'-' * 60}")
    print("  Artifact : GGUF model (neuralbroker/blitzkode)")
    if not args.dry_run and token:
        push_gguf(token, REPO_ROOT / "blitzkode.gguf", dry_run=False)
    else:
        push_gguf(token, REPO_ROOT / "blitzkode.gguf", dry_run=True)

    print(f"\n{'-' * 60}")
    print("  Artifact : Project documentation (neuralbroker/blitzkode)")
    if not push_project_docs(token, dry_run=args.dry_run):
        failures.append("neuralbroker/blitzkode docs")

    # Summary
    print(f"\n{'=' * 72}")
    if failures:
        print(f"PUSH FINISHED WITH FAILURES: {failures}")
        sys.exit(1)
    else:
        print("ALL PUSHES COMPLETE")
        print("\nHuggingFace repos:")
        for art in ARTIFACTS:
            print(f"  https://huggingface.co/{art['repo_id']}")
        print("  https://huggingface.co/neuralbroker/blitzkode  (GGUF)")
        print("  https://huggingface.co/neuralbroker/blitzkode/tree/main/docs  (docs)")


if __name__ == "__main__":
    main()