Update upload script: support all checkpoint-* folders per expert, upload as branches + tags

upload_classifiers.py

@@ -2,25 +2,38 @@
 Upload trained expert classifier checkpoints to the Hugging Face Hub.
 
 Scans for local checkpoint folders produced by train_expert_classifier.py
-and uploads each one to its corresponding Hub repo.
-
-Expected local structure (from training):
-    ./cls-expert-reentrancy/best_model/
-        adapter_config.json
-        adapter_model.safetensors
-        tokenizer.json
-        tokenizer_config.json
-        ...
-    ./cls-expert-access-control/best_model/
+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 folder gets uploaded to:
+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 found in current directory:
+    # Auto-detect and upload all expert folders + all checkpoints:
     python upload_classifiers.py
 
     # Upload from a custom base directory:
@@ -29,18 +42,26 @@ Usage:
     # 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 specific folders by path (overrides auto-detection):
-    python upload_classifiers.py --folder ./my-custom-checkpoint --hub_repo jhsu12/my-custom-model
+    # 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
@@ -64,9 +85,9 @@ IGNORE_PATTERNS = [
     "training_args*",
     "trainer_state*",
     "rng_state*",
-    "*.pt",           # Full optimizer states
-    "global_step*",   # Checkpoint step markers
-    "runs/",          # Tensorboard logs
+    "*.pt",
+    "global_step*",
+    "runs/",
 ]
 
 
@@ -92,18 +113,22 @@ def parse_args():
         help="Version suffix for Hub repo name (default: v1)"
     )
     parser.add_argument(
-        "--subfolder", type=str, default="best_model",
-        help="Subfolder inside each expert dir containing the checkpoint (default: best_model). "
-             "Set to empty string if checkpoint is directly in the expert folder."
+        "--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)"
+        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)"
@@ -119,30 +144,59 @@ def parse_args():
     return parser.parse_args()
 
 
-def find_checkpoint_folders(base_dir, subfolder="best_model", filter_experts=None):
+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 return list of
-    (slug, expert_name, folder_path) tuples.
+    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 subfolder:
-            checkpoint_dir = os.path.join(expert_dir, subfolder)
-        else:
-            checkpoint_dir = expert_dir
-
-        if filter_experts and slug not in filter_experts:
+        if not os.path.isdir(expert_dir):
             continue
 
-        if not os.path.isdir(checkpoint_dir):
-            continue
+        checkpoints = find_checkpoints_in_dir(expert_dir)
 
-        found.append((slug, expert_name, checkpoint_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
 
@@ -154,14 +208,26 @@ def validate_checkpoint(folder_path):
         if not os.path.isfile(os.path.join(folder_path, f)):
             missing.append(f)
 
-    # List what's actually there
     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, base_model="Qwen/Qwen2.5-Coder-3B-Instruct"):
+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}
@@ -188,7 +254,7 @@ Binary classifier that detects **{expert_name}** vulnerabilities in Solidity sma
 - **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
@@ -196,7 +262,6 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification, Bits
 from peft import PeftModel
 import torch
 
-# Load
 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)
@@ -209,7 +274,6 @@ model.eval()
 
 tokenizer = AutoTokenizer.from_pretrained("{hub_repo_id}", trust_remote_code=True)
 
-# Classify
 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():
