ayshajavd/graphcodebert-vuln-classifier

GraphCodeBERT Vulnerability Classifier

A multi-label code vulnerability detection model that identifies 31 vulnerability classes (30 CWEs + safe) mapped to the OWASP Top 10 2021 categories. Fine-tuned from CodeBERTa-small-v1 on 175K+ labeled code samples.

Quick Start

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_id = "ayshajavd/graphcodebert-vuln-classifier"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id)
model.eval()

code = """
import sqlite3
def get_user(username):
    query = f"SELECT * FROM users WHERE username = '{username}'"
    conn = sqlite3.connect('db.sqlite')
    return conn.execute(query).fetchone()
"""

inputs = tokenizer(code, return_tensors="pt", max_length=512, truncation=True, padding=True)
with torch.no_grad():
    logits = model(**inputs).logits
    probs = torch.sigmoid(logits).squeeze()

# Get predictions above threshold
TARGET_CWES = ["safe", "CWE-20", "CWE-22", "CWE-78", "CWE-79", "CWE-89", "CWE-94",
    "CWE-119", "CWE-125", "CWE-190", "CWE-200", "CWE-264", "CWE-269", "CWE-276",
    "CWE-284", "CWE-287", "CWE-310", "CWE-327", "CWE-330", "CWE-352", "CWE-362",
    "CWE-399", "CWE-401", "CWE-416", "CWE-434", "CWE-476", "CWE-502", "CWE-601",
    "CWE-787", "CWE-798", "CWE-918"]

threshold = 0.5
for i, (cwe, prob) in enumerate(zip(TARGET_CWES, probs)):
    if prob > threshold:
        print(f"{cwe}: {prob:.3f}")

Model Details

Supported Languages

Python, JavaScript, Java, C, C++, PHP, Go

The model was trained on a diverse multi-language dataset. Performance is strongest on C/C++ (largest training subset from BigVul) and Python/JavaScript (from the multi-language datasets).

Evaluation Results (Test Set — 5,000 samples)

Threshold Comparison

Per-Class Performance (threshold=0.3)

OWASP A01:2021 — Broken Access Control

OWASP A02:2021 — Cryptographic Failures

OWASP A03:2021 — Injection

OWASP A04:2021 — Insecure Design

OWASP A07–A10

Key Metric: Safe Code Detection

Model Strengths

• —Excellent recall on many vulnerability classes (0.75–1.0 for SQL injection, buffer overflow, XSS, code injection, etc.)
• —Strong safe code detection (F1=0.95) — reliably identifies secure code
• —High sensitivity — at threshold 0.3, catches most real vulnerabilities (macro recall=0.50)

Model Limitations

• —Low precision on rare classes — many false positives, especially on CWEs with few training examples
• —Precision can be improved by using threshold=0.5 (macro F1 improves to 0.125 but recall drops)
• —Classes with 0 test support cannot be evaluated

Design choice: For security applications, we prioritize recall (catching real vulnerabilities) over precision (reducing false positives). Missing a real vulnerability (false negative) is worse than flagging safe code (false positive).

Training Data

The model was trained on the code-security-vulnerability-dataset (175,419 samples), combining:

1. 1.[BigVul](https://huggingface.co/datasets/bstee615/bigvul) — 265K C/C++ vulnerable functions from real CVEs
2. 2.[CWE-enriched BigVul/PrimeVul](https://huggingface.co/datasets/mahdin70/cwe_enriched_balanced_bigvul_primevul) — Balanced CWE-labeled subset
3. 3.[Code Vulnerability Labeled](https://huggingface.co/datasets/lemon42-ai/Code_Vulnerability_Labeled_Dataset) — Multi-language (Python, JS, Java, PHP, Go)
4. 4.[CyberNative DPO](https://huggingface.co/datasets/CyberNative/Code_Vulnerability_Security_DPO) — Vulnerable/secure code pairs

Training Configuration

Limitations

1. 1.Class imbalance: Many rare CWE types have very few training examples, leading to high false positive rates
2. 2.Sequence length: Limited to 512 tokens — long functions may be truncated
3. 3.Language bias: Strongest on C/C++ due to BigVul's dominance. Go and PHP performance may be lower
4. 4.Single-function analysis: Analyzes individual functions, not cross-function or cross-file vulnerabilities
5. 5.Not a replacement: Should complement manual review and established SAST tools (Semgrep, CodeQL, etc.)

Interactive Demo

Try the model in our Code Security Analyzer Space — paste any code and get a full security report with OWASP mapping, severity scores, attack chain analysis, and suggested fixes.

Citation

@misc{graphcodebert-vuln-classifier,
  title={GraphCodeBERT Vulnerability Classifier: Multi-label CWE Detection Mapped to OWASP Top 10},
  author={ayshajavd},
  year={2025},
  url={https://huggingface.co/ayshajavd/graphcodebert-vuln-classifier}
}