- Source field: each vulnerability entry now includes a
sourcefield identifying its origin (cvelistv5, github, pysec, cnvd, csaf_*). - Dynamic dataset card for multi-source datasets: when generating a dataset from multiple sources (e.g.,
--sources cvelistv5,github,csaf_redhat,csaf_cisco,csaf_cisa,pysec), a dataset card is now automatically generated with a per-source breakdown table showing entry counts and percentages.
- Per-class metrics for severity trainer (
classify_severity.py):compute_metricsnow reports precision, recall, and F1 per class (Low/Medium/High/Critical) alongside overall accuracy and macro F1. - Best model checkpoint selection (
classify_severity.py): model selected by accuracy instead of eval_loss,save_total_limitincreased from 2 to 3.
- Moved all HuggingFace card templates (dataset cards, model cards) to a dedicated
vulntrain/cards/directory. - Updated dependencies.
- Data leakage fix: new
deduplicate_split()function groups entries by description text before splitting, preventing identical descriptions from appearing in both train and test sets. CNVD reuses boilerplate descriptions across different vulnerability IDs, which previously contaminated 15.6% of the test set. - Per-class metrics:
compute_metricsnow reports precision, recall, and F1 per class (Low/Medium/High) alongside overall accuracy and macro F1 at each evaluation epoch. - Class weighting options: new
--class-weightsflag (none,sqrt,balanced,focal) for experimenting with class imbalance strategies. Includes aFocalLossTrainerimplementation (Lin et al., 2017). Defaults to uniform loss after experiments showed all weighting strategies degrade Medium recall disproportionately. - Best model checkpoint selection:
metric_for_best_modelset toaccuracy(was defaulting toeval_loss),save_total_limitincreased from 2 to 3. - Dynamic model card: model card is now a template populated with actual eval metrics from
trainer.evaluate()after each training run. Documents per-class metrics, training configuration, and known limitations.
- CVE cross-references:
extract_cnvdnow extracts thecve_idfield fromcves.cve.cveNumber, enabling cross-referencing with CVE equivalents (~81% of entries). - Dataset card: new dataset card documenting severity distribution, CVE overlap rates, coverage decline post-RMSV (94% published in 2015 to 4% in 2023), and duplicate description caveat.
- CNVD severity model comparison validator: new
validators/severity_cnvd.pyscript to evaluate old and new models side by side on the same deduplicated test set, reporting per-class metrics, confusion matrices, and summary deltas.
- CodeCarbon tracker hang: replaced
@track_emissionsdecorator with explicitEmissionsTrackerscoped totrainer.train()only. Removedpush_to_hub=TruefromTrainingArgumentswhich caused the trainer to upload internally before returning. Applied to bothclassify_severity_cnvd.pyandclassify_severity.py.
- Added technical report:
docs/cnvd-severity-improvements.md. - Improved
docs/index.mdwith configuration section, CNVD-specific details, and validator usage.
Thanks to Eric Romang for his thorough independent analysis (VulnTrain#19) that prompted these improvements.
- New CLI options for severity classification trainer (
classify_severity.py):--no-codecarbon: Disable CodeCarbon emissions tracking.--no-push: Disable pushing the model and tokenizer to Hugging Face Hub.--no-cache: Disable cache for the model during training.
- CWE/Patch dataset improvements: Considered more fields to find vulnerability patches. Asynchronous requests to GitHub are now less aggressive.
- CWE Guesser dataset:
- Now uses the new vulnerability endpoint of Vulnerability-Lookup.
- References in security advisories without the
patchtag are also considered. - Repo ID is now a configurable parameter in the dataset generation script.
- URL handling improvements:
normalize_patch_urlfunction improved for better patch URL processing.- URLs with fragments are now properly handled.
- Concurrency: Reduced the number of default concurrent requests to 12 to avoid overloading external services.
- Updated Python dependencies, including PyTorch bump from 2.7.1 to 2.8.0.
- General dependency updates across the project.
- Minor code improvements and style updates (reformatted with
black).
- Dataset generation: Introduced a new script to build datasets of structured vulnerabilities enriched with CWE identifiers and corresponding patches. Each entry now includes the Git commit message and the full diff (Base64-encoded). #10 by @3LS3-1F
- Model generation: Added a new trainer for predicting CWE classifications from vulnerability descriptions and associated patches (commit messages). #10 by @3LS3-1F
Related resources shared via Hugging Face: https://huggingface.co/collections/CIRCL/vlai-for-cwe-guessing-68bab22e3d71b513146d13b3
- Improved documentation and reorganized modules for better clarity and maintainability.
- Updated dependencies to their latest stable versions.
- Dataset generation: Associating Git Fixes with Common Weakness Enumerations (CWEs) found in security advisories. ((#4)[#4])
- A documentation is now available. (8a345ca)
- Model generation: Added a boolean parameter in map_cvss_to_severity in order to switch between using the first non-null CVSS score or the mean of all available CVSS scores. (ff6616e)
- Dataset generation: Removed useless keys in extract_cnvd
- (b7d694)
This version adds support for creating new AI-ready datasets based on the China National Vulnerability Database (CNVD). It also introduces a new training module designed to classify vulnerabilities using text classification models tailored for CNVD data. By Léa
- Updated dependencies and fixed issues due to changes in transformers.
-
Updated dependencies.
- Dataset generation: CVSS are now extracted from GitHub and PySec security advisories.
- Dataset generation: CVSS, CPE, title and description (summary) are now extracted from CSAF document.
- Trainers: Support of roberta-base for the text classifier with improved settings for TrainingArguments.
- Validators: Validator for severity classification.
- Introduced a new trainer to automatically classify vulnerabilities based on their descriptions,
even when CVSS scores are unavailable. - Added CVSS parsing to the dataset generation script.
- Refactored the project structure for better organization.
- Improved CPE parsing.
- Enhanced the dataset generation script.
- Optimized the trainer for text generation on vulnerability descriptions.
- Improved command-line argument parsing.
- Improved the process of pushing the tokenizer and trainer to Hugging Face.
Fixed configuration module name.
Added support of configuration file.
The dataset generation step now uses data from GitHub Advisories, and the VulnExtractor cleans the summary and details fields.
Dataset generation: allow specifying a commit message when uploading to Hugging Face.
Validation: Added a simple validation script for model optimized for text generation. The script is able to pull a model and send tasks via a Pipeline
For the training step: added the choices of model: gpt2, distilgpt2, meta-llama/Llama-3.3-70B-Instruct, distilbert-base-uncased
Various improvements to the command line parsing.
Added a trainer. Experimenting distilbert-base-uncased (AutoModelForMaskedLM) and gpt2 (AutoModelForCausalLM). The goal is to generate text.
Various improvements to the dataset generator. And added a command line parser.
First release with upload of datasets to HuggingFace.
Datasets are build based on NIST data with enrichment from FKIE and vulnrichment