@@ -238,62 +302,135 @@ This is one of 5 expert classifiers in the
 """
 
 
-def upload_checkpoint(api, folder_path, hub_repo_id, expert_name, slug,
-                      private=False, commit_message=None, dry_run=False):
-    """Upload a single checkpoint folder to the Hub."""
-
-    print(f"\n{'─' * 60}")
-    print(f"  📦 {expert_name}")
-    print(f"     Local:  {folder_path}")
-    print(f"     Hub:    https://huggingface.co/{hub_repo_id}")
-
-    # Validate
+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)
-    print(f"     Files:  {len(files)} ({', '.join(files[:5])}{'...' if len(files) > 5 else ''})")
 
     if missing:
-        print(f"     ⚠️  Missing expected files: {missing}")
-        print(f"     Skipping — this doesn't look like a valid classifier checkpoint.")
+        print(f"       ⚠️  Missing files: {missing} — skipping")
         return False
 
     if dry_run:
-        print(f"     🔍 DRY RUN — would upload {len(files)} files to {hub_repo_id}")
+        print(f"       🔍 DRY RUN — {len(files)} files → {revision}")
         return True
 
-    # Create repo if needed
-    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 False
-
-    # Generate and save model card
-    readme_content = create_model_card(expert_name, slug, hub_repo_id)
-    readme_path = os.path.join(folder_path, "README.md")
-    readme_existed = os.path.exists(readme_path)
-    with open(readme_path, "w") as f:
-        f.write(readme_content)
-
-    # Upload
-    msg = commit_message or f"Upload {expert_name} classifier (QLoRA, Qwen2.5-Coder-3B)"
     try:
         api.upload_folder(
             folder_path=folder_path,
             repo_id=hub_repo_id,
+            revision=revision,
             ignore_patterns=IGNORE_PATTERNS,
-            commit_message=msg,
+            commit_message=commit_message,
         )
-        print(f"     🚀 Uploaded to https://huggingface.co/{hub_repo_id}")
+        return True
     except Exception as e:
-        print(f"     ❌ Upload failed: {e}")
+        print(f"       ❌ Upload failed: {e}")
         return False
-    finally:
-        # Clean up generated README if it didn't exist before
-        if not readme_existed and os.path.exists(readme_path):
-            os.remove(readme_path)
 
-    return True
+
+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():
@@ -304,7 +441,7 @@ def main():
     print("  Upload Expert Classifier Checkpoints to Hugging Face Hub")
     print("=" * 60)
 
-    # ── Mode 1: Upload a single specific folder ──────────────────────────────
+    # ── Mode 1: Upload a single folder / directory ────────────────────────────
     if args.folder:
         if not args.hub_repo:
             print("❌ --hub_repo is required when using --folder")
@@ -315,77 +452,95 @@ def main():
             print(f"❌ Folder not found: {folder}")
             sys.exit(1)
 
-        # Try to guess expert name from folder
         slug = os.path.basename(folder.rstrip("/"))
         expert_name = slug.replace("-", " ").replace("cls expert ", "").title()
 
-        success = upload_checkpoint(
-            api, folder, args.hub_repo, expert_name, slug,
-            private=args.private, commit_message=args.commit_message,
-            dry_run=args.dry_run,
+        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,
         )
-        sys.exit(0 if success else 1)
+        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}")
-    print(f"   Subfolder: {args.subfolder or '(root)'}")
 
-    checkpoints = find_checkpoint_folders(
-        base_dir,
-        subfolder=args.subfolder,
-        filter_experts=args.experts,
-    )
+    experts = find_all_experts(base_dir, filter_experts=args.experts)
 
-    if not checkpoints:
-        print(f"\n❌ No checkpoint folders found!")
+    if not experts:
+        print(f"\n❌ No expert checkpoint folders found!")
         print(f"\n   Expected structure:")
         print(f"   {base_dir}/")
         for slug in EXPERTS:
-            sub = f"/{args.subfolder}" if args.subfolder else ""
-            print(f"     cls-expert-{slug}{sub}/")
-            print(f"       adapter_config.json")
-            print(f"       adapter_model.safetensors")
-            print(f"       tokenizer.json")
-            print(f"       ...")
+            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 directory containing the folders")
-        print(f"   • Use --subfolder '' if checkpoints are in the expert folder root")
+        print(f"   • Use --base_dir to point to the parent directory")
+        print(f"   • Use --folder for a single non-standard directory")
         sys.exit(1)
 
-    print(f"\n   Found {len(checkpoints)} expert checkpoint(s):")
-    for slug, name, path in checkpoints:
-        print(f"     • {name} → {path}")
+    # 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)})")
 
-    # Upload each
-    succeeded = 0
-    failed = 0
+    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, folder_path in checkpoints:
+    for slug, expert_name, expert_dir, checkpoints in experts:
         hub_repo_id = f"{args.hub_namespace}/solidity-vuln-cls-{slug}-{args.version}"
 
-        success = upload_checkpoint(
-            api, folder_path, hub_repo_id, expert_name, slug,
-            private=args.private, commit_message=args.commit_message,
-            dry_run=args.dry_run,
+        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,
         )
-        if success:
-            succeeded += 1
-        else:
-            failed += 1
+        total_ok += ok
+        total_fail += fail
 
-    # Summary
+    # Final summary
     print(f"\n{'=' * 60}")
     action = "Would upload" if args.dry_run else "Uploaded"
-    print(f"  {action}: {succeeded}/{succeeded + failed}")
-    if failed:
-        print(f"  Failed: {failed}")
-    if not args.dry_run and succeeded > 0:
+    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, _ in checkpoints:
+        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}")