roberta-cwe-classifier-kelemia-v0.2

Kelemia v0.2 for CWE Classification

This model is a fine-tuned version of RoBERTa for classifying Common Weakness Enumeration (CWE) vulnerabilities.

The changes are significant, particularly the work to increase the number of entries in the dataset and the optimisation of hyperparameters.

Model Description

• Model Type: RoBERTa
• Language(s): English
• License: MIT
• Finetuned from Model: roberta-base

Intended Uses & Limitations

This model is intended for classifying software vulnerabilities according to the CWE standard. It should be used as part of a broader security analysis process and not as a standalone solution for identifying vulnerabilities.

Training and Evaluation Data

Dunateo/VulnDesc_CWE_Mapping

Example Usage

Here's an example of how to use this model for inference:

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load model and tokenizer
model_name = "Dunateo/roberta-cwe-classifier-kelemia-v0.2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Prepare input text
text = "The application stores sensitive user data in plaintext."

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

# Perform inference
with torch.no_grad():
    outputs = model(**inputs)

# Get prediction
probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
predicted_class = torch.argmax(probabilities, dim=-1).item()

print(f"Predicted CWE class: {predicted_class}")
print(f"Confidence: {probabilities[predicted_class].item():.4f}")

Label Dictionary

This model uses the following mapping for CWE classes:

{
  "0": "CWE-79",
  "1": "CWE-89",
  ...
}

import json
from huggingface_hub import hf_hub_download

label_dict_file = hf_hub_download(repo_id="Dunateo/roberta-cwe-classifier-kelemia-v0.2", filename="label_dict.json")
with open(label_dict_file, 'r') as f:
    label_dict = json.load(f)

id2label = {v: k for k, v in label_dict.items()}

print(f"Label : {id2label[predicted_class]}")

Now you can use label_dict to map prediction indices to CWE classes

Training procedure

Training hyperparameters

• Number of epochs: 3
• Learning rate: Scheduled from 1e-06 to 4e-5
• Batch size: 64
• Weight decay: 0.005
• Gradient Accumulation Steps: 1

Training results

• Training Loss: 0.834400 (final)
• Validation Loss: 0.942376 (final)
• Training Time: 2388.45 seconds (approximately 40 minutes)

Loss progression

Evaluation results

The model shows consistent improvement over the training period:

• Initial Training Loss: 1.246200
• Final Training Loss: 0.834400
• Initial Validation Loss: 1.183215
• Final Validation Loss: 0.942376

Performance analysis

• The model demonstrates a steady decrease in both training and validation loss, indicating good learning progress.
• The final validation loss (0.942376) being slightly higher than the final training loss (0.834400) suggests that the model generalizes well to unseen data, with a slight risk of overfitting.

Ethical considerations

This model should be used responsibly as part of a comprehensive security strategy. It should not be relied upon as the sole method for identifying or classifying vulnerabilities. False positives and negatives are possible, and results should be verified by security professionals.

Additional information

In all honesty, I think I'm close to the limit for classification as a single label, given that a vulnerability can be classified as several CWEs.

For more details on the CWE standard, please visit Common Weakness Enumeration.

My report on this : Fine-tuning blogpost.