File size: 841 Bytes
399dbf7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
---
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](https://huggingface.co/microsoft/codebert-base) for vulnerability detection.

## Usage
```python
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")
```