You can deploy this model and start earning money today!

Dunateo/
roberta-cwe-classifier-kelemia-v0.2

Name: roberta-cwe-classifier-kelemia-v0.2
Availability: InStock

Views

No views yet

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:

python

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:

json

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

python

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

Epoch	Training Loss	Validation Loss
1.0	1.246200	1.183215
2.0	0.990800	0.999081
3.0	0.834400	0.942376

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.