README.md

---
license: mit
datasets:
- mahdin70/balanced_merged_bigvul_primevul
metrics:
- accuracy
- f1
- recall
- precision
base_model:
- microsoft/codebert-base
pipeline_tag: text-classification
library_name: transformers
---

# CodeBERT-Primevul-BigVul Model Card

## Model Overview

`CodeBERT-Primevul-BigVul` is a multi-task model based on Microsoft's `codebert-base`, fine-tuned to detect vulnerabilities (`vul`) and classify Common Weakness Enumeration (CWE) types in code snippets. It was developed by [mahdin70](https://huggingface.co/mahdin70) and trained on a balanced dataset combining BigVul and PrimeVul datasets. The model performs binary classification for vulnerability detection and multi-class classification for CWE identification.

- **Model Repository**: [mahdin70/CodeBERT-Primevul-BigVul](https://huggingface.co/mahdin70/CodeBERT-Primevul-BigVul)
- **Base Model**: [microsoft/codebert-base](https://huggingface.co/microsoft/codebert-base)
- **Tasks**: Vulnerability Detection (Binary), CWE Classification (Multi-class)
- **License**: MIT (assumed; adjust if different)
- **Date**: Trained and uploaded as of April 22, 2025

## Model Architecture

The model extends `codebert-base` with two task-specific heads:
- **Vulnerability Head**: A linear layer mapping 768-dimensional hidden states to 2 classes (vulnerable or not).
- **CWE Head**: A linear layer mapping 768-dimensional hidden states to 135 classes (134 CWE types + 1 for "no CWE").

The architecture is implemented as a custom `MultiTaskCodeBERT` class in PyTorch, with the loss computed as the sum of cross-entropy losses for both tasks.

## Training Dataset

The model was trained on the `mahdin70/balanced_merged_bigvul_primevul` dataset, which combines:
- **BigVul**: A dataset of real-world vulnerabilities from open-source projects.
- **PrimeVul**: A dataset focused on prime vulnerabilities in code.

### Dataset Details
- **Splits**:
  - Train: 124,780 samples
  - Validation: 26,740 samples
  - Test: 26,738 samples
  
- **Features**:
  - `func`: Code snippet (text)
  - `vul`: Binary label (0 = non-vulnerable, 1 = vulnerable)
  - `CWE ID`: CWE identifier (e.g., CWE-89) or None for non-vulnerable samples
    
- **Preprocessing**:
  - CWE labels were encoded using a `LabelEncoder` with 134 unique CWE classes identified across the dataset.
  - Non-vulnerable samples assigned a CWE label of -1 (mapped to 0 in the model).

The dataset is balanced to ensure a fair representation of vulnerable and non-vulnerable samples, with a maximum of 10 samples per commit where applicable.

## Training Details

### Training Arguments
The model was trained using the Hugging Face `Trainer` API with the following arguments:
- **Evaluation Strategy**: Per epoch
- **Save Strategy**: Per epoch
- **Learning Rate**: 2e-5
- **Batch Size**: 8 (per device, train and eval)
- **Epochs**: 3
- **Weight Decay**: 0.01
- **Logging**: Every 10 steps, logged to `./logs`

### Training Environment
- **Hardware**: 2x NVIDIA Tesla T4 GPU
- **Framework**: PyTorch 2.5.1+cu121, Transformers 4.47.0
- **Duration**: ~6 hours, 23 minutes, 18 seconds (23,397 steps)

### Training Metrics
Validation metrics across epochs:

| Epoch | Training Loss | Validation Loss | Vul Accuracy | Vul Precision | Vul Recall | Vul F1   | CWE Accuracy |
|-------|---------------|-----------------|--------------|---------------|------------|----------|--------------|
| 1     | 0.4275        | 0.5737          | 0.9519       | 0.7753        | 0.4795     | 0.5925   | 0.0656       |
| 2     | 0.7608        | 0.5450          | 0.9537       | 0.7766        | 0.5133     | 0.6181   | 0.1349       |
| 3     | 0.5624        | 0.5443          | 0.9545       | 0.7669        | 0.5400     | 0.6338   | 0.1749       |


## Usage

### Installation
Install the required libraries:
```bash
pip install transformers torch datasets huggingface_hub
```

### Sample Code Snippet
Below is an example of how to use the model for inference on a code snippet:

```python
from transformers import AutoTokenizer, AutoModel
import torch

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
model = AutoModel.from_pretrained("mahdin70/CodeBERT-Primevul-BigVul", trust_remote_code=True)
model.eval()

# Example code snippet
code = """
bool DebuggerFunction::InitTabContents() {
Value* debuggee;
EXTENSION_FUNCTION_VALIDATE(args_->Get(0, &debuggee));

DictionaryValue* dict = static_cast<DictionaryValue*>(debuggee);
EXTENSION_FUNCTION_VALIDATE(dict->GetInteger(keys::kTabIdKey, &tab_id_));

contents_ = NULL;
TabContentsWrapper* wrapper = NULL;
bool result = ExtensionTabUtil::GetTabById(
tab_id_, profile(), include_incognito(), NULL, NULL, &wrapper, NULL);
if (!result || !wrapper) {
error_ = ExtensionErrorUtils::FormatErrorMessage(
keys::kNoTabError,
base::IntToString(tab_id_));
return false;
}
contents_ = wrapper->web_contents();

if (ChromeWebUIControllerFactory::GetInstance()->HasWebUIScheme(
contents_->GetURL())) {
error_ = ExtensionErrorUtils::FormatErrorMessage(
keys::kAttachToWebUIError,
contents_->GetURL().scheme());
return false;
}

return true;
}
"""

# Tokenize input
inputs = tokenizer(code, return_tensors="pt", padding="max_length", truncation=True, max_length=512)

# Move to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
inputs = {k: v.to(device) for k, v in inputs.items()}

# Get predictions
with torch.no_grad():
    outputs = model(**inputs)
    vul_logits = outputs["vul_logits"]
    cwe_logits = outputs["cwe_logits"]

    # Vulnerability prediction
    vul_pred = torch.argmax(vul_logits, dim=1).item()
    print(f"Vulnerability: {'Vulnerable' if vul_pred == 1 else 'Not Vulnerable'}")

    # CWE prediction (if vulnerable)
    if vul_pred == 1:
        cwe_pred = torch.argmax(cwe_logits, dim=1).item() - 1  # Subtract 1 as -1 is "no CWE"
        print(f"Predicted CWE: {cwe_pred if cwe_pred >= 0 else 'None'}")
```

### Output Example:
```bash
Vulnerability: Vulnerable
Predicted CWE: 120  # Maps to CWE-120 (Buffer Overflow), depending on encoder
```

## Notes
- The CWE prediction is an integer index (0 to 133). To map it to a specific CWE ID (e.g., CWE-120), you need the LabelEncoder used during training, available in the dataset preprocessing step.
- Ensure `trust_remote_code=True` as the model uses custom code from the repository.

## Limitations
- **CWE Accuracy**: The model has low CWE classification accuracy (17.49%), likely due to class imbalance or complexity in distinguishing similar CWE types.
- **Recall**: Moderate recall (54.00%) for vulnerability detection suggests some vulnerable samples may be missed.
- **Generalization**: Trained on BigVul and PrimeVul, performance may vary on out-of-domain codebases.

## Future Improvements
- Increase training epochs or dataset size to improve CWE accuracy.
- Experiment with class weighting to address CWE imbalance.
- Fine-tune on additional datasets for broader generalization.