Add scripts/push_to_hub.py

scripts/push_to_hub.py

@@ -0,0 +1,587 @@
+#!/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()