{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "provenance": [], "gpuType": "T4" }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# šŸ” Smart Contract Security Auditor ā€” GRPO V2 Training\n", "\n", "Train a specialized smart contract security auditor using **Group Relative Policy Optimization (GRPO)**\n", "on **50,902 real audit findings** from top security firms.\n", "\n", "**Model:** Qwen2.5-Coder-0.5B-Instruct ā†’ oxdev/security-auditor-grpo\n", "\n", "**Dataset:** [oxdev/smart-contract-security-audit-v2](https://huggingface.co/datasets/oxdev/smart-contract-security-audit-v2)\n", "\n", "**Hardware:** Free Colab T4 (16GB VRAM)\n", "\n", "---\n", "\n", "## Setup\n", "1. Go to **Runtime ā†’ Change runtime type ā†’ T4 GPU**\n", "2. Run all cells in order\n", "3. When prompted, enter your HuggingFace token (needs write access)\n", "4. Training takes ~4-6 hours on a T4 GPU with 2K samples" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Cell 1: Install dependencies\n", "!pip install -q torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121\n", "!pip install -q transformers>=4.51.0 trl>=1.2.0 datasets accelerate huggingface_hub\n", "print('\\nāœ… Dependencies installed!')\n", "\n", "import torch\n", "print(f'PyTorch: {torch.__version__}')\n", "print(f'CUDA available: {torch.cuda.is_available()}')\n", "if torch.cuda.is_available():\n", " print(f'GPU: {torch.cuda.get_device_name(0)}')\n", " print(f'VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Cell 2: Login to HuggingFace (needed to push model)\n", "from huggingface_hub import login\n", "login() # Will prompt for your token" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Cell 3: Configuration\n", "# ā•”ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•—\n", "# ā•‘ MODIFY THESE SETTINGS AS NEEDED ā•‘\n", "# ā•šā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•ā•\n", "\n", "MODEL_NAME = \"Qwen/Qwen2.5-Coder-0.5B-Instruct\" # Base model\n", "DATASET_ID = \"oxdev/smart-contract-security-audit-v2\" # 50K real findings\n", "HUB_MODEL_ID = \"oxdev/security-auditor-grpo\" # Where to push\n", "OUTPUT_DIR = \"/content/grpo_v2_output\" # Local output\n", "\n", "# Training hyperparameters (tuned for T4 16GB)\n", "SUBSET_SIZE = 2000 # Samples to train on (2K fits in ~4hrs on T4)\n", "BATCH_SIZE = 2 # Per-device batch size\n", "GRAD_ACCUM = 4 # Gradient accumulation ā†’ effective batch = 8\n", "NUM_GENERATIONS = 2 # GRPO generations per prompt\n", "MAX_COMPLETION_LENGTH = 512 # Max tokens per completion\n", "LEARNING_RATE = 1e-6\n", "BETA = 0.04 # KL penalty\n", "NUM_EPOCHS = 1\n", "SAVE_STEPS = 100\n", "\n", "print(f'Config ready: {SUBSET_SIZE} samples, batch={BATCH_SIZE}Ɨ{GRAD_ACCUM}, lr={LEARNING_RATE}')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Cell 4: Load and inspect dataset\n", "from datasets import load_dataset\n", "from collections import Counter\n", "\n", "print('Loading dataset...')\n", "dataset = load_dataset(DATASET_ID, split='train')\n", "print(f'Total: {len(dataset)} samples')\n", "print(f'Columns: {dataset.column_names}')\n", "print()\n", "\n", "# Show distributions\n", "sev_dist = Counter(dataset['severity'])\n", "cat_dist = Counter(dataset['category'])\n", "src_dist = Counter(dataset['source'])\n", "\n", "print('Severity distribution:')\n", "for sev, count in sorted(sev_dist.items(), key=lambda x: -x[1]):\n", " print(f' {sev:15s}: {count:6d} ({count/len(dataset)*100:.1f}%)')\n", "\n", "print(f'\\nCategory distribution (top 10):')\n", "for cat, count in sorted(cat_dist.items(), key=lambda x: -x[1])[:10]:\n", " print(f' {cat:20s}: {count:6d}')\n", "\n", "print(f'\\nSource distribution:')\n", "for src, count in sorted(src_dist.items(), key=lambda x: -x[1]):\n", " print(f' {src:20s}: {count:6d}')\n", "\n", "# Show a sample\n", "print(f'\\n--- Sample prompt (first 300 chars) ---')\n", "p = dataset[0]['prompt']\n", "user_msg = [m for m in p if m['role'] == 'user'][0]['content']\n", "print(user_msg[:300])" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Cell 5: Curate high-quality training subset\n", "print(f'Selecting top {SUBSET_SIZE} highest-value samples...')\n", "\n", "indices = []\n", "idx_set = set()\n", "\n", "# Priority 1: HIGH+CRITICAL severity with code (most valuable)\n", "for i, row in enumerate(dataset):\n", " if row['severity'] in ('high', 'critical') and row['has_code']:\n", " indices.append(i)\n", " idx_set.add(i)\n", "print(f' HIGH+CRITICAL with code: {len(indices)}')\n", "\n", "# Priority 2: Any with PoC reference\n", "for i, row in enumerate(dataset):\n", " if row['has_poc'] and i not in idx_set:\n", " indices.append(i)\n", " idx_set.add(i)\n", "print(f' + Has PoC: {len(indices)}')\n", "\n", "# Priority 3: MEDIUM with code (fill to cap)\n", "for i, row in enumerate(dataset):\n", " if row['severity'] == 'medium' and row['has_code'] and i not in idx_set:\n", " indices.append(i)\n", " idx_set.add(i)\n", " if len(indices) >= SUBSET_SIZE:\n", " break\n", "\n", "# If still short, add remaining HIGH+CRITICAL without code\n", "if len(indices) < SUBSET_SIZE:\n", " for i, row in enumerate(dataset):\n", " if row['severity'] in ('high', 'critical') and i not in idx_set:\n", " indices.append(i)\n", " idx_set.add(i)\n", " if len(indices) >= SUBSET_SIZE:\n", " break\n", "\n", "train_dataset = dataset.select(indices[:SUBSET_SIZE])\n", "print(f'\\nāœ… Final subset: {len(train_dataset)} samples')\n", "\n", "# Show final distribution\n", "final_sev = Counter(train_dataset['severity'])\n", "for sev, count in sorted(final_sev.items(), key=lambda x: -x[1]):\n", " print(f' {sev:15s}: {count:6d}')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Cell 6: Define reward functions\n", "import re\n", "\n", "def format_reward(prompts, completions, completion_ids=None, **kwargs):\n", " \"\"\"Reward for producing structured FINDING blocks and proper formatting.\"\"\"\n", " rewards = []\n", " for completion in completions:\n", " text = completion[0]['content'] if isinstance(completion, list) else str(completion)\n", " reward = 0.0\n", " if re.search(r'FINDING\\s*\\|', text):\n", " reward += 0.3\n", " fields = ['contract:', 'function:', 'bug_class:', 'confidence:']\n", " reward += 0.05 * sum(1 for f in fields if f in text)\n", " if re.search(r'```solidity', text):\n", " reward += 0.15\n", " section_keywords = ['description', 'impact', 'proof', 'fix', 'recommendation', 'mitigation']\n", " sect_count = sum(1 for kw in section_keywords if re.search(rf'(?i)(###?\\s*{kw}|{kw}:)', text))\n", " reward += 0.05 * min(sect_count, 3)\n", " if len(text) < 50: reward -= 0.3\n", " elif len(text) > 4000: reward -= 0.1\n", " rewards.append(max(-1.0, min(1.0, reward)))\n", " return rewards\n", "\n", "\n", "def _sev_rank(sev):\n", " return {'critical': 5, 'high': 4, 'medium': 3, 'low': 2, 'informational': 1, 'gas': 0}.get(sev, -1)\n", "\n", "def severity_reward(prompts, completions, completion_ids=None, severity=None, **kwargs):\n", " \"\"\"Reward for correctly identifying the severity level.\"\"\"\n", " rewards = []\n", " if severity is None:\n", " return [0.0] * len(completions)\n", " sev_list = severity if isinstance(severity, list) else [severity] * len(completions)\n", " for i, completion in enumerate(completions):\n", " text = completion[0]['content'] if isinstance(completion, list) else str(completion)\n", " gt_sev = sev_list[i] if i < len(sev_list) else 'unknown'\n", " if gt_sev == 'unknown':\n", " rewards.append(0.0); continue\n", " sev_match = re.search(r'(?i)(critical|high|medium|low|informational|gas)', text.lower())\n", " if not sev_match:\n", " rewards.append(-0.3)\n", " else:\n", " pred = sev_match.group(1).lower()\n", " diff = abs(_sev_rank(pred) - _sev_rank(gt_sev))\n", " rewards.append(1.0 if diff == 0 else 0.3 if diff == 1 else -0.5)\n", " return rewards\n", "\n", "\n", "CATEGORY_KEYWORDS = {\n", " 'reentrancy': ['reentrancy', 'reentrant', 're-enter', 'callback'],\n", " 'access-control': ['access control', 'unauthorized', 'permission', 'onlyowner', 'role', 'privilege'],\n", " 'oracle': ['oracle', 'price feed', 'chainlink', 'twap', 'price manipulation'],\n", " 'flash-loan': ['flash loan', 'flashloan'],\n", " 'overflow': ['overflow', 'underflow', 'arithmetic'],\n", " 'front-running': ['front-run', 'frontrun', 'sandwich', 'mev'],\n", " 'dos': ['denial of service', 'dos', 'gas limit', 'unbounded', 'out of gas'],\n", " 'token': ['erc20', 'erc721', 'token', 'fee-on-transfer', 'rebasing'],\n", " 'storage': ['storage collision', 'delegatecall', 'proxy', 'slot'],\n", " 'cross-chain': ['bridge', 'cross-chain', 'relay', 'message passing'],\n", " 'liquidation': ['liquidation', 'collateral', 'health factor'],\n", " 'signature': ['signature', 'ecrecover', 'replay', 'nonce', 'eip712'],\n", " 'initialization': ['initialize', 'constructor', 'uninitialized'],\n", " 'rounding': ['rounding', 'precision', 'truncation', 'decimal'],\n", " 'logic': ['logic error', 'incorrect calculation', 'business logic'],\n", "}\n", "\n", "def category_reward(prompts, completions, completion_ids=None, category=None, **kwargs):\n", " \"\"\"Reward for identifying the correct vulnerability category.\"\"\"\n", " rewards = []\n", " if category is None:\n", " return [0.0] * len(completions)\n", " cat_list = category if isinstance(category, list) else [category] * len(completions)\n", " for i, completion in enumerate(completions):\n", " text = completion[0]['content'] if isinstance(completion, list) else str(completion)\n", " gt_cat = cat_list[i] if i < len(cat_list) else 'other'\n", " if gt_cat in ('other', 'unknown'):\n", " rewards.append(0.0); continue\n", " gt_keywords = CATEGORY_KEYWORDS.get(gt_cat, [])\n", " if not gt_keywords:\n", " rewards.append(0.0); continue\n", " hits = sum(1 for kw in gt_keywords if kw in text.lower())\n", " if hits >= 2: rewards.append(1.0)\n", " elif hits == 1: rewards.append(0.5)\n", " else:\n", " any_hit = any(kw in text.lower() for kws in CATEGORY_KEYWORDS.values() for kw in kws)\n", " rewards.append(-0.2 if any_hit else -0.5)\n", " return rewards\n", "\n", "\n", "def quality_reward(prompts, completions, completion_ids=None, **kwargs):\n", " \"\"\"Reward for overall response quality: technical depth, actionability.\"\"\"\n", " rewards = []\n", " for completion in completions:\n", " text = completion[0]['content'] if isinstance(completion, list) else str(completion)\n", " reward = 0.0\n", " technical_terms = [\n", " 'msg.sender', 'tx.origin', 'delegatecall', 'selfdestruct',\n", " 'transfer', 'call.value', 'abi.encode', 'keccak256',\n", " 'require(', 'assert(', 'revert', 'mapping', 'storage',\n", " 'memory', 'calldata', 'modifier', 'interface', 'pragma',\n", " 'assembly', 'unchecked', 'payable', 'receive()', 'fallback()',\n", " ]\n", " reward += min(0.3, 0.03 * sum(1 for t in technical_terms if t in text))\n", " reasoning = ['because', 'therefore', 'this means', 'as a result',\n", " 'the attacker can', 'this allows', 'leading to',\n", " 'step 1', 'step 2', 'first,', 'then,', 'finally,']\n", " reward += min(0.3, 0.06 * sum(1 for r in reasoning if r.lower() in text.lower()))\n", " fix_ind = ['fix:', 'recommendation:', 'mitigation:', 'should', 'consider', 'instead']\n", " reward += min(0.2, 0.05 * sum(1 for f in fix_ind if f.lower() in text.lower()))\n", " if re.search(r'line\\s+\\d+|L\\d+|#L\\d+', text): reward += 0.1\n", " if re.search(r'function\\s+\\w+\\s*\\(', text): reward += 0.1\n", " generic = ['i cannot', \"i don't\", 'no vulnerabilities found', 'the code looks safe']\n", " if any(p in text.lower() for p in generic): reward -= 0.5\n", " rewards.append(max(-1.0, min(1.0, reward)))\n", " return rewards\n", "\n", "print('āœ… 4 reward functions defined: format, severity, category, quality')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Cell 7: Initialize GRPO Trainer\n", "from trl import GRPOTrainer, GRPOConfig\n", "\n", "config = GRPOConfig(\n", " output_dir=OUTPUT_DIR,\n", " num_train_epochs=NUM_EPOCHS,\n", " per_device_train_batch_size=BATCH_SIZE,\n", " gradient_accumulation_steps=GRAD_ACCUM,\n", " num_generations=NUM_GENERATIONS,\n", " max_completion_length=MAX_COMPLETION_LENGTH,\n", " learning_rate=LEARNING_RATE,\n", " beta=BETA,\n", " scale_rewards=True,\n", " reward_weights=[0.25, 0.25, 0.25, 0.25],\n", " gradient_checkpointing=True,\n", " bf16=True,\n", " logging_steps=10,\n", " logging_first_step=True,\n", " logging_strategy='steps',\n", " disable_tqdm=False, # Show progress bar in Colab\n", " save_strategy='steps',\n", " save_steps=SAVE_STEPS,\n", " save_total_limit=2,\n", " push_to_hub=False, # We push manually at the end\n", " log_completions=False,\n", " report_to='none',\n", " seed=42,\n", ")\n", "\n", "print('Initializing GRPOTrainer...')\n", "trainer = GRPOTrainer(\n", " model=MODEL_NAME,\n", " args=config,\n", " reward_funcs=[format_reward, severity_reward, category_reward, quality_reward],\n", " train_dataset=train_dataset,\n", ")\n", "print(f'āœ… GRPOTrainer ready! {len(train_dataset)} samples, ~{len(train_dataset) // (BATCH_SIZE * GRAD_ACCUM)} steps')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Cell 8: TRAIN! šŸš€\n", "# This takes 4-6 hours on T4. Colab will keep running if you stay connected.\n", "# Tip: Keep the tab open and active to prevent disconnection.\n", "\n", "import time\n", "start = time.time()\n", "print('šŸš€ Starting GRPO V2 training...')\n", "print(f'Estimated time: ~{len(train_dataset) / (BATCH_SIZE * GRAD_ACCUM) * 45 / 3600:.1f} hours')\n", "print()\n", "\n", "trainer.train()\n", "\n", "elapsed = time.time() - start\n", "print(f'\\nāœ… Training complete in {elapsed/3600:.1f} hours!')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Cell 9: Save and push to Hub\n", "import os\n", "from huggingface_hub import HfApi\n", "\n", "print(f'Saving model to {OUTPUT_DIR}...')\n", "trainer.save_model(OUTPUT_DIR)\n", "\n", "print(f'Pushing to Hub: {HUB_MODEL_ID}...')\n", "api = HfApi()\n", "api.create_repo(repo_id=HUB_MODEL_ID, exist_ok=True)\n", "\n", "# Upload model files (skip checkpoints and optimizer states to save time)\n", "api.upload_folder(\n", " folder_path=OUTPUT_DIR,\n", " repo_id=HUB_MODEL_ID,\n", " commit_message='GRPO V2 ā€” trained on real audit findings, 4 reward functions',\n", " ignore_patterns=['checkpoint-*', '*.pt'], # Skip checkpoints\n", ")\n", "\n", "print(f'\\nšŸŽ‰ Model pushed to https://huggingface.co/{HUB_MODEL_ID}')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Cell 10: Quick inference test\n", "from transformers import pipeline as hf_pipeline\n", "\n", "print('Loading trained model for inference...')\n", "pipe = hf_pipeline('text-generation', model=OUTPUT_DIR, device=0, torch_dtype=torch.bfloat16)\n", "\n", "test_contract = \"\"\"\n", "pragma solidity ^0.8.0;\n", "\n", "contract SimpleBank {\n", " mapping(address => uint256) public balances;\n", "\n", " function deposit() public payable {\n", " balances[msg.sender] += msg.value;\n", " }\n", "\n", " function withdraw(uint256 amount) public {\n", " require(balances[msg.sender] >= amount);\n", " (bool success, ) = msg.sender.call{value: amount}(\\\"\\\");\n", " require(success);\n", " balances[msg.sender] -= amount;\n", " }\n", "}\n", "\"\"\"\n", "\n", "messages = [\n", " {'role': 'system', 'content': 'You are an expert smart contract security auditor. Analyze the provided Solidity code for vulnerabilities.'},\n", " {'role': 'user', 'content': f'Audit this contract:\\n```solidity\\n{test_contract}\\n```'},\n", "]\n", "\n", "result = pipe(messages, max_new_tokens=512, do_sample=False, return_full_text=False)\n", "output = result[0]['generated_text']\n", "if isinstance(output, list):\n", " output = output[-1]['content']\n", "\n", "print('\\n=== Audit Result ===')\n", "print(output)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "---\n", "\n", "## šŸŽ‰ Done!\n", "\n", "Your V2 model is now pushed to the Hub. Test it interactively at:\n", "\n", "**Demo Space:** [oxdev/security-auditor-demo](https://huggingface.co/spaces/oxdev/security-auditor-demo)\n", "\n", "**Model:** [oxdev/security-auditor-grpo](https://huggingface.co/oxdev/security-auditor-grpo)\n", "\n", "### Next Steps\n", "- Train on more data: increase `SUBSET_SIZE` to 5000 or 10000\n", "- Use a bigger model: try `Qwen/Qwen2.5-Coder-1.5B-Instruct` (needs A100)\n", "- Fine-tune rewards: adjust weights in `reward_weights`\n", "- Try different hyperparameters: learning rate, beta, num_generations" ] } ] }