tanaymitra01/graphcodebert-vulnerability-detector

Graph CodeBERT — Solidity Vulnerability Detector

Fine-tuned `microsoft/graphcodebert-base` for Solidity smart-contract vulnerability type classification.

Part of SolidityGuard — used as a first-pass detector alongside Slither and an LLM auditor.

How it fits in SolidityGuard

1. 1.Input: Solidity source
2. 2.Detectors: Slither / patterns + this Graph CodeBERT model
3. 3.LLM agents: Scanner → Analyzer → Exploit Gen / Fix Suggester
4. 4.Output: Audit findings with severity and confidence

Model

• —Tokenizer → max_length=512
• —Graph CodeBERT encoder (~125M params)
• —Linear classification head → 12-class softmax → label + confidence

Training

1. 1.SmartBugs-Wild contracts (capped at 5,000)
2. 2.Labels from SmartBugs-Results tool consensus (≥2 tools agree on a category; else safe)
3. 3.Split 70 / 15 / 15 (train / val / test) with light augmentation
4. 4.Fine-tune microsoft/graphcodebert-base with early stopping on validation macro-F1
5. 5.Best checkpoint published here

Intended use

• —Input: Solidity source code (string)
• —Output: one of 12 labels + confidence
• —Best as a screening signal, not a sole security audit

Labels

Held-out test metrics

Labels are noisy (static-analysis consensus), so scores are moderate by design.

Quick start

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

repo = "tanaymitra01/graphcodebert-vulnerability-detector"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForSequenceClassification.from_pretrained(repo)
model.eval()

code = """
pragma solidity ^0.8.0;
contract Vault {
    mapping(address => uint) public bal;
    function withdraw() public {
        uint amount = bal[msg.sender];
        (bool ok,) = msg.sender.call{value: amount}("");
        require(ok);
        bal[msg.sender] = 0;
    }
}
"""
inputs = tok(code, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
    probs = torch.softmax(model(**inputs).logits, dim=-1)[0]
pred = int(probs.argmax())
print(model.config.id2label[pred], float(probs[pred]))

With SolidityGuard

export GRAPHCODEBERT_PATH=tanaymitra01/graphcodebert-vulnerability-detector

Where to get the weights

Files

• —model.safetensors — weights
• —config.json — RobertaForSequenceClassification config + label maps
• —label_map.json — label list / id maps used in training
• —README.md — this model card

Limitations

• —Tool-derived labels ≠ audited ground truth
• —Truncation at 512 tokens; large contracts lose context
• —Rare classes (e.g. access control) have low F1
• —Not a replacement for professional security review

Citation

@misc{solidityguard-graphcodebert,
  title  = {Graph CodeBERT Vulnerability Detector for Solidity},
  author = {Tanay Mitra},
  year   = {2026},
  url    = {https://huggingface.co/tanaymitra01/graphcodebert-vulnerability-detector}
}