初始化项目，由ModelHub XC社区提供模型
Model: oxdev/security-auditor-grpo Source: Original Platform
This commit is contained in:
ModelHub XC
36
.gitattributes
vendored
Normal file
36
.gitattributes
vendored
Normal file
@@ -0,0 +1,36 @@
|
||||
*.7z filter=lfs diff=lfs merge=lfs -text
|
||||
*.arrow filter=lfs diff=lfs merge=lfs -text
|
||||
*.bin filter=lfs diff=lfs merge=lfs -text
|
||||
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
||||
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
||||
*.ftz filter=lfs diff=lfs merge=lfs -text
|
||||
*.gz filter=lfs diff=lfs merge=lfs -text
|
||||
*.h5 filter=lfs diff=lfs merge=lfs -text
|
||||
*.joblib filter=lfs diff=lfs merge=lfs -text
|
||||
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
||||
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
||||
*.model filter=lfs diff=lfs merge=lfs -text
|
||||
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
||||
*.npy filter=lfs diff=lfs merge=lfs -text
|
||||
*.npz filter=lfs diff=lfs merge=lfs -text
|
||||
*.onnx filter=lfs diff=lfs merge=lfs -text
|
||||
*.ot filter=lfs diff=lfs merge=lfs -text
|
||||
*.parquet filter=lfs diff=lfs merge=lfs -text
|
||||
*.pb filter=lfs diff=lfs merge=lfs -text
|
||||
*.pickle filter=lfs diff=lfs merge=lfs -text
|
||||
*.pkl filter=lfs diff=lfs merge=lfs -text
|
||||
*.pt filter=lfs diff=lfs merge=lfs -text
|
||||
*.pth filter=lfs diff=lfs merge=lfs -text
|
||||
*.rar filter=lfs diff=lfs merge=lfs -text
|
||||
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
||||
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
||||
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
||||
*.tar filter=lfs diff=lfs merge=lfs -text
|
||||
*.tflite filter=lfs diff=lfs merge=lfs -text
|
||||
*.tgz filter=lfs diff=lfs merge=lfs -text
|
||||
*.wasm filter=lfs diff=lfs merge=lfs -text
|
||||
*.xz filter=lfs diff=lfs merge=lfs -text
|
||||
*.zip filter=lfs diff=lfs merge=lfs -text
|
||||
*.zst filter=lfs diff=lfs merge=lfs -text
|
||||
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
||||
tokenizer.json filter=lfs diff=lfs merge=lfs -text
|
||||
165
README.md
Normal file
165
README.md
Normal file
@@ -0,0 +1,165 @@
|
||||
---
|
||||
base_model: Qwen/Qwen2.5-Coder-0.5B-Instruct
|
||||
library_name: transformers
|
||||
model_name: security-auditor-grpo
|
||||
tags:
|
||||
- generated_from_trainer
|
||||
- grpo
|
||||
- trl
|
||||
- security
|
||||
- smart-contracts
|
||||
- solidity
|
||||
- audit
|
||||
- web3
|
||||
license: apache-2.0
|
||||
datasets:
|
||||
- oxdev/smart-contract-security-sft
|
||||
- oxdev/smart-contract-security-audit-v2
|
||||
pipeline_tag: text-generation
|
||||
language:
|
||||
- en
|
||||
---
|
||||
|
||||
# 🔐 Smart Contract Security Auditor (GRPO)
|
||||
|
||||
A specialized **smart contract security auditor** built on [Qwen2.5-Coder-0.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-Coder-0.5B-Instruct), fine-tuned using **Group Relative Policy Optimization (GRPO)** on real-world audit findings from top security firms.
|
||||
|
||||
## 🎯 What It Does
|
||||
|
||||
Given a Solidity smart contract, this model identifies security vulnerabilities and produces structured audit findings with:
|
||||
- Vulnerability classification (reentrancy, access control, oracle manipulation, etc.)
|
||||
- Severity assessment (Critical/High/Medium/Low)
|
||||
- Detailed description of the vulnerability
|
||||
- Impact analysis
|
||||
- Proof of concept exploit code
|
||||
- Recommended fixes
|
||||
|
||||
## Quick Start
|
||||
|
||||
```python
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"oxdev/security-auditor-grpo",
|
||||
use_cache=True, # Important: config has use_cache=False from training
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained("oxdev/security-auditor-grpo")
|
||||
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device="cuda")
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": "You are an expert smart contract security auditor. Analyze the provided Solidity code for vulnerabilities."},
|
||||
{"role": "user", "content": """Audit this contract:
|
||||
```solidity
|
||||
contract SimpleBank {
|
||||
mapping(address => uint256) public balances;
|
||||
function deposit() public payable { balances[msg.sender] += msg.value; }
|
||||
function withdraw(uint256 amount) public {
|
||||
require(balances[msg.sender] >= amount);
|
||||
(bool success, ) = msg.sender.call{value: amount}("");
|
||||
require(success);
|
||||
balances[msg.sender] -= amount;
|
||||
}
|
||||
}
|
||||
```"""},
|
||||
]
|
||||
|
||||
result = pipe(messages, max_new_tokens=512, do_sample=False, return_full_text=False)
|
||||
output = result[0]["generated_text"]
|
||||
if isinstance(output, list):
|
||||
output = output[-1]["content"]
|
||||
print(output)
|
||||
```
|
||||
|
||||
## 🔗 Try It Live
|
||||
|
||||
**Interactive Demo:** [oxdev/security-auditor-demo](https://huggingface.co/spaces/oxdev/security-auditor-demo) — Side-by-side comparison with base model, 7 test cases with known vulnerabilities, automated scoring.
|
||||
|
||||
## Training Details
|
||||
|
||||
### V1 (Current Model)
|
||||
- **Method:** GRPO (Group Relative Policy Optimization)
|
||||
- **Base Model:** Qwen2.5-Coder-0.5B-Instruct
|
||||
- **Dataset:** [oxdev/smart-contract-security-sft](https://huggingface.co/datasets/oxdev/smart-contract-security-sft) (327 synthetic samples)
|
||||
- **Hardware:** NVIDIA T4 (16GB)
|
||||
- **Epochs:** 2
|
||||
- **Reward Functions:** Format compliance, finding rate
|
||||
- **Results:**
|
||||
- Format reward: 0.025 → 0.40 (**16× improvement**)
|
||||
- Finding rate: 0% → 50-75%
|
||||
- Mean reward: -0.34 → -0.006
|
||||
|
||||
### V2 (Pending — Colab Notebook Ready)
|
||||
- **Dataset:** [oxdev/smart-contract-security-audit-v2](https://huggingface.co/datasets/oxdev/smart-contract-security-audit-v2) (50,902 real audit findings)
|
||||
- **Sources:** SkywardNomad92/smart-contract-audit-findings, samscrack/cyfrin-audit-findings, Solodit API
|
||||
- **4 Reward Functions:** Format (0.25), Severity matching (0.25), Category matching (0.25), Quality (0.25)
|
||||
- **Train on Colab:** Open [`train_grpo_v2_colab.ipynb`](https://huggingface.co/oxdev/security-auditor-grpo/blob/main/train_grpo_v2_colab.ipynb) in Google Colab with a free T4 GPU
|
||||
|
||||
## Vulnerability Categories Covered
|
||||
|
||||
| Category | Keywords |
|
||||
|----------|----------|
|
||||
| Reentrancy | reentrancy, reentrant, callback |
|
||||
| Access Control | unauthorized, permission, onlyowner |
|
||||
| Oracle Manipulation | price feed, chainlink, twap |
|
||||
| Flash Loan | flash loan, flashloan |
|
||||
| Overflow/Underflow | overflow, underflow, arithmetic |
|
||||
| Front-running | front-run, sandwich, MEV |
|
||||
| DoS | denial of service, gas limit, unbounded |
|
||||
| Token Issues | fee-on-transfer, rebasing, ERC20 |
|
||||
| Storage | storage collision, delegatecall, proxy |
|
||||
| Cross-chain | bridge, relay, message passing |
|
||||
| Liquidation | liquidation, collateral, health factor |
|
||||
| Signature | ecrecover, replay, nonce, EIP712 |
|
||||
| Initialization | uninitialized, constructor |
|
||||
| Rounding | precision, truncation, decimal |
|
||||
|
||||
## Architecture
|
||||
|
||||
- **Model:** Qwen2ForCausalLM
|
||||
- **Parameters:** 0.5B
|
||||
- **Hidden Size:** 896
|
||||
- **Layers:** 24
|
||||
- **Attention Heads:** 14 (2 KV heads)
|
||||
- **Context Length:** 32,768 tokens
|
||||
- **Chat Template:** ChatML (`<|im_start|>` / `<|im_end|>`)
|
||||
|
||||
## ⚠️ Important Notes
|
||||
|
||||
1. **Set `use_cache=True`** when loading for inference — the saved config has `use_cache=False` from training, which makes generation 10-20× slower
|
||||
2. **This is a 0.5B model** — it's fast but not as capable as larger models. Use it for quick triage, not as a replacement for professional audits
|
||||
3. **V1 was trained on 327 samples** — V2 training on 50K real findings will significantly improve quality
|
||||
|
||||
## Files
|
||||
|
||||
| File | Description |
|
||||
|------|-------------|
|
||||
| `model.safetensors` | V1 trained model weights (1.8GB) |
|
||||
| `train_grpo_job.py` | V1 training script |
|
||||
| `train_grpo_v2.py` | V2 training script (4 reward functions) |
|
||||
| `train_grpo_v2_colab.ipynb` | V2 Colab notebook (free T4 GPU) |
|
||||
| `checkpoint-300/` | V1 training checkpoint |
|
||||
| `checkpoint-326/` | V1 final checkpoint |
|
||||
|
||||
## Related Resources
|
||||
|
||||
- **GitHub:** [0xedev/skills](https://github.com/0xedev/skills) — Pashov Audit Group AI-powered security skills
|
||||
- **V2 Dataset:** [oxdev/smart-contract-security-audit-v2](https://huggingface.co/datasets/oxdev/smart-contract-security-audit-v2)
|
||||
- **Demo Space:** [oxdev/security-auditor-demo](https://huggingface.co/spaces/oxdev/security-auditor-demo)
|
||||
|
||||
## Framework Versions
|
||||
|
||||
- TRL: 1.2.0
|
||||
- Transformers: 5.6.2
|
||||
- PyTorch: 2.6.0+cu126
|
||||
- Datasets: 4.8.4
|
||||
|
||||
## Citations
|
||||
|
||||
```bibtex
|
||||
@article{shao2024deepseekmath,
|
||||
title = {{DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models}},
|
||||
author = {Zhihong Shao and Peiyi Wang and Qihao Zhu and Runxin Xu and Junxiao Song and others},
|
||||
year = 2024,
|
||||
eprint = {arXiv:2402.03300},
|
||||
}
|
||||
```
|
||||
54
chat_template.jinja
Normal file
54
chat_template.jinja
Normal file
@@ -0,0 +1,54 @@
|
||||
{%- if tools %}
|
||||
{{- '<|im_start|>system\n' }}
|
||||
{%- if messages[0]['role'] == 'system' %}
|
||||
{{- messages[0]['content'] }}
|
||||
{%- else %}
|
||||
{{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}
|
||||
{%- endif %}
|
||||
{{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
|
||||
{%- for tool in tools %}
|
||||
{{- "\n" }}
|
||||
{{- tool | tojson }}
|
||||
{%- endfor %}
|
||||
{{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
|
||||
{%- else %}
|
||||
{%- if messages[0]['role'] == 'system' %}
|
||||
{{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
|
||||
{%- else %}
|
||||
{{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- for message in messages %}
|
||||
{%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
|
||||
{{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
|
||||
{%- elif message.role == "assistant" %}
|
||||
{{- '<|im_start|>' + message.role }}
|
||||
{%- if message.content %}
|
||||
{{- '\n' + message.content }}
|
||||
{%- endif %}
|
||||
{%- for tool_call in message.tool_calls %}
|
||||
{%- if tool_call.function is defined %}
|
||||
{%- set tool_call = tool_call.function %}
|
||||
{%- endif %}
|
||||
{{- '\n<tool_call>\n{"name": "' }}
|
||||
{{- tool_call.name }}
|
||||
{{- '", "arguments": ' }}
|
||||
{{- tool_call.arguments | tojson }}
|
||||
{{- '}\n</tool_call>' }}
|
||||
{%- endfor %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- elif message.role == "tool" %}
|
||||
{%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
|
||||
{{- '<|im_start|>user' }}
|
||||
{%- endif %}
|
||||
{{- '\n<tool_response>\n' }}
|
||||
{{- message.content }}
|
||||
{{- '\n</tool_response>' }}
|
||||
{%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|im_start|>assistant\n' }}
|
||||
{%- endif %}
|
||||
54
checkpoint-300/chat_template.jinja
Normal file
54
checkpoint-300/chat_template.jinja
Normal file
@@ -0,0 +1,54 @@
|
||||
{%- if tools %}
|
||||
{{- '<|im_start|>system\n' }}
|
||||
{%- if messages[0]['role'] == 'system' %}
|
||||
{{- messages[0]['content'] }}
|
||||
{%- else %}
|
||||
{{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}
|
||||
{%- endif %}
|
||||
{{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
|
||||
{%- for tool in tools %}
|
||||
{{- "\n" }}
|
||||
{{- tool | tojson }}
|
||||
{%- endfor %}
|
||||
{{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
|
||||
{%- else %}
|
||||
{%- if messages[0]['role'] == 'system' %}
|
||||
{{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
|
||||
{%- else %}
|
||||
{{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- for message in messages %}
|
||||
{%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
|
||||
{{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
|
||||
{%- elif message.role == "assistant" %}
|
||||
{{- '<|im_start|>' + message.role }}
|
||||
{%- if message.content %}
|
||||
{{- '\n' + message.content }}
|
||||
{%- endif %}
|
||||
{%- for tool_call in message.tool_calls %}
|
||||
{%- if tool_call.function is defined %}
|
||||
{%- set tool_call = tool_call.function %}
|
||||
{%- endif %}
|
||||
{{- '\n<tool_call>\n{"name": "' }}
|
||||
{{- tool_call.name }}
|
||||
{{- '", "arguments": ' }}
|
||||
{{- tool_call.arguments | tojson }}
|
||||
{{- '}\n</tool_call>' }}
|
||||
{%- endfor %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- elif message.role == "tool" %}
|
||||
{%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
|
||||
{{- '<|im_start|>user' }}
|
||||
{%- endif %}
|
||||
{{- '\n<tool_response>\n' }}
|
||||
{{- message.content }}
|
||||
{{- '\n</tool_response>' }}
|
||||
{%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|im_start|>assistant\n' }}
|
||||
{%- endif %}
|
||||
57
checkpoint-300/config.json
Normal file
57
checkpoint-300/config.json
Normal file
@@ -0,0 +1,57 @@
|
||||
{
|
||||
"architectures": [
|
||||
"Qwen2ForCausalLM"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": null,
|
||||
"dtype": "float32",
|
||||
"eos_token_id": 151645,
|
||||
"hidden_act": "silu",
|
||||
"hidden_size": 896,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 4864,
|
||||
"layer_types": [
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention"
|
||||
],
|
||||
"max_position_embeddings": 32768,
|
||||
"max_window_layers": 24,
|
||||
"model_type": "qwen2",
|
||||
"num_attention_heads": 14,
|
||||
"num_hidden_layers": 24,
|
||||
"num_key_value_heads": 2,
|
||||
"pad_token_id": 151643,
|
||||
"rms_norm_eps": 1e-06,
|
||||
"rope_parameters": {
|
||||
"rope_theta": 1000000.0,
|
||||
"rope_type": "default"
|
||||
},
|
||||
"sliding_window": null,
|
||||
"tie_word_embeddings": true,
|
||||
"transformers_version": "5.6.2",
|
||||
"use_cache": false,
|
||||
"use_sliding_window": false,
|
||||
"vocab_size": 151936
|
||||
}
|
||||
13
checkpoint-300/generation_config.json
Normal file
13
checkpoint-300/generation_config.json
Normal file
@@ -0,0 +1,13 @@
|
||||
{
|
||||
"do_sample": true,
|
||||
"eos_token_id": [
|
||||
151645,
|
||||
151643
|
||||
],
|
||||
"pad_token_id": 151643,
|
||||
"repetition_penalty": 1.05,
|
||||
"temperature": 0.7,
|
||||
"top_k": 20,
|
||||
"top_p": 0.8,
|
||||
"transformers_version": "5.6.2"
|
||||
}
|
||||
3
checkpoint-300/model.safetensors
Normal file
3
checkpoint-300/model.safetensors
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:275ebac6d53742a54d972d5a2fdf93a64ab774cf50af3c817e02a1376655c840
|
||||
size 1976163472
|
||||
3
checkpoint-300/optimizer.pt
Normal file
3
checkpoint-300/optimizer.pt
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3577491d619a3e9b2d76cba84e6eee9cdffd5bb2784ebc6a1e3453f2ce9f8021
|
||||
size 3952505274
|
||||
3
checkpoint-300/rng_state.pth
Normal file
3
checkpoint-300/rng_state.pth
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:fa9a1789e81962b242729edabc19959b88ccde1eb3dfdbc7cd826e14f85a76f9
|
||||
size 14244
|
||||
3
checkpoint-300/scheduler.pt
Normal file
3
checkpoint-300/scheduler.pt
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:262427cf509faa4beebbf93a0c170cf18cb00c5f988d14d843ea44ed3b3c2cae
|
||||
size 1064
|
||||
3
checkpoint-300/tokenizer.json
Normal file
3
checkpoint-300/tokenizer.json
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3fd169731d2cbde95e10bf356d66d5997fd885dd8dbb6fb4684da3f23b2585d8
|
||||
size 11421892
|
||||
32
checkpoint-300/tokenizer_config.json
Normal file
32
checkpoint-300/tokenizer_config.json
Normal file
@@ -0,0 +1,32 @@
|
||||
{
|
||||
"add_prefix_space": false,
|
||||
"backend": "tokenizers",
|
||||
"bos_token": null,
|
||||
"clean_up_tokenization_spaces": false,
|
||||
"eos_token": "<|im_end|>",
|
||||
"errors": "replace",
|
||||
"extra_special_tokens": [
|
||||
"<|im_start|>",
|
||||
"<|im_end|>",
|
||||
"<|object_ref_start|>",
|
||||
"<|object_ref_end|>",
|
||||
"<|box_start|>",
|
||||
"<|box_end|>",
|
||||
"<|quad_start|>",
|
||||
"<|quad_end|>",
|
||||
"<|vision_start|>",
|
||||
"<|vision_end|>",
|
||||
"<|vision_pad|>",
|
||||
"<|image_pad|>",
|
||||
"<|video_pad|>"
|
||||
],
|
||||
"is_local": false,
|
||||
"local_files_only": false,
|
||||
"model_max_length": 32768,
|
||||
"pad_token": "<|endoftext|>",
|
||||
"padding_side": "left",
|
||||
"split_special_tokens": false,
|
||||
"tokenizer_class": "Qwen2Tokenizer",
|
||||
"truncation_side": "left",
|
||||
"unk_token": null
|
||||
}
|
||||
1803
checkpoint-300/trainer_state.json
Normal file
1803
checkpoint-300/trainer_state.json
Normal file
File diff suppressed because it is too large
Load Diff
3
checkpoint-300/training_args.bin
Normal file
3
checkpoint-300/training_args.bin
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:b51f3856815802830b7add9e23ddd089207e5c9941078dd606f120af0f983d09
|
||||
size 6776
|
||||
54
checkpoint-326/chat_template.jinja
Normal file
54
checkpoint-326/chat_template.jinja
Normal file
@@ -0,0 +1,54 @@
|
||||
{%- if tools %}
|
||||
{{- '<|im_start|>system\n' }}
|
||||
{%- if messages[0]['role'] == 'system' %}
|
||||
{{- messages[0]['content'] }}
|
||||
{%- else %}
|
||||
{{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}
|
||||
{%- endif %}
|
||||
{{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
|
||||
{%- for tool in tools %}
|
||||
{{- "\n" }}
|
||||
{{- tool | tojson }}
|
||||
{%- endfor %}
|
||||
{{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
|
||||
{%- else %}
|
||||
{%- if messages[0]['role'] == 'system' %}
|
||||
{{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
|
||||
{%- else %}
|
||||
{{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- for message in messages %}
|
||||
{%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
|
||||
{{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
|
||||
{%- elif message.role == "assistant" %}
|
||||
{{- '<|im_start|>' + message.role }}
|
||||
{%- if message.content %}
|
||||
{{- '\n' + message.content }}
|
||||
{%- endif %}
|
||||
{%- for tool_call in message.tool_calls %}
|
||||
{%- if tool_call.function is defined %}
|
||||
{%- set tool_call = tool_call.function %}
|
||||
{%- endif %}
|
||||
{{- '\n<tool_call>\n{"name": "' }}
|
||||
{{- tool_call.name }}
|
||||
{{- '", "arguments": ' }}
|
||||
{{- tool_call.arguments | tojson }}
|
||||
{{- '}\n</tool_call>' }}
|
||||
{%- endfor %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- elif message.role == "tool" %}
|
||||
{%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
|
||||
{{- '<|im_start|>user' }}
|
||||
{%- endif %}
|
||||
{{- '\n<tool_response>\n' }}
|
||||
{{- message.content }}
|
||||
{{- '\n</tool_response>' }}
|
||||
{%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|im_start|>assistant\n' }}
|
||||
{%- endif %}
|
||||
57
checkpoint-326/config.json
Normal file
57
checkpoint-326/config.json
Normal file
@@ -0,0 +1,57 @@
|
||||
{
|
||||
"architectures": [
|
||||
"Qwen2ForCausalLM"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": null,
|
||||
"dtype": "float32",
|
||||
"eos_token_id": 151645,
|
||||
"hidden_act": "silu",
|
||||
"hidden_size": 896,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 4864,
|
||||
"layer_types": [
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention"
|
||||
],
|
||||
"max_position_embeddings": 32768,
|
||||
"max_window_layers": 24,
|
||||
"model_type": "qwen2",
|
||||
"num_attention_heads": 14,
|
||||
"num_hidden_layers": 24,
|
||||
"num_key_value_heads": 2,
|
||||
"pad_token_id": 151643,
|
||||
"rms_norm_eps": 1e-06,
|
||||
"rope_parameters": {
|
||||
"rope_theta": 1000000.0,
|
||||
"rope_type": "default"
|
||||
},
|
||||
"sliding_window": null,
|
||||
"tie_word_embeddings": true,
|
||||
"transformers_version": "5.6.2",
|
||||
"use_cache": false,
|
||||
"use_sliding_window": false,
|
||||
"vocab_size": 151936
|
||||
}
|
||||
13
checkpoint-326/generation_config.json
Normal file
13
checkpoint-326/generation_config.json
Normal file
@@ -0,0 +1,13 @@
|
||||
{
|
||||
"do_sample": true,
|
||||
"eos_token_id": [
|
||||
151645,
|
||||
151643
|
||||
],
|
||||
"pad_token_id": 151643,
|
||||
"repetition_penalty": 1.05,
|
||||
"temperature": 0.7,
|
||||
"top_k": 20,
|
||||
"top_p": 0.8,
|
||||
"transformers_version": "5.6.2"
|
||||
}
|
||||
3
checkpoint-326/model.safetensors
Normal file
3
checkpoint-326/model.safetensors
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3c64421ad1b2b08b8f84687a636657a68a9f6c9ef639c6c2dc449cb93d2c4219
|
||||
size 1976163472
|
||||
3
checkpoint-326/optimizer.pt
Normal file
3
checkpoint-326/optimizer.pt
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:2029a2d6a0a790f550621622c3c990b0d5f81492de3d7642988e1fa042d3a073
|
||||
size 3952505274
|
||||
3
checkpoint-326/rng_state.pth
Normal file
3
checkpoint-326/rng_state.pth
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:7709cb910b037a9984c235d2fc7fe7fd99ccb8982993a4ca396269149709e777
|
||||
size 14244
|
||||
3
checkpoint-326/scheduler.pt
Normal file
3
checkpoint-326/scheduler.pt
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:f3aa6a7bb4866149bd0f10ff54d9da7e2c37c93aa3f37a2a6471c11bd6760f19
|
||||
size 1064
|
||||
3
checkpoint-326/tokenizer.json
Normal file
3
checkpoint-326/tokenizer.json
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3fd169731d2cbde95e10bf356d66d5997fd885dd8dbb6fb4684da3f23b2585d8
|
||||
size 11421892
|
||||
32
checkpoint-326/tokenizer_config.json
Normal file
32
checkpoint-326/tokenizer_config.json
Normal file
@@ -0,0 +1,32 @@
|
||||
{
|
||||
"add_prefix_space": false,
|
||||
"backend": "tokenizers",
|
||||
"bos_token": null,
|
||||
"clean_up_tokenization_spaces": false,
|
||||
"eos_token": "<|im_end|>",
|
||||
"errors": "replace",
|
||||
"extra_special_tokens": [
|
||||
"<|im_start|>",
|
||||
"<|im_end|>",
|
||||
"<|object_ref_start|>",
|
||||
"<|object_ref_end|>",
|
||||
"<|box_start|>",
|
||||
"<|box_end|>",
|
||||
"<|quad_start|>",
|
||||
"<|quad_end|>",
|
||||
"<|vision_start|>",
|
||||
"<|vision_end|>",
|
||||
"<|vision_pad|>",
|
||||
"<|image_pad|>",
|
||||
"<|video_pad|>"
|
||||
],
|
||||
"is_local": false,
|
||||
"local_files_only": false,
|
||||
"model_max_length": 32768,
|
||||
"pad_token": "<|endoftext|>",
|
||||
"padding_side": "left",
|
||||
"split_special_tokens": false,
|
||||
"tokenizer_class": "Qwen2Tokenizer",
|
||||
"truncation_side": "left",
|
||||
"unk_token": null
|
||||
}
|
||||
1948
checkpoint-326/trainer_state.json
Normal file
1948
checkpoint-326/trainer_state.json
Normal file
File diff suppressed because it is too large
Load Diff
3
checkpoint-326/training_args.bin
Normal file
3
checkpoint-326/training_args.bin
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:b51f3856815802830b7add9e23ddd089207e5c9941078dd606f120af0f983d09
|
||||
size 6776
|
||||
57
config.json
Normal file
57
config.json
Normal file
@@ -0,0 +1,57 @@
|
||||
{
|
||||
"architectures": [
|
||||
"Qwen2ForCausalLM"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": null,
|
||||
"dtype": "float32",
|
||||
"eos_token_id": 151645,
|
||||
"hidden_act": "silu",
|
||||
"hidden_size": 896,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 4864,
|
||||
"layer_types": [
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"full_attention"
|
||||
],
|
||||
"max_position_embeddings": 32768,
|
||||
"max_window_layers": 24,
|
||||
"model_type": "qwen2",
|
||||
"num_attention_heads": 14,
|
||||
"num_hidden_layers": 24,
|
||||
"num_key_value_heads": 2,
|
||||
"pad_token_id": 151643,
|
||||
"rms_norm_eps": 1e-06,
|
||||
"rope_parameters": {
|
||||
"rope_theta": 1000000.0,
|
||||
"rope_type": "default"
|
||||
},
|
||||
"sliding_window": null,
|
||||
"tie_word_embeddings": true,
|
||||
"transformers_version": "5.6.2",
|
||||
"use_cache": false,
|
||||
"use_sliding_window": false,
|
||||
"vocab_size": 151936
|
||||
}
|
||||
13
generation_config.json
Normal file
13
generation_config.json
Normal file
@@ -0,0 +1,13 @@
|
||||
{
|
||||
"do_sample": true,
|
||||
"eos_token_id": [
|
||||
151645,
|
||||
151643
|
||||
],
|
||||
"pad_token_id": 151643,
|
||||
"repetition_penalty": 1.05,
|
||||
"temperature": 0.7,
|
||||
"top_k": 20,
|
||||
"top_p": 0.8,
|
||||
"transformers_version": "5.6.2"
|
||||
}
|
||||
3
model.safetensors
Normal file
3
model.safetensors
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3c64421ad1b2b08b8f84687a636657a68a9f6c9ef639c6c2dc449cb93d2c4219
|
||||
size 1976163472
|
||||
3
tokenizer.json
Normal file
3
tokenizer.json
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3fd169731d2cbde95e10bf356d66d5997fd885dd8dbb6fb4684da3f23b2585d8
|
||||
size 11421892
|
||||
32
tokenizer_config.json
Normal file
32
tokenizer_config.json
Normal file
@@ -0,0 +1,32 @@
|
||||
{
|
||||
"add_prefix_space": false,
|
||||
"backend": "tokenizers",
|
||||
"bos_token": null,
|
||||
"clean_up_tokenization_spaces": false,
|
||||
"eos_token": "<|im_end|>",
|
||||
"errors": "replace",
|
||||
"extra_special_tokens": [
|
||||
"<|im_start|>",
|
||||
"<|im_end|>",
|
||||
"<|object_ref_start|>",
|
||||
"<|object_ref_end|>",
|
||||
"<|box_start|>",
|
||||
"<|box_end|>",
|
||||
"<|quad_start|>",
|
||||
"<|quad_end|>",
|
||||
"<|vision_start|>",
|
||||
"<|vision_end|>",
|
||||
"<|vision_pad|>",
|
||||
"<|image_pad|>",
|
||||
"<|video_pad|>"
|
||||
],
|
||||
"is_local": false,
|
||||
"local_files_only": false,
|
||||
"model_max_length": 32768,
|
||||
"pad_token": "<|endoftext|>",
|
||||
"padding_side": "left",
|
||||
"split_special_tokens": false,
|
||||
"tokenizer_class": "Qwen2Tokenizer",
|
||||
"truncation_side": "left",
|
||||
"unk_token": null
|
||||
}
|
||||
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()
|
||||
356
train_grpo_v2.py
Normal file
356
train_grpo_v2.py
Normal file
@@ -0,0 +1,356 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
train_grpo_v2.py — GRPO training on 50K real audit findings.
|
||||
|
||||
V2 improvements over V1:
|
||||
- 155x more data (50,902 vs 327)
|
||||
- 4 reward functions with ground-truth severity/category matching
|
||||
- Reference-based semantic similarity reward
|
||||
- Better exploration via higher num_generations
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
from collections import Counter
|
||||
|
||||
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-audit-v2"
|
||||
OUTPUT_DIR = "/tmp/grpo_v2_output"
|
||||
HUB_MODEL_ID = "oxdev/security-auditor-grpo"
|
||||
|
||||
|
||||
# ─── Reward Function 1: Structure & Format (weight: 0.25) ────────────────────
|
||||
|
||||
def format_reward(prompts, completions, completion_ids=None, **kwargs):
|
||||
"""Reward for producing structured FINDING blocks and proper formatting."""
|
||||
rewards = []
|
||||
for completion in completions:
|
||||
text = completion[0]["content"] if isinstance(completion, list) else str(completion)
|
||||
reward = 0.0
|
||||
|
||||
# FINDING block present
|
||||
if re.search(r'FINDING\s*\|', text):
|
||||
reward += 0.3
|
||||
# Required fields
|
||||
fields = ['contract:', 'function:', 'bug_class:', 'confidence:']
|
||||
field_count = sum(1 for f in fields if f in text)
|
||||
reward += 0.05 * field_count # up to 0.2 more
|
||||
|
||||
# Has code block
|
||||
if re.search(r'```solidity', text):
|
||||
reward += 0.15
|
||||
|
||||
# Has structured sections
|
||||
section_keywords = ['description', 'impact', 'proof', 'fix', 'recommendation', 'mitigation']
|
||||
section_count = sum(1 for kw in section_keywords if re.search(rf'(?i)(###?\s*{kw}|{kw}:)', text))
|
||||
reward += 0.05 * min(section_count, 3) # up to 0.15
|
||||
|
||||
# Penalize very short or very long
|
||||
if len(text) < 50:
|
||||
reward -= 0.3
|
||||
elif len(text) > 4000:
|
||||
reward -= 0.1
|
||||
|
||||
rewards.append(max(-1.0, min(1.0, reward)))
|
||||
return rewards
|
||||
|
||||
|
||||
# ─── Reward Function 2: Severity Match (weight: 0.25) ────────────────────────
|
||||
|
||||
def severity_reward(prompts, completions, completion_ids=None, severity=None, **kwargs):
|
||||
"""Reward for correctly identifying the severity level."""
|
||||
rewards = []
|
||||
|
||||
if severity is None:
|
||||
return [0.0] * len(completions)
|
||||
|
||||
# Handle batch: severity may be a list
|
||||
if isinstance(severity, list):
|
||||
sev_list = severity
|
||||
else:
|
||||
sev_list = [severity] * len(completions)
|
||||
|
||||
for i, completion in enumerate(completions):
|
||||
text = completion[0]["content"] if isinstance(completion, list) else str(completion)
|
||||
text_lower = text.lower()
|
||||
|
||||
gt_sev = sev_list[i] if i < len(sev_list) else "unknown"
|
||||
if gt_sev == "unknown":
|
||||
rewards.append(0.0)
|
||||
continue
|
||||
|
||||
# Extract predicted severity
|
||||
pred_sev = None
|
||||
sev_match = re.search(r'(?i)(critical|high|medium|low|informational|gas)', text_lower)
|
||||
if sev_match:
|
||||
pred_sev = sev_match.group(1).lower()
|
||||
|
||||
if pred_sev is None:
|
||||
rewards.append(-0.3)
|
||||
elif pred_sev == gt_sev:
|
||||
rewards.append(1.0) # Exact match
|
||||
elif abs(_sev_rank(pred_sev) - _sev_rank(gt_sev)) == 1:
|
||||
rewards.append(0.3) # Off by one level
|
||||
else:
|
||||
rewards.append(-0.5) # Way off
|
||||
|
||||
return rewards
|
||||
|
||||
|
||||
def _sev_rank(sev):
|
||||
ranks = {"critical": 5, "high": 4, "medium": 3, "low": 2, "informational": 1, "gas": 0}
|
||||
return ranks.get(sev, -1)
|
||||
|
||||
|
||||
# ─── Reward Function 3: Vulnerability Category (weight: 0.25) ────────────────
|
||||
|
||||
CATEGORY_KEYWORDS = {
|
||||
"reentrancy": ["reentrancy", "reentrant", "re-enter", "callback"],
|
||||
"access-control": ["access control", "unauthorized", "permission", "onlyowner", "role", "privilege"],
|
||||
"oracle": ["oracle", "price feed", "chainlink", "twap", "price manipulation"],
|
||||
"flash-loan": ["flash loan", "flashloan"],
|
||||
"overflow": ["overflow", "underflow", "arithmetic"],
|
||||
"front-running": ["front-run", "frontrun", "sandwich", "mev"],
|
||||
"dos": ["denial of service", "dos", "gas limit", "unbounded", "out of gas"],
|
||||
"token": ["erc20", "erc721", "token", "fee-on-transfer", "rebasing"],
|
||||
"storage": ["storage collision", "delegatecall", "proxy", "slot"],
|
||||
"cross-chain": ["bridge", "cross-chain", "relay", "message passing"],
|
||||
"liquidation": ["liquidation", "collateral", "health factor"],
|
||||
"signature": ["signature", "ecrecover", "replay", "nonce", "eip712"],
|
||||
"initialization": ["initialize", "constructor", "uninitialized"],
|
||||
"rounding": ["rounding", "precision", "truncation", "decimal"],
|
||||
"logic": ["logic error", "incorrect calculation", "business logic"],
|
||||
}
|
||||
|
||||
def category_reward(prompts, completions, completion_ids=None, category=None, **kwargs):
|
||||
"""Reward for identifying the correct vulnerability category."""
|
||||
rewards = []
|
||||
|
||||
if category is None:
|
||||
return [0.0] * len(completions)
|
||||
|
||||
if isinstance(category, list):
|
||||
cat_list = category
|
||||
else:
|
||||
cat_list = [category] * len(completions)
|
||||
|
||||
for i, completion in enumerate(completions):
|
||||
text = completion[0]["content"] if isinstance(completion, list) else str(completion)
|
||||
text_lower = text.lower()
|
||||
|
||||
gt_cat = cat_list[i] if i < len(cat_list) else "other"
|
||||
if gt_cat == "other" or gt_cat == "unknown":
|
||||
# Can't evaluate — neutral reward
|
||||
rewards.append(0.0)
|
||||
continue
|
||||
|
||||
# Check if the model mentions keywords from the ground truth category
|
||||
gt_keywords = CATEGORY_KEYWORDS.get(gt_cat, [])
|
||||
if not gt_keywords:
|
||||
rewards.append(0.0)
|
||||
continue
|
||||
|
||||
hits = sum(1 for kw in gt_keywords if kw in text_lower)
|
||||
if hits >= 2:
|
||||
rewards.append(1.0)
|
||||
elif hits == 1:
|
||||
rewards.append(0.5)
|
||||
else:
|
||||
# Check if it mentions ANY vulnerability category (at least trying)
|
||||
any_hit = any(kw in text_lower for kws in CATEGORY_KEYWORDS.values() for kw in kws)
|
||||
rewards.append(-0.2 if any_hit else -0.5)
|
||||
|
||||
return rewards
|
||||
|
||||
|
||||
# ─── Reward Function 4: Content Quality (weight: 0.25) ───────────────────────
|
||||
|
||||
def quality_reward(prompts, completions, completion_ids=None, **kwargs):
|
||||
"""Reward for overall response quality: technical depth, actionability."""
|
||||
rewards = []
|
||||
for completion in completions:
|
||||
text = completion[0]["content"] if isinstance(completion, list) else str(completion)
|
||||
reward = 0.0
|
||||
|
||||
# Technical indicators
|
||||
technical_terms = [
|
||||
'msg.sender', 'tx.origin', 'delegatecall', 'selfdestruct',
|
||||
'transfer', 'call.value', 'abi.encode', 'keccak256',
|
||||
'require(', 'assert(', 'revert', 'mapping', 'storage',
|
||||
'memory', 'calldata', 'modifier', 'interface', 'pragma',
|
||||
'assembly', 'unchecked', 'payable', 'receive()', 'fallback()',
|
||||
]
|
||||
tech_count = sum(1 for t in technical_terms if t in text)
|
||||
reward += min(0.3, 0.03 * tech_count)
|
||||
|
||||
# Explanation quality (has reasoning)
|
||||
reasoning_indicators = [
|
||||
'because', 'therefore', 'this means', 'as a result',
|
||||
'the attacker can', 'this allows', 'leading to',
|
||||
'step 1', 'step 2', 'first,', 'then,', 'finally,',
|
||||
]
|
||||
reasoning_count = sum(1 for r in reasoning_indicators if r.lower() in text.lower())
|
||||
reward += min(0.3, 0.06 * reasoning_count)
|
||||
|
||||
# Actionable fix provided
|
||||
fix_indicators = ['fix:', 'recommendation:', 'mitigation:', 'should', 'consider', 'instead']
|
||||
fix_count = sum(1 for f in fix_indicators if f.lower() in text.lower())
|
||||
reward += min(0.2, 0.05 * fix_count)
|
||||
|
||||
# Code reference specificity
|
||||
if re.search(r'line\s+\d+|L\d+|#L\d+', text):
|
||||
reward += 0.1
|
||||
if re.search(r'function\s+\w+\s*\(', text):
|
||||
reward += 0.1
|
||||
|
||||
# Penalize generic/unhelpful responses
|
||||
generic_phrases = ['i cannot', 'i don\'t', 'no vulnerabilities found', 'the code looks safe']
|
||||
if any(p in text.lower() for p in generic_phrases):
|
||||
reward -= 0.5
|
||||
|
||||
rewards.append(max(-1.0, min(1.0, reward)))
|
||||
return rewards
|
||||
|
||||
|
||||
# ─── Main ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
logger.info("=" * 60)
|
||||
logger.info("GRPO V2 Training — 50K Real Audit Findings")
|
||||
logger.info(f"Model: {MODEL_NAME}")
|
||||
logger.info(f"Dataset: {DATASET_ID}")
|
||||
logger.info(f"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}")
|
||||
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...")
|
||||
dataset = load_dataset(DATASET_ID, split="train")
|
||||
logger.info(f"Dataset: {len(dataset)} samples, columns={dataset.column_names}")
|
||||
|
||||
# For GRPO we only need 'prompt' column + metadata columns for reward
|
||||
# The reward functions access metadata via kwargs passed from the dataset
|
||||
|
||||
# Log severity distribution
|
||||
sev_dist = Counter(dataset['severity'])
|
||||
logger.info(f"Severity distribution: {dict(sev_dist)}")
|
||||
|
||||
# Subsample — 5K highest-value samples for A10G (fits in ~6hrs)
|
||||
# Focus on HIGH+CRITICAL with code — most valuable training signal
|
||||
logger.info("Selecting high-quality training subset (5K for A10G)...")
|
||||
indices = []
|
||||
idx_set = set()
|
||||
|
||||
# Priority 1: HIGH+CRITICAL severity with code (most valuable)
|
||||
for i, row in enumerate(dataset):
|
||||
if row['severity'] in ('high', 'critical') and row['has_code']:
|
||||
indices.append(i)
|
||||
idx_set.add(i)
|
||||
logger.info(f" HIGH+CRITICAL with code: {len(indices)}")
|
||||
|
||||
# Priority 2: Any with PoC reference
|
||||
for i, row in enumerate(dataset):
|
||||
if row['has_poc'] and i not in idx_set:
|
||||
indices.append(i)
|
||||
idx_set.add(i)
|
||||
logger.info(f" + Has PoC: {len(indices)}")
|
||||
|
||||
# Priority 3: MEDIUM with code (fill to 5K cap)
|
||||
for i, row in enumerate(dataset):
|
||||
if row['severity'] == 'medium' and row['has_code'] and i not in idx_set:
|
||||
indices.append(i)
|
||||
idx_set.add(i)
|
||||
if len(indices) >= 5000:
|
||||
break
|
||||
logger.info(f" Final subset: {len(indices)} samples")
|
||||
|
||||
train_dataset = dataset.select(indices)
|
||||
|
||||
# Log final stats
|
||||
final_sev = Counter(train_dataset['severity'])
|
||||
final_src = Counter(train_dataset['source'])
|
||||
logger.info(f"Training severity: {dict(final_sev)}")
|
||||
logger.info(f"Training sources: {dict(final_src)}")
|
||||
|
||||
# GRPO Config — tuned for 0.5B on T4 (16GB VRAM)
|
||||
config = GRPOConfig(
|
||||
output_dir=OUTPUT_DIR,
|
||||
num_train_epochs=1, # 1 epoch over 15K samples = plenty
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4, # effective batch = 8
|
||||
num_generations=2,
|
||||
max_completion_length=768, # more room for detailed findings
|
||||
learning_rate=1e-6, # slightly higher lr for more data
|
||||
beta=0.04, # small KL penalty to prevent mode collapse with large dataset
|
||||
scale_rewards=True,
|
||||
reward_weights=[0.25, 0.25, 0.25, 0.25], # equal weight across 4 rewards
|
||||
gradient_checkpointing=True,
|
||||
bf16=True,
|
||||
logging_steps=10,
|
||||
logging_first_step=True,
|
||||
logging_strategy="steps",
|
||||
disable_tqdm=True,
|
||||
save_strategy="steps",
|
||||
save_steps=200,
|
||||
save_total_limit=2,
|
||||
push_to_hub=False,
|
||||
log_completions=False,
|
||||
report_to="none",
|
||||
seed=42,
|
||||
)
|
||||
|
||||
logger.info("Initializing GRPOTrainer with 4 reward functions...")
|
||||
trainer = GRPOTrainer(
|
||||
model=MODEL_NAME,
|
||||
args=config,
|
||||
reward_funcs=[format_reward, severity_reward, category_reward, quality_reward],
|
||||
train_dataset=train_dataset,
|
||||
)
|
||||
logger.info("GRPOTrainer initialized!")
|
||||
|
||||
logger.info("Starting training...")
|
||||
trainer.train()
|
||||
logger.info("Training complete!")
|
||||
|
||||
# Save
|
||||
logger.info(f"Saving model to {OUTPUT_DIR}...")
|
||||
trainer.save_model(OUTPUT_DIR)
|
||||
|
||||
# Push to Hub
|
||||
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)
|
||||
try:
|
||||
api.create_repo(repo_id=HUB_MODEL_ID, exist_ok=True)
|
||||
except Exception as e:
|
||||
logger.warning(f"create_repo: {e}")
|
||||
api.upload_folder(
|
||||
folder_path=OUTPUT_DIR,
|
||||
repo_id=HUB_MODEL_ID,
|
||||
commit_message="GRPO v2 — trained on 50K real audit findings, 4 reward functions",
|
||||
)
|
||||
logger.info(f"✅ Model pushed to https://huggingface.co/{HUB_MODEL_ID}")
|
||||
except Exception as e:
|
||||
logger.error(f"Push failed: {e}")
|
||||
else:
|
||||
logger.warning("No HF_TOKEN — model saved locally only")
|
||||
|
||||
logger.info("DONE")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
482
train_grpo_v2_colab.ipynb
Normal file
482
train_grpo_v2_colab.ipynb
Normal file
@@ -0,0 +1,482 @@
|
||||
{
|
||||
"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"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
3
training_args.bin
Normal file
3
training_args.bin
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:b51f3856815802830b7add9e23ddd089207e5c9941078dd606f120af0f983d09
|
||||
size 6776
|
||||
Reference in New Issue
Block a user