Arshit1616/vultrix-multilabel-v2

Vultrix multilabel-v2

Scans a Solidity smart contract and estimates, for 8 vulnerability types, how likely it is to have each one. This model powers the Vultrix scanner, a local web app built as a personal project.

Project code (web app, backend, training): github.com/Arshit-Verma/Vultrix

It is GraphCodeBERT fine-tuned on 35,533 real Ethereum contracts from the SmartBugs dataset (processed version).

Labels

#	Label	Meaning
0	Other	Other issues analysis tools report (uninitialised storage, locked ether, …)
1	access_control	Sensitive functions anyone can call
2	arithmetic	Integer overflow / underflow (mainly Solidity < 0.8)
3	denial_service	Patterns that can lock the contract
4	front_running	Results depend on transaction order
5	reentrancy	External call before the contract updates its state
6	time_manipulation	Logic based on block.timestamp
7	unchecked_low_calls	call / send results that are never checked

Each label is an independent yes/no (sigmoid). A contract can have several labels or none.

How to use it

Contracts are usually longer than the 512 tokens the model reads at once. Prepare them exactly as in training: unify line endings, remove comments, cut the tokens into 510-token windows, score every window, and keep the highest score per label. Then compare each label with its own threshold from thresholds.json.

import json, re, torch
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer, AutoModelForSequenceClassification

path = snapshot_download("Arshit1616/vultrix-multilabel-v2")
tok = AutoTokenizer.from_pretrained(path)
model = AutoModelForSequenceClassification.from_pretrained(path).eval()
thresholds = json.load(open(f"{path}/thresholds.json"))

def clean_code(code):
    code = code.replace("\r\n", "\n").replace("\r", "\n")
    code = re.sub(r"/\*.*?\*/", "", code, flags=re.S)
    code = re.sub(r"//[^\n]*", "", code)
    return re.sub(r"\n\s*\n+", "\n", code).strip()

code = open("Contract.sol", encoding="utf-8").read()
ids = tok.encode(clean_code(code), add_special_tokens=False)
windows = [[tok.bos_token_id] + ids[i:i + 510] + [tok.eos_token_id]
           for i in range(0, max(len(ids), 1), 510)]
with torch.no_grad():
    probs = torch.stack([torch.sigmoid(model(torch.tensor([w])).logits[0]) for w in windows]).max(0).values

for i, p in enumerate(probs.tolist()):
    name = model.config.id2label[i]
    print(f"{name:22} {p:.3f}", "FLAGGED" if p >= thresholds[name] else "")

Per-label thresholds (tuned on the validation set for best F1):

other	access_control	arithmetic	denial_service	front_running	reentrancy	time_manipulation	unchecked_low_calls
0.67	0.75	0.48	0.91	0.84	0.86	0.87	0.94

Results

On 4,459 test contracts the model never saw during training. The baseline always answers "arithmetic", the most common label.

Label	Contracts with it	Precision	Recall	F1	Baseline F1
Other	628	0.537	0.635	0.582	0.000
access_control	72	0.619	0.722	0.667	0.000
arithmetic	2,680	0.815	0.879	0.846	0.751
denial_service	282	0.710	0.780	0.743	0.000
front_running	117	0.518	0.487	0.502	0.000
reentrancy	545	0.841	0.754	0.795	0.000
time_manipulation	100	0.635	0.660	0.647	0.000
unchecked_low_calls	116	0.707	0.560	0.625	0.000
Macro-F1				0.676	0.094
Micro-F1		0.752	0.799	0.775	0.596

Training

• Data: Arshit1616/vultrix-smartbugs-processed, the training split minus 138 contracts that duplicate validation or test ones.
• Method: each contract's highest window score per label, trained with binary cross-entropy. Rare labels are weighted up (negatives ÷ positives, capped between 1 and 10). Up to 8 random windows per contract per step.
• Settings: AdamW, learning rate 2e-5, weight decay 0.01, 6% warmup, 3 epochs, 16 contracts per update, fp16, gradient checkpointing, seed 42. The best epoch was picked by validation macro-F1.
• Hardware: 2 × NVIDIA T4 (Kaggle).

Limits

• Tool labels: the labels mean "at least 2 of 9 automated analysis tools flagged this", not a human audit. The model learns the tools' blind spots and false alarms too.
• Old Solidity: the training contracts are mostly 2019-era 0.4.x. Modern 0.8 code is less familiar to the model, e.g. "arithmetic" is still predicted although 0.8 checks overflow by default.
• Weakest labels: front_running and unchecked_low_calls.
• No line numbers: it says what it suspects, not where.
• A first screen, not a security audit: don't treat a clean result as proof a contract is safe.

Credits and licence

• This model: Apache-2.0.
• Base model: microsoft/graphcodebert-base, from Microsoft's CodeBERT repository (MIT License, Copyright (c) Microsoft Corporation). Guo et al., GraphCodeBERT: Pre-training Code Representations with Data Flow, ICLR 2021.
• Training data: derived from SmartBugs Wild and SmartBugs results (Apache-2.0). Durieux et al., Empirical Review of Automated Analysis Tools on 47,587 Ethereum Smart Contracts, ICSE 2020.