upload_classifiers.py

"""
Upload trained expert classifier checkpoints to the Hugging Face Hub.

Scans for local checkpoint folders produced by train_expert_classifier.py
and uploads ALL checkpoints (checkpoint-100, checkpoint-200, etc.) for each
expert to a single Hub repo. Each checkpoint is uploaded as a tagged revision
so you can load any one by name.

Expected local structure (from HF Trainer with save_strategy="epoch"):
    ./cls-expert-reentrancy/
        checkpoint-50/
            adapter_config.json
            adapter_model.safetensors
            tokenizer.json
            ...
        checkpoint-100/
            ...
        checkpoint-150/
            ...
    ./cls-expert-access-control/
        checkpoint-50/
            ...
        ...

Each expert folder gets a single Hub repo:
    jhsu12/solidity-vuln-cls-reentrancy-v1
    jhsu12/solidity-vuln-cls-access-control-v1
    ...

With tagged revisions for each checkpoint:
    checkpoint-50, checkpoint-100, checkpoint-150, ...

The LAST checkpoint (highest step number) is also uploaded to `main`.

Usage:
    # Auto-detect and upload all expert folders + all checkpoints:
    python upload_classifiers.py

    # Upload from a custom base directory:
    python upload_classifiers.py --base_dir /path/to/training/output

    # Upload only specific experts:
    python upload_classifiers.py --experts reentrancy access-control

    # Upload only the latest checkpoint per expert (skip older ones):
    python upload_classifiers.py --latest_only

    # Dry run — show what would be uploaded without actually doing it:
    python upload_classifiers.py --dry_run

    # Upload to a different Hub namespace:
    python upload_classifiers.py --hub_namespace myorg

    # Upload a single arbitrary folder:
    python upload_classifiers.py --folder ./my-checkpoint --hub_repo jhsu12/my-model

    # If your checkpoints don't follow the cls-expert-* naming,
    # upload a whole expert directory with all its checkpoint-* subfolders:
    python upload_classifiers.py --folder ./my-expert-dir --hub_repo jhsu12/my-model --all_checkpoints
"""

import argparse
import os
import re
import sys
import json
from huggingface_hub import HfApi, create_repo

# ── Expert registry (must match train_expert_classifier.py) ──────────────────
EXPERTS = {
    "reentrancy": "Reentrancy",
    "access-control": "Access Control",
    "integer-overflow-underflow": "Integer Overflow/Underflow",
    "timestamp-dependence": "Timestamp Dependence",
    "unchecked-low-level-calls": "Unchecked Low-Level Calls",
}

# Files we expect in a valid checkpoint
EXPECTED_FILES = ["adapter_config.json", "adapter_model.safetensors"]

# Files to skip during upload (training artifacts, not needed for inference)
IGNORE_PATTERNS = [
    "optimizer*",
    "scheduler*",
    "training_args*",
    "trainer_state*",
    "rng_state*",
    "*.pt",
    "global_step*",
    "runs/",
]


def parse_args():
    parser = argparse.ArgumentParser(
        description="Upload expert classifier checkpoints to Hugging Face Hub."
    )
    parser.add_argument(
        "--base_dir", type=str, default=".",
        help="Base directory containing cls-expert-* folders (default: current dir)"
    )
    parser.add_argument(
        "--experts", type=str, nargs="*", default=None,
        help="Only upload these experts (by slug, e.g. 'reentrancy access-control'). "
             "Default: upload all found."
    )
    parser.add_argument(
        "--hub_namespace", type=str, default="jhsu12",
        help="Hub namespace/username (default: jhsu12)"
    )
    parser.add_argument(
        "--version", type=str, default="v1",
        help="Version suffix for Hub repo name (default: v1)"
    )
    parser.add_argument(
        "--latest_only", action="store_true", default=False,
        help="Only upload the latest checkpoint (highest step) per expert"
    )
    parser.add_argument(
        "--folder", type=str, default=None,
        help="Upload a single specific folder (overrides auto-detection). "
             "Can be a checkpoint dir or a parent dir containing checkpoint-* subdirs."
    )
    parser.add_argument(
        "--hub_repo", type=str, default=None,
        help="Target Hub repo ID for --folder mode (e.g. jhsu12/my-model)"
    )
    parser.add_argument(
        "--all_checkpoints", action="store_true", default=False,
        help="When using --folder, scan it for checkpoint-* subdirs and upload all"
    )
    parser.add_argument(
        "--private", action="store_true", default=False,
        help="Create Hub repos as private (default: public)"
    )
    parser.add_argument(
        "--dry_run", action="store_true", default=False,
        help="Show what would be uploaded without actually uploading"
    )
    parser.add_argument(
        "--commit_message", type=str, default=None,
        help="Custom commit message (default: auto-generated)"
    )
    return parser.parse_args()


