483 lines
21 KiB
Plaintext
483 lines
21 KiB
Plaintext
{
|
||
"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"
|
||
]
|
||
}
|
||
]
|
||
}
|