UnixCoder-Primevul-BigVul

UnixCoder-Primevul-BigVul Model Card

Model Overview

UnixCoder-Primevul-BigVul is a multi-task model based on Microsoft's unixcoder-base, fine-tuned to detect vulnerabilities (vul) and classify Common Weakness Enumeration (CWE) types in code snippets. It was developed by 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/UnixCoder-Primevul-BigVul
• Base Model: microsoft/unixcoder-base
• Tasks: Vulnerability Detection (Binary), CWE Classification (Multi-class)
• License: MIT (assumed; adjust if different)
• Date: Trained and uploaded as of March 11, 2025

Model Architecture

The model extends unixcoder-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 MultiTaskUnixCoder 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:

• Output Directory: ./unixcoder_multitask
• 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
• WANDB: Disabled

Training Environment

• Hardware: NVIDIA Tesla T4 GPU
• Framework: PyTorch 2.5.1+cu121, Transformers 4.47.0
• Duration: ~6 hours, 34 minutes, 53 seconds (23,397 steps)

Training Metrics

Validation metrics across epochs:

• Final Training Loss: 0.4430 (average over all steps)

Evaluation

The model was evaluated on the test split (26,738 samples) with the following metrics:

• Vulnerability Detection:
  • Accuracy: 0.9571
  • Precision: 0.7947
  • Recall: 0.5437
  • F1 Score: 0.6457
• CWE Classification (on vulnerable samples):
  • Accuracy: 0.3288

The model excels at identifying non-vulnerable code (high accuracy) but has moderate recall for vulnerabilities and lower CWE classification accuracy, indicating room for improvement in CWE prediction.

Usage

Installation

Install the required libraries:

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:

from transformers import AutoTokenizer, AutoModel
import torch

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("microsoft/unixcoder-base")
model = AutoModel.from_pretrained("mahdin70/UnixCoder-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:

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 struggles with precise CWE classification (32.88% accuracy), likely due to class imbalance or complexity in distinguishing similar CWE types.
• Recall: Moderate recall (54.37%) 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 for better CWE accuracy.
• Experiment with class weighting to address CWE imbalance.
• Fine-tune on additional datasets for broader generalization.