CoolFace
Modelpublic

mahdin70/graphcodebert-devign-code-vulnerability-detector

sourceHugging Facemitupdated 2y agoView on Hugging Face
2likes38downloads
README.md124 linesDownload Raw Back to root
1---2library_name: transformers3tags:4- Code5- Vulnerability6- Detection7- C/C++8datasets:9- DetectVul/devign10language:11- en12base_model:13- microsoft/graphcodebert-base14license: mit15metrics:16- accuracy17- precision18- f119- recall20---21 22## GraphCodeBERT for Code Vulnerability Detection23 24## Model Summary25This model is a fine-tuned version of **microsoft/graphcodebert-base**, optimized for detecting vulnerabilities in code. It is trained on the **DetectVul/devign** dataset.26The model takes in a code snippet and classifies it as either **safe (0)** or **vulnerable (1)**.27 28## Model Details29 30- **Developed by:** Mukit Mahdin31- **Finetuned from:** `microsoft/graphcodebert-base`32- **Language(s):** English (for code comments & metadata), C/C++33- **License:** MIT34- **Task:** Code vulnerability detection35- **Dataset Used:** `DetectVul/devign`36- **Architecture:** Transformer-based sequence classification37 38## Uses39 40### Direct Use41This model can be used for **static code analysis**, security audits, and automatic vulnerability detection in software repositories. It is useful for:42- **Developers**: To analyze their code for potential security flaws.43- **Security Teams**: To scan repositories for known vulnerabilities.44- **Researchers**: To study vulnerability detection in AI-powered systems.45 46### Downstream Use47This model can be integrated into **IDE plugins**, **CI/CD pipelines**, or **security scanners** to provide real-time vulnerability detection.48 49### Out-of-Scope Use50- The model is **not meant to replace human security experts**.51- It may not generalize well to **languages other than C/C++**.52- False positives/negatives may occur due to dataset limitations.53 54## Bias, Risks, and Limitations55- **False Positives & False Negatives:** The model may flag safe code as vulnerable or miss actual vulnerabilities.56- **Limited to C/C++:** The model was trained on a dataset primarily composed of **C and C++ code**. It may not perform well on other languages.57- **Dataset Bias:** The training data may not cover all possible vulnerabilities.58 59### Recommendations60Users should **not rely solely on the model** for security assessments. Instead, it should be used alongside **manual code review and static analysis tools**.61 62## How to Get Started with the Model63Use the code below to load the model and run inference on a sample code snippet:64 65```python66from transformers import AutoTokenizer, AutoModelForSequenceClassification67import torch68 69# Load the fine-tuned model70tokenizer = AutoTokenizer.from_pretrained("microsoft/graphcodebert-base")71model = AutoModelForSequenceClassification.from_pretrained("mahdin70/graphcodebert-devign-code-vulnerability-detector")72 73# Sample code snippet74code_snippet = '''75void process(char *input) {76    char buffer[50];77    strcpy(buffer, input); // Potential buffer overflow78}79'''80 81# Tokenize the input82inputs = tokenizer(code_snippet, return_tensors="pt", truncation=True, padding="max_length", max_length=512)83 84# Run inference85with torch.no_grad():86    outputs = model(**inputs)87    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)88    predicted_label = torch.argmax(predictions, dim=1).item()89 90# Output the result91print("Vulnerable Code" if predicted_label == 1 else "Safe Code")92```93 94## Training Details95 96### Training Data97- **Dataset:** `DetectVul/devign`98- **Classes:** `0 (Safe)`, `1 (Vulnerable)`99- **Size:** `21800` Code Snippets100 101### Training Procedure102- **Optimizer:** AdamW103- **Loss Function:** CrossEntropyLoss104- **Batch Size:** 16105- **Learning Rate:** 2e-05106- **Epochs:** 3107- **Hardware Used:** 2x T4 GPU108 109### Metrics110| Metric  | Score |111|------------|-------------|112| **Train Loss** | 0.6112 |113| **Evaluation Loss** | 0.605983 |114| **Accuracy** | 64.27% |115| **F1 Score** | 51.8% |116| **Precision** | 68.04% |117| **Recall** | 41.9% |118 119## Environmental Impact120 121| Factor  | Value |122|-----------|----------|123| **GPU Used** | 2x T4 GPU |124| **Training Time** | ~1 hour |