README.md

---
base_model: hfl/chinese-macbert-base
datasets:
- CIRCL/Vulnerability-CNVD
library_name: transformers
license: apache-2.0
metrics:
- accuracy
tags:
- generated_from_trainer
- text-classification
- classification
- nlp
- chinese
- vulnerability
pipeline_tag: text-classification
language: zh
model-index:
- name: vulnerability-severity-classification-chinese-macbert-base
  results: []
---

# VLAI: A RoBERTa-Based Model for Automated Vulnerability Severity Classification

This model, named **VLAI**, is a fine-tuned version of [hfl/chinese-macbert-base](https://huggingface.co/hfl/chinese-macbert-base) on the dataset [CIRCL/Vulnerability-CNVD](https://huggingface.co/datasets/CIRCL/Vulnerability-CNVD).

The model was presented in the paper [VLAI: A RoBERTa-Based Model for Automated Vulnerability Severity Classification](https://huggingface.co/papers/2507.03607).

**Abstract:** VLAI is a transformer-based model that predicts software vulnerability severity levels directly from text descriptions. Built on RoBERTa, VLAI is fine-tuned on over 600,000 real-world vulnerabilities and achieves over 82% accuracy in predicting severity categories, enabling faster and more consistent triage ahead of manual CVSS scoring. The model and dataset are open-source and integrated into the Vulnerability-Lookup service.

For more information, visit the [Vulnerability-Lookup project page](https://vulnerability.circl.lu) or the [ML-Gateway GitHub repository](https://github.com/vulnerability-lookup/ML-Gateway), which demonstrates its usage in a FastAPI server.

It achieves the following results on the evaluation set:
- Loss: 0.5994
- Accuracy: 0.7900

## How to use

You can use this model directly with the Hugging Face `transformers` library for text classification:

```python
from transformers import pipeline

classifier = pipeline(
    "text-classification",
    model="CIRCL/vulnerability-severity-classification-chinese-macbert-base"
)

# Example usage for a Chinese vulnerability description
description_chinese = "TOTOLINK A3600R是中国吉翁电子（TOTOLINK）公司的一款6天线1200M无线路由器。TOTOLINK A3600R存在缓冲区溢出漏洞，该漏洞源于/cgi-bin/cstecgi.cgi文件的UploadCustomModule函数中的File参数未能正确验证输入数据的长度大小，攻击者可利用该漏洞在系统上执行任意代码或者导致拒绝服务。"
result_chinese = classifier(description_chinese)
print(result_chinese)
# Expected output example: [{'label': '高', 'score': 0.9802}]
```

## Training procedure

### Training hyperparameters

The following hyperparameters were used during training:
- learning_rate: 3e-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

| Training Loss | Epoch | Step  | Validation Loss | Accuracy |
|:-------------:|:-----:|:-----:|:---------------:|:--------:|
| 0.65          | 1.0   | 3388  | 0.5772          | 0.7561   |
| 0.582         | 2.0   | 6776  | 0.5656          | 0.7620   |
| 0.5284        | 3.0   | 10164 | 0.5274          | 0.7881   |
| 0.3406        | 4.0   | 13552 | 0.5555          | 0.7869   |
| 0.3224        | 5.0   | 16940 | 0.5994          | 0.7900   |


### Framework versions

- Transformers 4.51.3
- Pytorch 2.7.1+cu126
- Datasets 3.6.0
- Tokenizers 0.21.1