--- 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