def find_checkpoints_in_dir(expert_dir):
    """
    Find all checkpoint-* subdirectories inside an expert folder.
    Returns list of (step_number, folder_name, folder_path) sorted by step.
    """
    checkpoints = []
    if not os.path.isdir(expert_dir):
        return checkpoints

    for name in os.listdir(expert_dir):
        path = os.path.join(expert_dir, name)
        if not os.path.isdir(path):
            continue

        # Match checkpoint-NNN pattern
        match = re.match(r"checkpoint-(\d+)$", name)
        if match:
            step = int(match.group(1))
            checkpoints.append((step, name, path))

    # Sort by step number (ascending)
    checkpoints.sort(key=lambda x: x[0])
    return checkpoints


def find_all_experts(base_dir, filter_experts=None):
    """
    Scan base_dir for cls-expert-* folders and find all checkpoints in each.
    Returns list of (slug, expert_name, expert_dir, checkpoints) tuples.
    """
    found = []
    base_dir = os.path.abspath(base_dir)

    for slug, expert_name in EXPERTS.items():
        if filter_experts and slug not in filter_experts:
            continue

        dir_name = f"cls-expert-{slug}"
        expert_dir = os.path.join(base_dir, dir_name)

        if not os.path.isdir(expert_dir):
            continue

        checkpoints = find_checkpoints_in_dir(expert_dir)

        # Also check if the expert_dir itself is a valid checkpoint (no subdirs)
        if not checkpoints:
            missing, _ = validate_checkpoint(expert_dir)
            if not missing:
                checkpoints = [(0, ".", expert_dir)]

        if checkpoints:
            found.append((slug, expert_name, expert_dir, checkpoints))

    return found


def validate_checkpoint(folder_path):
    """Check if a folder contains the required checkpoint files."""
    missing = []
    for f in EXPECTED_FILES:
        if not os.path.isfile(os.path.join(folder_path, f)):
            missing.append(f)

    files = os.listdir(folder_path) if os.path.isdir(folder_path) else []
    return missing, files


def create_model_card(expert_name, slug, hub_repo_id, checkpoint_tags,
                      base_model="Qwen/Qwen2.5-Coder-3B-Instruct"):
    """Generate a README.md model card for the Hub repo."""

    tags_table = ""
    if len(checkpoint_tags) > 1:
        tags_table = "\n## Available Checkpoints\n\n"
        tags_table += "Load a specific checkpoint with `revision=`:\n"
        tags_table += "```python\n"
        tags_table += f'model = PeftModel.from_pretrained(base, "{hub_repo_id}", revision="checkpoint-200")\n'
        tags_table += "```\n\n"
        tags_table += "| Tag | Step |\n|-----|------|\n"
        for tag_name, step in checkpoint_tags:
            is_main = " ← `main`" if step == checkpoint_tags[-1][1] else ""
            tags_table += f"| `{tag_name}` | {step}{is_main} |\n"

    return f"""---
library_name: peft
base_model: {base_model}
tags:
  - peft
  - lora
  - sequence-classification
  - solidity
  - smart-contract
  - vulnerability-detection
  - {slug}
pipeline_tag: text-classification
license: apache-2.0
---

# Solidity Vulnerability Classifier — {expert_name}

Binary classifier that detects **{expert_name}** vulnerabilities in Solidity smart contracts.

## Model Details

- **Base model**: [{base_model}](https://huggingface.co/{base_model})
- **Method**: QLoRA (4-bit NF4) + classification head
- **Task**: Sequence Classification (2 labels: safe / vulnerable)
- **LoRA rank**: 16, targeting q_proj, k_proj, v_proj, o_proj
- **Classification head**: `modules_to_save=["score"]`
{tags_table}
## Usage

```python
from transformers import AutoTokenizer, AutoModelForSequenceClassification, BitsAndBytesConfig
from peft import PeftModel
import torch

base_model = "{base_model}"
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)

model = AutoModelForSequenceClassification.from_pretrained(
    base_model, num_labels=2, quantization_config=bnb_config,
    device_map="auto", trust_remote_code=True, ignore_mismatched_sizes=True)
model = PeftModel.from_pretrained(model, "{hub_repo_id}")
model.eval()

tokenizer = AutoTokenizer.from_pretrained("{hub_repo_id}", trust_remote_code=True)

code = "pragma solidity ^0.8.0; contract Example {{ ... }}"
inputs = tokenizer(code, return_tensors="pt", truncation=True, max_length=1536).to(model.device)
with torch.no_grad():
    logits = model(**inputs).logits
probs = torch.softmax(logits, dim=-1)
print(f"Safe: {{probs[0][0]:.2%}}, Vulnerable: {{probs[0][1]:.2%}}")
```

Or use the inference script:
```bash
python inference_classifier.py --checkpoint {hub_repo_id} --file contract.sol
```

## Part of

This is one of 5 expert classifiers in the
[Solidity Vulnerability Detector](https://huggingface.co/jhsu12/solidity-vulnerability-detector) system.

| Expert | Hub Repo |
|--------|----------|
| Reentrancy | `jhsu12/solidity-vuln-cls-reentrancy-v1` |
| Access Control | `jhsu12/solidity-vuln-cls-access-control-v1` |
| Integer Overflow/Underflow | `jhsu12/solidity-vuln-cls-integer-overflow-underflow-v1` |
| Timestamp Dependence | `jhsu12/solidity-vuln-cls-timestamp-dependence-v1` |
| Unchecked Low-Level Calls | `jhsu12/solidity-vuln-cls-unchecked-low-level-calls-v1` |
"""


