""" 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()