Vulnerability Severity Classification Distilbert Base Uncased

🚀 vulnerability-severity-classification-distilbert-base-uncased

This model is fine - tuned from DistilBERT to classify vulnerability severity based on descriptions, achieving good accuracy on the evaluation set.

🚀 Quick Start

This model is a fine - tuned version of distilbert/distilbert-base-uncased on the dataset CIRCL/vulnerability-scores.

It achieves the following results on the evaluation set:

• Loss: 0.6447
• Accuracy: 0.7595

✨ Features

It is a classification model and is aimed to assist in classifying vulnerabilities by severity based on their descriptions.

💻 Usage Examples

Basic Usage

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

labels = ["low", "medium", "high", "critical"]

model_name = "vulnerability-severity-classification-distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.eval()

test_description = "langchain_experimental 0.0.14 allows an attacker to bypass the CVE-2023-36258 fix and execute arbitrary code via the PALChain in the python exec method."
inputs = tokenizer(test_description, return_tensors="pt", truncation=True, padding=True)

# Run inference
with torch.no_grad():
    outputs = model(**inputs)
    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)


# Print results
print("Predictions:", predictions)
predicted_class = torch.argmax(predictions, dim=-1).item()
print("Predicted severity:", labels[predicted_class])

🔧 Technical Details

Training hyperparameters

The following hyperparameters were used during training:

• learning_rate: 2e-05
• train_batch_size: 32
• eval_batch_size: 32
• seed: 42
• optimizer: Use OptimizerNames.ADAMW_TORCH with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
• lr_scheduler_type: linear
• num_epochs: 5

Training results

Framework versions

• Transformers 4.49.0
• Pytorch 2.6.0+cu124
• Datasets 3.3.2
• Tokenizers 0.21.0

📄 License

This project is licensed under the Apache - 2.0 license.