hellotung/Code_And_Res_copy

Top-K-CWE: Taxonomy-Aware Multi-Class CWE Classification on MegaVul

Reproducible code for the paper "Top-K-CWE: Taxonomy-Aware Multi-Class CWE Classification on MegaVul with Class-Balanced Learning."

We fine-tune pre-trained code encoders (CodeBERT, UniXcoder, CodeT5) to classify a vulnerable C/C++ function into one of K CWE categories drawn from the MITRE Top-25. On top of a strong class-imbalance recipe (Focal / Class- Balanced Focal + weighted oversampling) we add a taxonomy-aware learning layer that injects the CWE hierarchy (CWE-1000 pillars) into training:

• M1 — multi-task coarse head. An auxiliary classifier over the M coarse CWE groups, trained jointly with the fine head (L = L_fine + λ·L_coarse).
• M2 — hierarchy-aware soft labels. Mis-predicting a sibling CWE (same coarse parent, e.g. CWE-787 vs CWE-125) is penalised less than a cross-group error, by spreading a little label mass over siblings.

Both reduce exactly to the flat baseline when switched off, so the ablation isolates the taxonomy effect. CodeBERT and UniXcoder are RoBERTa-family encoders (pooled at [CLS]/<s>); CodeT5 is a T5 encoder-decoder, so we load only its encoder (T5EncoderModel) and pool with a masked mean. The dataset is split 80/10/10 chronologically by commit date to avoid temporal leakage.

Status. The pipeline is fully implemented and unit-tested. Runs have not been executed yet; metric cells in the paper are marked TBM and are filled in by the evaluation code (outputs/<run>/eval_test/metrics.json).

---

1. Quick Start

# 1) Environment
conda create -n topkcwe python=3.10 -y
conda activate topkcwe
pip install -r requirements.txt

# 2) Download MegaVul (manual step - see DATA.md), place megavul.json in data/raw/

# 3) Build the K-class dataset (Top-25 CWEs with >=100 samples, chronological split)
python -m src.data_prep \
    --raw_path data/raw/megavul.json \
    --out_dir  data/processed \
    --min_samples 100 --top25_only
#    -> always check data/processed/stats.json for the exact K and per-split counts.

# 4a) Train the main system (UniXcoder + CB-Focal + oversample)
python -m src.train --config configs/unixcoder_cbfocal.yaml --seed 42

# 4b) Train the full taxonomy-aware system (M1 + M2)
python -m src.train --config configs/unixcoder_cbfocal_mtl_hier.yaml --seed 42

# 5) Evaluate (macro/weighted-F1, top-1/3, MCC, coarse-group metrics, confusion)
python -m src.evaluate --config configs/unixcoder_cbfocal_mtl_hier.yaml \
                       --ckpt outputs/unixcoder_cbfocal_mtl_hier/best.pt --split test

# 6) Non-pretrained external baseline (CPU, a few minutes)
python -m src.baselines --data_dir data/processed --out_dir outputs/baseline_svm --model svm

Reproduce the whole study (phased; multi-seed; baselines) and aggregate:

bash scripts/run_all_experiments.sh                  # ENCODER=unixcoder by default
python scripts/collect_results.py --root outputs --out outputs/summary.csv

Hardware & budget (single NVIDIA RTX 3090, 24 GB)

A single run (train + eval, 10 epochs over 9.6K functions, batch 32, 512 tokens, fp16) takes **1–1.5 h** and usually early-stops around epoch 6–8. The full grid — 6 imbalance-core runs (single seed) + 2 main candidates × 3 seeds + 3 taxonomy-aware configs × 3 seeds + the SVM baseline — is ~25–35 GPU-hours. A complete single-seed table is available much sooner (~12–16 h).

If a config OOMs at batch 32 (most likely CodeT5 or the MTL head): set batch_size: 16 and grad_accum_steps: 2 in that YAML to keep the same effective batch size of 32.

---

2. Project Layout

code/
├── README.md / DATA.md / requirements.txt
├── configs/
│   ├── base.yaml                       <- shared defaults (+ taxonomy flags)
│   ├── {codebert,unixcoder,codet5}_{ce,focal}.yaml
│   ├── {unixcoder,codet5}_cbfocal.yaml          <- main candidates
│   └── {unixcoder,codet5}_cbfocal_{mtl,hier,mtl_hier}.yaml  <- taxonomy-aware
├── src/
│   ├── data_prep.py        <- filter + chronological 80/10/10 split
│   ├── dataset.py          <- PyTorch Dataset + collate
│   ├── models.py           <- encoder (+ T5 encoder) + fine head + optional coarse head (M1)
│   ├── losses.py           <- Focal, CB-Focal, weighted CE, hierarchy-aware soft labels (M2)
│   ├── cwe_hierarchy.py    <- CWE -> coarse-group mapping (CWE-1000) + helpers
│   ├── samplers.py         <- WeightedRandomSampler for oversampling
│   ├── train.py            <- training loop (MTL / hier / grad-accum / early stop)
│   ├── evaluate.py         <- fine + coarse metrics, in/cross-group error split, plots
│   ├── baselines.py        <- TF-IDF+LinearSVC (default) and BiLSTM (optional)
│   ├── utils.py / cwe_top25.py
├── scripts/
│   ├── run_all_experiments.sh   <- phased grid, 3 seeds, baselines
│   ├── collect_results.py       <- mean ± std summary across seeds (paper tables)
│   └── plot_confusion.py
└── tests/                       <- 37 unit tests (run: pytest -q)
    ├── test_data_prep.py        <- temporal split + label normalisation
    ├── test_losses.py           <- Focal/CB-Focal + hierarchy-aware loss math
    ├── test_cwe_hierarchy.py    <- mapping / coarse map / sibling matrix
    ├── test_models.py           <- pooling (CLS vs masked-mean)
    └── test_dataset.py

The CWE→coarse mapping in src/cwe_hierarchy.py is released so reviewers can inspect and contest it; the paper flags it as a construct-validity threat.

---

3. Reproducibility

• Fixed seeds (42, 1, 7); the two main candidates and the taxonomy-aware configs are run at all three and reported as mean ± std.
• All preprocessing is deterministic given (raw_path, min_samples, top25_only).
• The chronological split is by commit_date (ascending), then 80/10/10 across the global timeline. We accept the resulting mild class drift (reported in stats.json and discussed in the paper).
• tests/ runs without any network/model download (pooling and loss math are tested directly; a fake tokenizer stands in for HuggingFace).

4. License

Code released under MIT. MegaVul is distributed by its original authors under their own terms.