metadata
language: en
license: mit
tags:
- code
- vulnerability-detection
- defect-detection
- codebert
datasets:
- devign
base_model: microsoft/codebert-base
CodeBERT for Defect Detection
Fine-tuned from microsoft/codebert-base for vulnerability detection.
Usage
from transformers import RobertaTokenizer, RobertaForSequenceClassification
import torch
tokenizer = RobertaTokenizer.from_pretrained("phamtungthuy/codebert-defect-detection")
model = RobertaForSequenceClassification.from_pretrained("phamtungthuy/codebert-defect-detection")
code = "int main() { char buf[10]; strcpy(buf, argv[1]); }"
inputs = tokenizer(code, return_tensors="pt", truncation=True, max_length=400)
prob = torch.sigmoid(model(**inputs).logits)
print("Vulnerable" if prob > 0.5 else "Secure")