vikramdgx's picture
Upload README.md with huggingface_hub
04ee056 verified
|
Raw
History Blame Contribute Delete
16.1 kB
---
library_name: peft
license: apache-2.0
base_model: Qwen/Qwen2.5-Coder-7B-Instruct
tags:
- security
- vulnerability-detection
- injection
- code-analysis
- lora
- qwen2
- cybersecurity
- sast
- cwe
datasets:
- hitoshura25/cvefixes
- claudios/DiverseVul
- CyberNative/Code_Vulnerability_Security_DPO
language:
- en
- code
pipeline_tag: text-generation
model-index:
- name: injection-vulnerability-detector
results:
- task:
type: text-generation
name: Injection Vulnerability Detection
dataset:
name: Injection Vulnerability Novel Test Set (leakage-controlled)
type: injection-vuln-novel
metrics:
- name: F1
type: f1
value: 0.6803519061583577
verified: false
- name: Precision
type: precision
value: 0.8169014084507042
verified: false
- name: Recall
type: recall
value: 0.5829145728643216
verified: false
---
# Injection Vulnerability Detector v1.0
A fine-tuned **Qwen2.5-Coder-7B-Instruct** model specialized in detecting injection
vulnerabilities in source code. Built with LoRA (Low-Rank Adaptation) for efficient
fine-tuning on consumer/workstation hardware.
**Key contribution:** Existing vulnerability datasets lack CWE-specific labels for
injection subtypes β€” most group all injections generically, and rare types like LDAP,
XPath, EL injection, and SSTI have near-zero representation. We solved this with a
**template-based synthetic data generator** (151 templates, 9 CWEs) that produces
CWE-specific injection samples at scale without requiring an LLM, combined with
filtered real-world vulnerability data from open-source datasets.
## What it does
Given a code snippet, the model identifies whether it contains an **injection vulnerability**
and, if so, classifies the specific CWE (Common Weakness Enumeration) type. It outputs
structured JSON with a verdict, CWE ID, vulnerability type, tainted data-flow analysis,
an explanation, and a fix suggestion.
### Supported vulnerability types (9 CWEs)
| CWE | Vulnerability Type |
|-----|-------------------|
| CWE-89 | SQL Injection |
| CWE-78 | OS Command Injection |
| CWE-94 | Code Injection |
| CWE-611 | XXE (XML External Entity) |
| CWE-90 | LDAP Injection |
| CWE-643 | XPath Injection |
| CWE-917 | Expression Language Injection |
| CWE-1336 | Server-Side Template Injection (SSTI) |
| CWE-113 | CRLF Injection |
## Evaluation results
**Our primary metrics are computed on novel code only** (template clones removed; the full-set figure is shown for comparison) β€” test samples verified
to be structurally distinct from all training data. Of 1572 total
test samples, 1007 (64.1%) were identified as structural
clones of training data and excluded (see Evaluation Methodology below).
### Evaluation: three levels of rigor
| Evaluation set | Precision | Recall | F1 |
|---|---|---|---|
| Full test set (includes template clones) | 96.2% | 83.6% | **89.5%** |
| Novel code, all injection-related CWEs (n=565) | 82.8% | 46.6% | 59.7% |
| **Novel code, in-scope 9 target CWEs (n=496)** | **81.7%** | **58.3%** | **68.0%** |
Our **primary metric is the in-scope novel row** β€” performance on the 9 CWEs the model targets, on code it has never seen in any structural form. The full-set figure is what most detectors report; we consider it inflated by template-clone memorisation (see Evaluation Methodology). The all-CWE novel row additionally counts out-of-scope vulnerability types the model does not target.
### Primary metrics (novel code, in-scope 9 CWEs, n=496)
| Metric | Score |
|--------|-------|
| **Precision** | 81.7% |
| **Recall** | 58.3% |
| **F1 Score** | 68.0% |
| **Accuracy** | 78.0% |
| JSON Parse Rate | 94.9% |
### Per-CWE performance (novel code)
| CWE | Vulnerability | Precision | Recall | F1 | Samples |
|-----|--------------|-----------|--------|-----|---------|
| CWE-89 | SQL Injection | 76.5% | 67.2% | 71.6% | 58 |
| CWE-78 | OS Command Injection | 75.9% | 42.3% | 54.3% | 52 |
| CWE-94 | Code Injection | 77.4% | 46.2% | 57.8% | 52 |
| CWE-611 | XXE (XML External Entity) | 100.0% | 75.0% | 85.7% | 16 |
| CWE-90 | LDAP Injection | 100.0% | 60.0% | 75.0% | 5 |
| CWE-643 | XPath Injection | 100.0% | 100.0% | 100.0% | 4 |
| CWE-917 | Expression Language Injection | 100.0% | 100.0% | 100.0% | 2 |
| CWE-1336 | Server-Side Template Injection | 100.0% | 100.0% | 100.0% | 6 |
| CWE-113 | CRLF Injection | 100.0% | 100.0% | 100.0% | 4 |
### High-confidence mode (zero false positives)
The `tainted_flow` field doubles as a confidence signal. Accepting a `VULNERABLE`
verdict **only when the model names a concrete source and sink** filters out its
weakest calls:
| Mode | Precision | Recall | F1 |
|---|---|---|---|
| Standard (all VULNERABLE verdicts) | 81.7% | 58.3% | 68.0% |
| High-confidence (concrete `tainted_flow` required) | **100.0%** | 22.1% | 36.2% |
On the in-scope novel test set this yielded **44 true positives and 0 false
positives** (n=496). Use high-confidence mode where alert fatigue matters more than
coverage (CI gating, auto-filing issues); use standard mode for triage sweeps where
a human reviews each finding.
```python
import json
r = json.loads(response)
flow = r.get("tainted_flow") or {}
high_confidence = r.get("verdict") == "VULNERABLE" and flow.get("source") and flow.get("sink")
```
### Evaluation methodology
Standard group-aware splitting (by CVE ID / project) prevents the *same* code
from appearing in both train and test, but it does not prevent **template
clones** β€” structurally identical code differing only in variable names,
string literals, and numeric constants β€” from leaking across the split. A
skeleton-hashing analysis revealed that 64.1% of the initial test set
were structural clones of training samples, scoring near-perfect F1 from
memorisation rather than generalisation.
**Skeleton hashing method:** strip comments, string literals, identifiers, and
numeric literals from each code sample, then SHA-256 hash the normalised
skeleton. Any test sample whose skeleton matches a training sample is
classified as a clone and excluded from the metrics reported above.
The scores above therefore reflect performance on genuinely novel code that
the model has never seen in any structural form during training. We report
these numbers β€” not the inflated full-test-set figures β€” because they are
what matters for real-world deployment.
## Usage
### Quick start with transformers + PEFT
```python
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load base model + LoRA adapter
base_model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen2.5-Coder-7B-Instruct",
torch_dtype="auto",
device_map="auto",
)
model = PeftModel.from_pretrained(base_model, "vikramdgx/injection-vulnerability-detector")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B-Instruct")
# Analyze code β€” use the EXACT prompt format the model was trained on
code_snippet = '''import sqlite3
def get_user(username):
conn = sqlite3.connect("app.db")
query = f"SELECT * FROM users WHERE name = '{username}'"
return conn.execute(query).fetchone()
'''
SYSTEM_PROMPT = (
"You are a code security analyzer specialized in detecting injection "
"vulnerabilities. Analyze the provided code and respond with a JSON "
"object containing: verdict, cwe, vulnerability_type, tainted_flow, "
"explanation, and fix_suggestion."
)
user_message = (
"Analyze the following code for injection vulnerabilities. "
"Respond with JSON.\n\n"
+ code_snippet
)
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user_message},
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=768, do_sample=False)
response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)
```
### Expected output format
The model returns structured JSON matching its training schema:
```json
{
"verdict": "VULNERABLE",
"cwe": "CWE-89",
"vulnerability_type": "SQL Injection",
"tainted_flow": {
"source": "username parameter",
"sink": "conn.execute(query)",
"sanitizer": null
},
"explanation": "User input is directly interpolated into SQL query string via f-string without parameterization, enabling SQL injection.",
"fix_suggestion": "Use parameterized queries: conn.execute('SELECT * FROM users WHERE name = ?', (username,))"
}
```
For safe code, the model returns:
```json
{
"verdict": "SAFE",
"cwe": null,
"vulnerability_type": null,
"tainted_flow": null,
"explanation": "No injection vulnerability is present; untrusted input is not passed unsanitized to a sensitive sink.",
"fix_suggestion": null
}
```
## Training details
### Hardware
- **NVIDIA DGX Spark** (Grace Blackwell GB10)
- 128 GB unified LPDDR5X memory
- CUDA 13.0, compute capability sm_121
### Configuration
| Parameter | Value |
|-----------|-------|
| Base model | Qwen/Qwen2.5-Coder-7B-Instruct |
| Base model checkpoint used | `unsloth/Qwen2.5-Coder-7B-Instruct` (Unsloth mirror of the same Apache-2.0 weights; the adapter loads against either) |
| Method | LoRA (PEFT) |
| LoRA rank (r) | 16 |
| LoRA alpha | 32 |
| Target modules | q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj |
| Trainable parameters | 40.3M (0.53% of total) |
| Epochs | 3 |
| Batch size | 2 (gradient accumulation: 4, effective 8) |
| Learning rate | 2e-4 |
| Max sequence length | 4096 |
| Optimizer | AdamW (8-bit) |
| LR scheduler | Cosine |
| Packing | Enabled |
| Loss | Completion-only |
### Why we built a custom dataset
Existing open-source vulnerability datasets (CVEfixes, DiverseVul, BigVul, etc.)
provide valuable real-world code samples, but they share a common gap: **they lack
fine-grained, CWE-specific labels for injection subtypes.** Most label code broadly
as "vulnerable" or group all injections under generic categories like CWE-74 or
CWE-20 β€” without distinguishing SQL injection (CWE-89) from command injection
(CWE-78) from SSTI (CWE-1336). For niche injection types like LDAP injection
(CWE-90), XPath injection (CWE-643), or EL injection (CWE-917), labeled samples
in public datasets are extremely scarce (often single digits).
This makes it impossible to train a detector that both identifies vulnerabilities
AND classifies the specific CWE type β€” which is what security engineers actually
need for triage.
### Our approach: template-based synthetic generation
To solve this, we built a **procedural code generator** that produces CWE-specific
injection samples at scale β€” no LLM in the loop, pure template expansion:
- **151 handcrafted code templates** covering all 9 target CWEs, each with
`@@PLACEHOLDER@@` tokens for variable names, function names, table names,
database fields, and code patterns
- **Randomized variable pools** (200+ variable names, 100+ function names, 50+
table names per CWE) ensure each generated sample is syntactically unique
- **Paired generation**: every vulnerable template has a corresponding safe version
that uses parameterized queries, input validation, or proper escaping β€” teaching
the model the difference, not just the pattern
- **~5,000 synthetic samples** generated in seconds, balanced across all 9 CWEs
including the rare ones that have near-zero representation in public datasets
This approach is deterministic, reproducible, and produces exactly the CWE
distribution the model needs β€” no data collection bottleneck, no labeling errors,
no class imbalance.
### Dataset composition
The final training set combines three sources:
1. **Real-world vulnerabilities** β€” filtered from open-source CVE/vulnerability
datasets, keeping only injection-related CWEs. Provides realistic code patterns
from production software.
2. **Template-based synthetic data** (~5,000 samples) β€” our original procedural
generator. Fills the CWE-specific gap that public datasets leave open,
especially for rare injection types (LDAP, XPath, EL, SSTI, Header Injection).
3. **Hard negatives** β€” safe code samples including post-patch fixes and
non-vulnerable functions from vulnerability-adjacent codebases. Teaches the
model what secure code looks like.
Group-aware splitting ensures no data leakage between train and test sets (samples
sharing a CVE ID or project stay together). Additionally, a skeleton-hashing pass
excludes structural template clones from the evaluation set (see Evaluation
Methodology above).
## Dataset sources and attribution
This model was trained on data from the following open-source datasets, combined
with our original synthetic generation. We gratefully acknowledge the dataset creators:
| Dataset | Source | License | Role |
|---------|--------|---------|------|
| [CVEfixes](https://huggingface.co/datasets/hitoshura25/cvefixes) | Bhandari et al. | Apache 2.0 (data: CC BY 4.0) | Real-world vuln + patch pairs |
| [DiverseVul](https://huggingface.co/datasets/claudios/DiverseVul) | Chen & Bhatt (RAID 2023) | Not specified on HF card | Safe code (non-vuln functions) |
| [Code Vulnerability Security DPO](https://huggingface.co/datasets/CyberNative/Code_Vulnerability_Security_DPO) | CyberNative AI | Apache 2.0 | Injection code examples |
| **Template-based synthetic generator** | **Original work** | **Apache 2.0** | **CWE-specific injection samples (core contribution)** |
**License notes:**
- The DiverseVul dataset does not declare an explicit license on its HuggingFace
card as of this writing. This model's own adapter
weights are original work released under Apache 2.0, but users should check the
current licensing status of these upstream datasets before commercial deployment.
## Limitations
- **Injection-only scope:** This model detects 9 injection-related CWEs. It does not
cover other vulnerability classes (buffer overflow, authentication, crypto, etc.).
- **Synthetic training bias:** The model is trained partly on template-generated code.
While this solves the CWE-distribution problem, performance on novel real-world
patterns is lower than on template-similar code, reflecting the generalisation gap
that template-based training introduces. See Evaluation Results for exact numbers.
- **Rare CWEs are synthetic-validated only:** The five rare injection types (LDAP,
XPath, EL, SSTI, Header) have near-zero representation in public vulnerability
datasets. Our evaluation on novel code therefore cannot validate these CWEs β€”
their coverage relies entirely on synthetic templates. Real-world performance on
these types is unknown.
- **CWE classification accuracy:** While detection (vulnerable vs. safe) is reliable,
the specific CWE label assigned to a detected vulnerability may be incorrect in some
cases β€” particularly between similar injection types (e.g. CWE-94 code injection vs CWE-1336 template injection, or CWE-78 vs CWE-94).
- **Code context:** The model analyzes individual functions/snippets. It cannot trace
data flow across files or understand application-level sanitization.
- **Language coverage:** Primarily trained on Python, Java, PHP, C/C++, and JavaScript.
Performance on other languages may vary.
- **Not a replacement for manual review:** Use as a triage/prioritization tool alongside
established SAST tooling and expert code review.
## License
This model adapter is released under the **Apache 2.0** license. The base model
(Qwen2.5-Coder-7B-Instruct) is also Apache 2.0.
## Citation
If you use this model in your research, please cite:
```bibtex
@misc{injection-detector-v10,
title={Injection Vulnerability Detector v1.0},
author={Thrivikram Gujarathi},
year={2026},
publisher={HuggingFace},
url={https://huggingface.co/vikramdgx/injection-vulnerability-detector}
}
```