初始化项目，由ModelHub XC社区提供模型
Model: oxdev/security-auditor-grpo Source: Original Platform
This commit is contained in:
ModelHub XC
252
train_grpo_job.py
Normal file
252
train_grpo_job.py
Normal file
@@ -0,0 +1,252 @@
|
||||
#!/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 via HfApi.upload_folder().
|
||||
"""
|
||||
|
||||
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-0.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, **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)
|
||||
|
||||
n = len(rewards) if rewards else 1
|
||||
logger.info(f"[reward] finding_rate={finding_count/n:.2f} compile_rate={compile_count/n:.2f} exploit_rate={pass_count/n:.2f}")
|
||||
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(f"CUDA available: {torch.cuda.is_available()}")
|
||||
if torch.cuda.is_available():
|
||||
logger.info(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
|
||||
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 — NO hub_model_id, NO log_completions, NO push_to_hub
|
||||
# This prevents ANY Hub calls during __init__ or training
|
||||
config = GRPOConfig(
|
||||
output_dir=OUTPUT_DIR,
|
||||
num_train_epochs=2,
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=2,
|
||||
num_generations=2,
|
||||
max_completion_length=512,
|
||||
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,
|
||||
# CRITICAL: all Hub-related settings OFF to prevent 401 at init
|
||||
push_to_hub=False,
|
||||
log_completions=False,
|
||||
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("GRPOTrainer initialized successfully!")
|
||||
|
||||
logger.info("Starting training...")
|
||||
trainer.train()
|
||||
logger.info("Training complete!")
|
||||
|
||||
# Save locally
|
||||
logger.info(f"Saving model to {OUTPUT_DIR}...")
|
||||
trainer.save_model(OUTPUT_DIR)
|
||||
|
||||
# Manual push to hub using HfApi — safer and more explicit
|
||||
hf_token = os.environ.get("HF_TOKEN")
|
||||
if hf_token:
|
||||
logger.info(f"Pushing to hub: {HUB_MODEL_ID}")
|
||||
try:
|
||||
from huggingface_hub import HfApi
|
||||
api = HfApi(token=hf_token)
|
||||
# Create repo if needed (ignore error if exists)
|
||||
try:
|
||||
api.create_repo(repo_id=HUB_MODEL_ID, exist_ok=True)
|
||||
except Exception as e:
|
||||
logger.warning(f"create_repo warning (may already exist): {e}")
|
||||
# Upload entire output folder
|
||||
api.upload_folder(
|
||||
folder_path=OUTPUT_DIR,
|
||||
repo_id=HUB_MODEL_ID,
|
||||
commit_message="GRPO training complete — smart contract security auditor",
|
||||
)
|
||||
logger.info(f"✅ Model pushed to https://huggingface.co/{HUB_MODEL_ID}")
|
||||
except Exception as e:
|
||||
logger.error(f"Push failed: {e}")
|
||||
logger.info(f"Model saved locally at {OUTPUT_DIR}")
|
||||
else:
|
||||
logger.warning("No HF_TOKEN found — model saved locally only")
|
||||
logger.info(f"Model at: {OUTPUT_DIR}")
|
||||
|
||||
logger.info("=" * 60)
|
||||
logger.info("DONE")
|
||||
logger.info("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user