scripts/push_to_hub.py

#!/usr/bin/env python3
"""Push a PEFT/LoRA adapter checkpoint to the Hugging Face Hub.



Usage examples

--------------

    # Validate everything without pushing:

    python scripts/push_to_hub.py --dry-run



    # Push with defaults (reads HF_TOKEN from environment):

    python scripts/push_to_hub.py



    # Explicit options:

    python scripts/push_to_hub.py \\

        --checkpoint checkpoints/available-lora-0.5b-full/final \\

        --repo-id neuralbroker/blitzkode-lora-0.5b \\

        --commit-message "Add trained adapter v2.1"



    # Private repo push with explicit token:

    python scripts/push_to_hub.py --private --token hf_...

"""

from __future__ import annotations

import argparse
import json
import os
import sys
import textwrap
from pathlib import Path

# ---------------------------------------------------------------------------
# Constants / defaults
# ---------------------------------------------------------------------------
REPO_ROOT = Path(__file__).resolve().parents[1]

DEFAULT_CHECKPOINT = REPO_ROOT / "checkpoints" / "available-lora-0.5b-full" / "final"
DEFAULT_REPO_ID = "neuralbroker/blitzkode-lora-0.5b"
DEFAULT_REPO_TYPE = "model"
DEFAULT_COMMIT_MSG = "Upload BlitzKode LoRA adapter"

# Files that must be present for a valid PEFT adapter
REQUIRED_FILES = ["adapter_config.json", "adapter_model.safetensors"]


# ---------------------------------------------------------------------------
# CLI argument parsing
# ---------------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=DEFAULT_CHECKPOINT,
        metavar="PATH",
        help=(
            f"Path to the adapter checkpoint directory. "
            f"(default: {DEFAULT_CHECKPOINT})"
        ),
    )
    parser.add_argument(
        "--repo-id",
        default=DEFAULT_REPO_ID,
        metavar="OWNER/REPO",
        help=f"HuggingFace repo to push to. (default: {DEFAULT_REPO_ID})",
    )
    parser.add_argument(
        "--repo-type",
        default=DEFAULT_REPO_TYPE,
        choices=("model", "dataset", "space"),
        help=f"Repository type. (default: {DEFAULT_REPO_TYPE})",
    )
    parser.add_argument(
        "--private",
        action="store_true",
        help="Create the repository as private. (default: public)",
    )
    parser.add_argument(
        "--token",
        default=None,
        metavar="HF_TOKEN",
        help=(
            "HuggingFace API write token. "
            "Falls back to the HF_TOKEN environment variable if not set."
        ),
    )
    parser.add_argument(
        "--create-repo",
        action=argparse.BooleanOptionalAction,
        default=True,
        help=(
            "Create the HuggingFace repo if it does not exist. "
            "Use --no-create-repo to skip. (default: True)"
        ),
    )
    parser.add_argument(
        "--commit-message",
        default=DEFAULT_COMMIT_MSG,
        metavar="MSG",
        help=f"Commit message for the Hub upload. (default: '{DEFAULT_COMMIT_MSG}')",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help=(
            "Validate the checkpoint and configuration but do NOT push anything "
            "to Hugging Face Hub. Useful for CI or pre-flight checks."
        ),
    )
    return parser.parse_args()


# ---------------------------------------------------------------------------
# Dependency check
# ---------------------------------------------------------------------------
def check_huggingface_hub() -> None:
    """Abort with a helpful message if huggingface_hub is not installed."""
    try:
        import huggingface_hub  # noqa: F401  # type: ignore[import]
    except ImportError:
        print(
            "\n[ERROR] The `huggingface_hub` package is not installed.\n"
            "Install it with one of the following commands:\n\n"
            "    pip install huggingface_hub\n"
            "    pip install -r requirements-training.txt\n",
            file=sys.stderr,
        )
        sys.exit(1)