def upload_single_checkpoint(api, folder_path, hub_repo_id, tag_name,
                              commit_message, revision="main", dry_run=False):
    """Upload one checkpoint folder to a specific branch/revision of a Hub repo."""
    missing, files = validate_checkpoint(folder_path)

    if missing:
        print(f"       ⚠️  Missing files: {missing} — skipping")
        return False

    if dry_run:
        print(f"       🔍 DRY RUN — {len(files)} files → {revision}")
        return True

    try:
        api.upload_folder(
            folder_path=folder_path,
            repo_id=hub_repo_id,
            revision=revision,
            ignore_patterns=IGNORE_PATTERNS,
            commit_message=commit_message,
        )
        return True
    except Exception as e:
        print(f"       ❌ Upload failed: {e}")
        return False


def upload_expert(api, slug, expert_name, checkpoints, hub_repo_id,
                  private=False, latest_only=False, commit_message=None,
                  dry_run=False):
    """
    Upload all checkpoints for a single expert.
    - Latest checkpoint → main branch
    - Each checkpoint → tagged revision (checkpoint-NNN)
    """
    print(f"\n{'━' * 60}")
    print(f"  📦 {expert_name}")
    print(f"     Hub:         https://huggingface.co/{hub_repo_id}")
    print(f"     Checkpoints: {len(checkpoints)} found")

    for step, name, path in checkpoints:
        _, files = validate_checkpoint(path)
        print(f"       • {name} ({len(files)} files)")

    if latest_only:
        checkpoints = [checkpoints[-1]]  # Keep only the highest step
        print(f"     --latest_only: uploading only {checkpoints[0][1]}")

    # Create repo
    if not dry_run:
        try:
            create_repo(hub_repo_id, private=private, exist_ok=True)
            print(f"     ✅ Repo ready")
        except Exception as e:
            print(f"     ❌ Failed to create repo: {e}")
            return 0, len(checkpoints)

    succeeded = 0
    failed = 0

    # Upload each checkpoint
    for i, (step, name, path) in enumerate(checkpoints):
        is_latest = (i == len(checkpoints) - 1)
        tag_name = name  # e.g. "checkpoint-100"

        # Latest goes to main; all get a tag
        if is_latest:
            print(f"\n     🚀 {name} (step {step}) → main + tag '{tag_name}'")

            # Upload to main
            msg = commit_message or f"Upload {expert_name} classifier — {name} (latest)"
            ok = upload_single_checkpoint(
                api, path, hub_repo_id, tag_name, msg,
                revision="main", dry_run=dry_run,
            )

            if ok:
                # Also create model card on main
                if not dry_run:
                    checkpoint_tags = [(ckpt_name, s) for s, ckpt_name, _ in checkpoints]
                    readme = create_model_card(expert_name, slug, hub_repo_id, checkpoint_tags)
                    try:
                        api.upload_file(
                            path_or_fileobj=readme.encode("utf-8"),
                            path_in_repo="README.md",
                            repo_id=hub_repo_id,
                            commit_message=f"Add model card for {expert_name}",
                        )
                    except Exception:
                        pass  # Non-fatal

                # Create tag for this checkpoint
                if not dry_run and name != ".":
                    try:
                        api.create_tag(
                            hub_repo_id, tag=tag_name,
                            tag_message=f"{expert_name} — {name} (step {step})",
                        )
                        print(f"       🏷️  Tagged as '{tag_name}'")
                    except Exception:
                        pass  # Tag may already exist

                succeeded += 1
            else:
                failed += 1
        else:
            print(f"\n     📤 {name} (step {step}) → branch '{tag_name}'")

            # Create a branch for this checkpoint
            if not dry_run:
                try:
                    api.create_branch(hub_repo_id, branch=tag_name)
                except Exception:
                    pass  # Branch may already exist

            msg = commit_message or f"Upload {expert_name} classifier — {name}"
            ok = upload_single_checkpoint(
                api, path, hub_repo_id, tag_name, msg,
                revision=tag_name, dry_run=dry_run,
            )

            if ok:
                succeeded += 1
                if not dry_run:
                    print(f"       ✅ Uploaded to branch '{tag_name}'")
            else:
                failed += 1

    return succeeded, failed


