|
|
--- |
|
|
language: |
|
|
- en |
|
|
- code |
|
|
tags: |
|
|
- security |
|
|
- vulnerability-detection |
|
|
- codebert |
|
|
- classification |
|
|
license: mit |
|
|
--- |
|
|
|
|
|
# codebert_vulnerability_scanner |
|
|
|
|
|
## Overview |
|
|
|
|
|
`codebert_vulnerability_scanner` is a fine-tuned RoBERTa model (specifically based on Microsoft's CodeBERT) designed to detect potential security vulnerabilities in source code snippets. It treats vulnerability detection as a binary classification task, labeling code as either `SAFE` or `VULNERABLE`. |
|
|
|
|
|
## Model Architecture |
|
|
|
|
|
This model utilizes the `RobertaForSequenceClassification` architecture. It was pre-trained on the CodeSearchNet dataset (a large collection of function-level code across multiple programming languages) and subsequently fine-tuned on a curated dataset of C and C++ functions labeled with Common Weakness Enumerations (CWEs), such as buffer overflows and memory leaks. |
|
|
|
|
|
- **Base Model:** `microsoft/codebert-base` |
|
|
- **Head:** A linear classification head on top of the pooled output. |
|
|
- **Input:** Source code functions (tokenized). |
|
|
- **Output:** Logits for two classes: SAFE (0) and VULNERABLE (1). |
|
|
|
|
|
## Intended Use |
|
|
|
|
|
This model is intended primarily for DevSecOps workflows and static analysis research. |
|
|
|
|
|
- **Automated Code Review:** Scanning pull requests for high-risk code patterns before merging. |
|
|
- **Security Auditing:** Quickly analyzing large legacy codebases to prioritize manual security reviews. |
|
|
- **Research:** Benchmarking against traditional static analysis security testing (SAST) tools. |
|
|
|
|
|
### How to use |
|
|
|
|
|
```python |
|
|
from transformers import AutoTokenizer, AutoModelForSequenceClassification |
|
|
import torch |
|
|
|
|
|
model_name = "your_username/codebert_vulnerability_scanner" |
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
model = AutoModelForSequenceClassification.from_pretrained(model_name) |
|
|
|
|
|
# Example C function snippet |
|
|
code_snippet = """ |
|
|
void vulnerable_function(char *user_input) { |
|
|
char buffer[64]; |
|
|
strcpy(buffer, user_input); // Potential buffer overflow |
|
|
} |
|
|
""" |
|
|
|
|
|
inputs = tokenizer(code_snippet, return_tensors="pt", truncation=True, max_length=512) |
|
|
|
|
|
with torch.no_grad(): |
|
|
logits = model(**inputs).logits |
|
|
|
|
|
predicted_class_id = logits.argmax().item() |
|
|
labels = model.config.id2label |
|
|
print(f"Prediction: {labels[predicted_class_id]}") |
|
|
# Expected output: VULNERABLE |