# ---------------------------------------------------------------------------
# Checkpoint validation
# ---------------------------------------------------------------------------
def validate_checkpoint(checkpoint: Path) -> dict:
    """Ensure the checkpoint directory is valid.



    Checks that the directory exists and contains every file listed in

    REQUIRED_FILES. Returns the parsed ``adapter_config.json`` dict.

    """
    if not checkpoint.exists():
        print(
            f"\n[ERROR] Checkpoint directory not found: {checkpoint}\n"
            "Run training first, e.g.:\n"
            "    python scripts/train_available.py\n",
            file=sys.stderr,
        )
        sys.exit(1)

    if not checkpoint.is_dir():
        print(
            f"\n[ERROR] Checkpoint path is not a directory: {checkpoint}\n",
            file=sys.stderr,
        )
        sys.exit(1)

    missing = [f for f in REQUIRED_FILES if not (checkpoint / f).exists()]
    if missing:
        print(
            f"\n[ERROR] Missing required files in {checkpoint}:\n"
            + "\n".join(f"  - {f}" for f in missing)
            + "\n\nIs this a valid PEFT adapter checkpoint?\n",
            file=sys.stderr,
        )
        sys.exit(1)

    config_path = checkpoint / "adapter_config.json"
    try:
        adapter_config: dict = json.loads(config_path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        print(
            f"\n[ERROR] adapter_config.json is not valid JSON: {exc}\n",
            file=sys.stderr,
        )
        sys.exit(1)
    except OSError as exc:
        print(
            f"\n[ERROR] Could not read adapter_config.json: {exc}\n",
            file=sys.stderr,
        )
        sys.exit(1)

    return adapter_config


# ---------------------------------------------------------------------------
# Token resolution
# ---------------------------------------------------------------------------
def resolve_token(args_token: str | None, *, dry_run: bool = False) -> str:
    """Return the HF token, or abort with instructions if none is found."""
    token = args_token or os.environ.get("HF_TOKEN", "")
    if token:
        return token

    if dry_run:
        # A token is not needed for dry runs, return a placeholder.
        return "__dry_run_placeholder__"

    print(
        "\n[ERROR] No HuggingFace API token found.\n"
        "Provide a write token using one of these methods:\n\n"
        "  1. CLI flag:\n"
        "       python scripts/push_to_hub.py --token hf_YOUR_TOKEN\n\n"
        "  2. Environment variable (recommended):\n"
        "       Windows CMD : set HF_TOKEN=hf_YOUR_TOKEN\n"
        "       PowerShell  : $env:HF_TOKEN = 'hf_YOUR_TOKEN'\n"
        "       Linux/macOS : export HF_TOKEN=hf_YOUR_TOKEN\n\n"
        "  3. HuggingFace CLI login (persists across sessions):\n"
        "       pip install huggingface_hub\n"
        "       huggingface-cli login\n\n"
        "Generate a token at: https://huggingface.co/settings/tokens\n"
        "Make sure the token has **write** access to the target repo.\n",
        file=sys.stderr,
    )
    sys.exit(1)


# ---------------------------------------------------------------------------
# Model card / README generation
# ---------------------------------------------------------------------------
def build_model_card(adapter_config: dict, repo_id: str) -> str:
    """Generate the HuggingFace-compatible README.md content for the adapter repo."""
    base_model = adapter_config.get(
        "base_model_name_or_path", "Qwen/Qwen2.5-0.5B-Instruct"
    )
    lora_r = adapter_config.get("r", 16)
    lora_alpha = adapter_config.get("lora_alpha", 32)
    lora_dropout = adapter_config.get("lora_dropout", 0.05)
    target_modules: list = adapter_config.get("target_modules", [])
    modules_str = (
        ", ".join(f"`{m}`" for m in target_modules)
        if target_modules
        else "`q_proj`, `k_proj`, `v_proj`, `o_proj`"
    )

    # YAML frontmatter -------------------------------------------------------
    frontmatter = textwrap.dedent(f"""\

        ---

        language:

          - en

        license: mit

        library_name: peft

        tags:

          - code-generation

          - lora

          - qwen2.5

          - blitzkode

          - coding-assistant

          - fine-tuned

          - peft

        base_model: {base_model}

        pipeline_tag: text-generation

        ---

        """)

    # README body ------------------------------------------------------------
    body = textwrap.dedent(f"""\

        # BlitzKode LoRA Adapter (0.5B)



        **BlitzKode** is a local AI coding assistant fine-tuned from

        **[{base_model}](https://huggingface.co/{base_model})** using LoRA

        (Low-Rank Adaptation). This repository contains the PEFT adapter — the

        research-friendly version that can be hot-loaded on top of the base model.



        > **Creator:** [Sajad (neuralbroker)](https://github.com/neuralbroker)

        > **GitHub:** <https://github.com/neuralbroker/blitzkode>

        > **Production GGUF:** [`neuralbroker/blitzkode`](https://huggingface.co/neuralbroker/blitzkode)



        ---



        ## Model Details



        | Property | Value |

        |---|---|

        | **Adapter version** | 2.1 |

        | **Base model** | `{base_model}` |

        | **LoRA rank (r)** | {lora_r} |

        | **LoRA alpha** | {lora_alpha} |

        | **LoRA dropout** | {lora_dropout} |

        | **Target modules** | {modules_str} |

        | **Training steps** | 50 |

        | **Final loss** | ~0.48 |

        | **Library** | PEFT |

        | **License** | MIT |



        ---



        ## Training Pipeline



        This adapter was produced by a **4-stage fine-tuning pipeline** applied

        to the Qwen2.5 family:



        | Stage | Method | Purpose |

        |---|---|---|

        | 1 | SFT | Supervised fine-tuning on 71 curated algorithmic coding problems |

        | 2 | Reward-SFT | Continued SFT with heuristic reward signals for code correctness and formatting |

        | 3 | DPO | Direct Preference Optimization on handcrafted chosen/rejected pairs |

        | 4 | LoRA SFT (this adapter) | Final LoRA fine-tune (r={lora_r}) on 99 samples; base model Qwen2.5-0.5B |



        ### Training Dataset (199 total samples)



        | Subset | Count | Source | License |

        |---|---|---|---|

        | Curated algorithmic problems | 71 | Custom (local) — arrays, strings, trees, DP, graphs | MIT |

        | MetaMathQA samples | 100 | [`meta-math/MetaMathQA`](https://huggingface.co/datasets/meta-math/MetaMathQA) | CC BY 4.0 |

        | Python/JavaScript patterns | 28 | Custom (local) — decorators, context managers, data classes | MIT |

        | **Total** | **199** | | |



        ---



        ## Usage



        ### Load with PEFT



        ```python

        from peft import PeftModel

        from transformers import AutoModelForCausalLM, AutoTokenizer



        base_model_id = "{base_model}"

        adapter_repo  = "{repo_id}"



        tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)

        model = AutoModelForCausalLM.from_pretrained(

            base_model_id,

            torch_dtype="auto",

            device_map="auto",

            trust_remote_code=True,

        )

        model = PeftModel.from_pretrained(model, adapter_repo)

        model.eval()

        ```



        ### Generate code



        ```python

        prompt = (

            "<|im_start|>system\\n"

            "You are BlitzKode, a precise AI coding assistant created by Sajad.\\n"

            "<|im_end|>\\n"

            "<|im_start|>user\\n"

            "Write a Python function for binary search with full edge-case handling.\\n"

            "<|im_end|>\\n"

            "<|im_start|>assistant\\n"

        )



        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        outputs = model.generate(

            **inputs,

            max_new_tokens=300,

            temperature=0.7,

            do_sample=True,

            repetition_penalty=1.1,

        )

        print(tokenizer.decode(outputs[0], skip_special_tokens=True))

        ```



        ### Merge adapter into base model (for export)



        ```python

        merged = model.merge_and_unload()

        merged.save_pretrained("blitzkode-0.5b-merged")

        tokenizer.save_pretrained("blitzkode-0.5b-merged")

        ```



        ---



        ## Prompt Format



        BlitzKode uses the **ChatML** template standard for Qwen models:



        ```

        <|im_start|>system

        You are BlitzKode, a precise AI coding assistant created by Sajad.<|im_end|>

        <|im_start|>user

        {{your question}}<|im_end|>

        <|im_start|>assistant

        ```



        ---



        ## Limitations



        - **Text-only** — no image/multimodal support.

        - **0.5B parameters** — smaller and faster than the 1.5B GGUF variant; may be

          less accurate on complex algorithmic tasks.

        - **2048-token context** — not suitable for long repository-level analysis.

        - **Review all outputs** — generated code must be tested before use in production.

        - **Not security-audited** — do not use for cryptographic or safety-critical code

          without thorough expert review.

        - **Math reasoning** — MetaMathQA training improves basic reasoning but does not

          substitute a dedicated math model.



        ---



        ## Relation to the Production Model



        | Variant | Repo | Size | Runtime | Use case |

        |---|---|---|---|---|

        | GGUF (1.5B, F16) | [`neuralbroker/blitzkode`](https://huggingface.co/neuralbroker/blitzkode) | ~3 GB | llama.cpp / llama-cpp-python | Production; CPU/GPU, no Python ML stack needed |

        | LoRA adapter (0.5B) | `{repo_id}` (this repo) | ~100 MB | PEFT + Transformers | Research; merging, further fine-tuning, quantization |



        ---



        ## License



        **MIT** — see [LICENSE](https://github.com/neuralbroker/blitzkode/blob/main/LICENSE).



        You must also comply with the upstream

        [{base_model}](https://huggingface.co/{base_model}) license

        when redistributing any derived weights.



        ---



        ## Citation



        ```bibtex

        @software{{blitzkode2025,

          author  = {{Sajad}},

          title   = {{BlitzKode: A Local AI Coding Assistant}},

          year    = {{2025}},

          url     = {{https://github.com/neuralbroker/blitzkode}}

        }}

        ```

        """)

    return frontmatter + "\n" + body


# ---------------------------------------------------------------------------
# Main push routine
# ---------------------------------------------------------------------------
def push(args: argparse.Namespace) -> None:  # noqa: C901
    check_huggingface_hub()

    # Import here so the check above can give a clean error first.
    from huggingface_hub import HfApi  # type: ignore[import]
    from huggingface_hub.utils import HfHubHTTPError  # type: ignore[import]

    sep = "=" * 70
    print(sep)
    print("BlitzKode — Push LoRA Adapter to Hugging Face Hub")
    if args.dry_run:
        print("(DRY RUN — nothing will be pushed)")
    print(sep)

    # ------------------------------------------------------------------
    # Step 1: Validate checkpoint
    # ------------------------------------------------------------------
    print(f"\n[1/5] Validating checkpoint directory …")
    print(f"      Path: {args.checkpoint}")
    adapter_config = validate_checkpoint(args.checkpoint)

    base_model = adapter_config.get("base_model_name_or_path", "unknown")
    lora_r = adapter_config.get("r", "?")
    lora_alpha = adapter_config.get("lora_alpha", "?")
    target_modules = adapter_config.get("target_modules", [])
    files_found = sorted(p.name for p in args.checkpoint.iterdir() if p.is_file())

    print(f"      base_model     : {base_model}")
    print(f"      lora r / alpha : {lora_r} / {lora_alpha}")
    print(f"      target_modules : {target_modules}")
    print(f"      files          : {files_found}")
    print("      [OK] Checkpoint is valid.")

    # ------------------------------------------------------------------
    # Step 2: Resolve token
    # ------------------------------------------------------------------
    print("\n[2/5] Resolving HuggingFace token …")
    token = resolve_token(args.token, dry_run=args.dry_run)
    if args.dry_run:
        print("      [OK] Token check skipped (dry run).")
    else:
        masked = token[:8] + "..." if len(token) > 8 else "***"
        print(f"      [OK] Token resolved (starts with: {masked})")

    # ------------------------------------------------------------------
    # Dry-run exit
    # ------------------------------------------------------------------
    if args.dry_run:
        print()
        print(sep)
        print("DRY RUN COMPLETE — all validations passed, nothing was pushed.")
        print(f"  Checkpoint  : {args.checkpoint}")
        print(f"  Target repo : https://huggingface.co/{args.repo_id}")
        print(f"  Repo type   : {args.repo_type}")
        print(f"  Private     : {args.private}")
        print(f"  Files ready : {files_found}")
        print(sep)
        return

    api = HfApi(token=token)

    # ------------------------------------------------------------------
    # Step 3: Create repo (if requested)
    # ------------------------------------------------------------------
    if args.create_repo:
        print(f"\n[3/5] Creating / verifying repo: {args.repo_id} …")
        try:
            repo_url = api.create_repo(
                repo_id=args.repo_id,
                repo_type=args.repo_type,
                private=args.private,
                exist_ok=True,   # silently succeed if repo already exists
            )
            print(f"      [OK] Repo ready: {repo_url}")
        except HfHubHTTPError as exc:
            print(
                f"\n[ERROR] Failed to create / access repo '{args.repo_id}':\n"
                f"  {exc}\n"
                "Check that your token has write access and the repo name is correct.\n",
                file=sys.stderr,
            )
            sys.exit(1)
    else:
        print("\n[3/5] Skipping repo creation (--no-create-repo).")

    # ------------------------------------------------------------------
    # Step 4: Upload checkpoint folder
    # ------------------------------------------------------------------
    print(f"\n[4/5] Uploading checkpoint folder → {args.repo_id} …")
    print(f"      Commit message: \"{args.commit_message}\"")
    try:
        commit_info = api.upload_folder(
            folder_path=str(args.checkpoint),
            repo_id=args.repo_id,
            repo_type=args.repo_type,
            commit_message=args.commit_message,
        )
        commit_ref = getattr(commit_info, "oid", None) or str(commit_info)
        print(f"      [OK] Folder uploaded. Commit: {commit_ref}")
    except HfHubHTTPError as exc:
        print(
            f"\n[ERROR] Folder upload failed:\n  {exc}\n",
            file=sys.stderr,
        )
        sys.exit(1)

    # ------------------------------------------------------------------
    # Step 5: Upload model card README.md
    # ------------------------------------------------------------------
    print("\n[5/5] Uploading model card (README.md) …")
    readme_content = build_model_card(adapter_config, args.repo_id)
    try:
        api.upload_file(
            path_or_fileobj=readme_content.encode("utf-8"),
            path_in_repo="README.md",
            repo_id=args.repo_id,
            repo_type=args.repo_type,
            commit_message="Update model card README.md",
        )
        print("      [OK] README.md uploaded.")
    except HfHubHTTPError as exc:
        # Non-fatal: the adapter files are already uploaded.
        print(
            f"\n[WARN] Could not upload README.md (adapter files were uploaded OK):\n"
            f"  {exc}\n"
            "You can upload the model card manually from the Hub web interface.\n",
            file=sys.stderr,
        )

    # ------------------------------------------------------------------
    # Summary
    # ------------------------------------------------------------------
    repo_url = f"https://huggingface.co/{args.repo_id}"
    print()
    print(sep)
    print("PUSH COMPLETE")
    print(f"  Repo URL      : {repo_url}")
    print(f"  Checkpoint    : {args.checkpoint}")
    print(f"  Files pushed  : {files_found}")
    print(f"  Base model    : {base_model}")
    print(f"  LoRA r/alpha  : {lora_r}/{lora_alpha}")
    print(f"  Commit msg    : {args.commit_message}")
    print(sep)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
def main() -> None:
    args = parse_args()
    push(args)


if __name__ == "__main__":
    main()