def main():
    args = parse_args()
    api = HfApi()

    print("=" * 60)
    print("  Upload Expert Classifier Checkpoints to Hugging Face Hub")
    print("=" * 60)

    # ── Mode 1: Upload a single folder / directory ────────────────────────────
    if args.folder:
        if not args.hub_repo:
            print("❌ --hub_repo is required when using --folder")
            sys.exit(1)

        folder = os.path.abspath(args.folder)
        if not os.path.isdir(folder):
            print(f"❌ Folder not found: {folder}")
            sys.exit(1)

        slug = os.path.basename(folder.rstrip("/"))
        expert_name = slug.replace("-", " ").replace("cls expert ", "").title()

        if args.all_checkpoints:
            # Treat folder as parent dir containing checkpoint-* subdirs
            checkpoints = find_checkpoints_in_dir(folder)
            if not checkpoints:
                print(f"❌ No checkpoint-* subdirectories found in {folder}")
                sys.exit(1)
        else:
            # Treat folder as a single checkpoint
            checkpoints = [(0, os.path.basename(folder), folder)]

        ok, fail = upload_expert(
            api, slug, expert_name, checkpoints, args.hub_repo,
            private=args.private, latest_only=args.latest_only,
            commit_message=args.commit_message, dry_run=args.dry_run,
        )
        print(f"\n  Uploaded: {ok}, Failed: {fail}")
        sys.exit(0 if fail == 0 else 1)

    # ── Mode 2: Auto-detect expert folders ────────────────────────────────────
    base_dir = os.path.abspath(args.base_dir)
    print(f"\n🔍 Scanning: {base_dir}")

    experts = find_all_experts(base_dir, filter_experts=args.experts)

    if not experts:
        print(f"\n❌ No expert checkpoint folders found!")
        print(f"\n   Expected structure:")
        print(f"   {base_dir}/")
        for slug in EXPERTS:
            print(f"     cls-expert-{slug}/")
            print(f"       checkpoint-50/")
            print(f"         adapter_config.json")
            print(f"         adapter_model.safetensors")
            print(f"         ...")
            print(f"       checkpoint-100/")
            print(f"         ...")
        print(f"\n   Tips:")
        print(f"   • Run train_expert_classifier.py first to generate checkpoints")
        print(f"   • Use --base_dir to point to the parent directory")
        print(f"   • Use --folder for a single non-standard directory")
        sys.exit(1)

    # Summary of what was found
    total_ckpts = sum(len(ckpts) for _, _, _, ckpts in experts)
    print(f"\n   Found {len(experts)} expert(s) with {total_ckpts} total checkpoint(s):")
    for slug, name, expert_dir, ckpts in experts:
        steps = [str(s) for s, _, _ in ckpts]
        print(f"     • {name}: {len(ckpts)} checkpoints (steps: {', '.join(steps)})")

    if args.latest_only:
        print(f"\n   --latest_only: will upload only the latest checkpoint per expert")

    # Upload each expert
    total_ok = 0
    total_fail = 0

    for slug, expert_name, expert_dir, checkpoints in experts:
        hub_repo_id = f"{args.hub_namespace}/solidity-vuln-cls-{slug}-{args.version}"

        ok, fail = upload_expert(
            api, slug, expert_name, checkpoints, hub_repo_id,
            private=args.private, latest_only=args.latest_only,
            commit_message=args.commit_message, dry_run=args.dry_run,
        )
        total_ok += ok
        total_fail += fail

    # Final summary
    print(f"\n{'=' * 60}")
    action = "Would upload" if args.dry_run else "Uploaded"
    print(f"  {action}: {total_ok} checkpoint(s) across {len(experts)} expert(s)")
    if total_fail:
        print(f"  Failed: {total_fail}")

    if not args.dry_run and total_ok > 0:
        print(f"\n  Hub repos:")
        for slug, name, _, ckpts in experts:
            repo_id = f"{args.hub_namespace}/solidity-vuln-cls-{slug}-{args.version}"
            ckpt_names = [cn for _, cn, _ in ckpts]
            print(f"    https://huggingface.co/{repo_id}")
            if len(ckpt_names) > 1 and not args.latest_only:
                print(f"      Branches: {', '.join(ckpt_names)}")
                print(f"      main = {ckpt_names[-1]} (latest)")

    print(f"\n  To load a specific checkpoint:")
    print(f'    PeftModel.from_pretrained(base, "jhsu12/solidity-vuln-cls-reentrancy-v1", revision="checkpoint-100")')
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()