| 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") | |
| ``` | |