Upload train_grpo_job.py with huggingface_hub

train_grpo_job.py

@@ -0,0 +1,220 @@
+#!/usr/bin/env python3
+"""
+train_grpo_job.py — Self-contained GRPO training job for HF Jobs.
+
+Loads dataset from HF Hub, runs GRPO training with custom reward functions,
+pushes model to Hub on completion.
+"""
+
+import logging
+import os
+import re
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+
+import torch
+from datasets import load_dataset
+from trl import GRPOTrainer, GRPOConfig
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger(__name__)
+
+# ─── Config ───────────────────────────────────────────────────────────────────
+MODEL_NAME = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
+DATASET_ID = "oxdev/smart-contract-security-sft"
+OUTPUT_DIR = "/tmp/grpo_output"
+HUB_MODEL_ID = "oxdev/security-auditor-grpo"
+
+FORGE_AVAILABLE = shutil.which("forge") is not None
+
+# ─── Reward Functions ─────────────────────────────────────────────────────────
+
+def extract_finding_block(text: str) -> dict | None:
+    pattern = re.compile(
+        r'FINDING\s*\|\s*contract:\s*(\S+)\s*\|\s*function:\s*(\S+)\s*\|'
+        r'\s*bug_class:\s*(\S+)\s*\|\s*confidence:\s*(\d+)',
+        re.IGNORECASE
+    )
+    match = pattern.search(text)
+    if not match:
+        return None
+    return {
+        "contract": match.group(1),
+        "function": match.group(2),
+        "bug_class": match.group(3),
+        "confidence": int(match.group(4)),
+    }
+
+
+def extract_solidity_poc(text: str) -> str | None:
+    pattern = re.compile(r'```solidity\s*\n(.*?)```', re.DOTALL)
+    matches = pattern.findall(text)
+    if not matches:
+        return None
+    for code in matches:
+        if "is Test" in code or "function test_" in code:
+            return code.strip()
+    return max(matches, key=len).strip() if matches else None
+
+
+def _check_solidity_syntax(code: str) -> bool:
+    required = [r'pragma\s+solidity', r'contract\s+\w+', r'function\s+\w+']
+    return all(re.search(p, code) for p in required)
+
+
+def run_forge_test(poc_code: str, timeout: int = 30) -> dict:
+    if not FORGE_AVAILABLE:
+        return {
+            "compiled": False,
+            "test_passed": False,
+            "syntax_valid": _check_solidity_syntax(poc_code),
+        }
+
+    tmpdir = tempfile.mkdtemp(prefix="forge_poc_")
+    try:
+        test_dir = Path(tmpdir) / "test"
+        test_dir.mkdir()
+        (Path(tmpdir) / "foundry.toml").write_text('[profile.default]\nsrc = "src"\nout = "out"\nlibs = ["lib"]\nsolc_version = "0.8.24"\n')
+        (Path(tmpdir) / "src").mkdir()
+
+        try:
+            subprocess.run(
+                ["forge", "install", "foundry-rs/forge-std", "--no-git", "--no-commit"],
+                cwd=tmpdir, capture_output=True, timeout=60,
+            )
+        except Exception:
+            pass
+
+        (Path(tmpdir) / "remappings.txt").write_text("forge-std/=lib/forge-std/src/\n")
+        (test_dir / "PoC.t.sol").write_text(poc_code)
+
+        build = subprocess.run(["forge", "build"], cwd=tmpdir, capture_output=True, text=True, timeout=timeout)
+        if build.returncode != 0:
+            return {"compiled": False, "test_passed": False}
+
+        test = subprocess.run(["forge", "test", "-vv"], cwd=tmpdir, capture_output=True, text=True, timeout=timeout)
+        return {"compiled": True, "test_passed": test.returncode == 0 and "PASS" in test.stdout}
+
+    except Exception:
+        return {"compiled": False, "test_passed": False}
+    finally:
+        shutil.rmtree(tmpdir, ignore_errors=True)
+
+
+def security_audit_reward(completions, log_extra=None, log_metric=None, **kwargs):
+    """Primary reward: FINDING block + PoC compilation + exploit verification."""
+    rewards = []
+    finding_count = compile_count = pass_count = 0
+
+    for completion in completions:
+        text = completion[0]["content"] if isinstance(completion, list) else str(completion)
+        reward = -1.0
+
+        finding = extract_finding_block(text)
+        if finding:
+            finding_count += 1
+            reward = 0.0
+            poc = extract_solidity_poc(text)
+            if poc:
+                reward = 0.2
+                result = run_forge_test(poc)
+                if result.get("compiled") or result.get("syntax_valid", False):
+                    compile_count += 1
+                    reward = 0.5
+                    if result.get("test_passed"):
+                        pass_count += 1
+                        reward = 1.0
+        elif any(kw in text.lower() for kw in ["vulnerability", "exploit", "bug", "finding"]):
+            reward = -0.5
+
+        rewards.append(reward)
+
+    if log_metric and rewards:
+        log_metric("finding_rate", finding_count / len(rewards))
+        log_metric("compile_rate", compile_count / len(rewards))
+        log_metric("exploit_rate", pass_count / len(rewards))
+
+    return rewards
+
+
+def format_reward(completions, **kwargs):
+    """Secondary reward: structural format compliance."""
+    rewards = []
+    for completion in completions:
+        text = completion[0]["content"] if isinstance(completion, list) else str(completion)
+        reward = 0.0
+        if re.search(r'FINDING\s*\|', text):
+            fields = sum(bool(re.search(p, text)) for p in [r'path:', r'proof:', r'description:', r'fix:'])
+            reward = 0.3 + (0.05 * fields)
+        if re.search(r'```solidity', text):
+            reward += 0.1
+        rewards.append(reward)
+    return rewards
+
+
+# ─── Main ─────────────────────────────────────────────────────────────────────
+
+def main():
+    logger.info("=" * 60)
+    logger.info("GRPO Training — Smart Contract Security Auditor")
+    logger.info(f"Model: {MODEL_NAME}")
+    logger.info(f"Dataset: {DATASET_ID}")
+    logger.info(f"Forge available: {FORGE_AVAILABLE}")
+    logger.info(f"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}")
+    logger.info("=" * 60)
+
+    # Load dataset
+    logger.info("Loading dataset from HF Hub...")
+    dataset = load_dataset(DATASET_ID, split="train")
+    logger.info(f"Dataset: {len(dataset)} samples, columns={dataset.column_names}")
+
+    # Configure GRPO
+    config = GRPOConfig(
+        output_dir=OUTPUT_DIR,
+        num_train_epochs=2,
+        per_device_train_batch_size=2,
+        num_generations=4,
+        max_completion_length=1536,
+        learning_rate=5e-7,
+        beta=0.0,
+        scale_rewards=True,
+        reward_weights=[0.7, 0.3],
+        gradient_checkpointing=True,
+        bf16=True,
+        logging_steps=5,
+        logging_first_step=True,
+        logging_strategy="steps",
+        disable_tqdm=True,
+        save_strategy="steps",
+        save_steps=50,
+        save_total_limit=2,
+        log_completions=True,
+        push_to_hub=True,
+        hub_model_id=HUB_MODEL_ID,
+        report_to="none",
+        seed=42,
+    )
+
+    # Train
+    logger.info("Initializing GRPOTrainer...")
+    trainer = GRPOTrainer(
+        model=MODEL_NAME,
+        args=config,
+        reward_funcs=[security_audit_reward, format_reward],
+        train_dataset=dataset,
+    )
+
+    logger.info("Starting training...")
+    trainer.train()
+
+    logger.info("Saving model...")
+    trainer.save_model(OUTPUT_DIR)
+    trainer.push_to_hub()
+
+    logger.info(f"✅ Done! Model pushed to https://huggingface.co/{HUB_MODEL_ID}")
+
+
+if __name__ == "__main__":
+    main()