exploitintel / cve-cwe-gemma4-12b

cve-cwe-gemma4-12b

A Gemma 4 12B fine-tune that maps a CVE description to its CWE ID(s) .

📖 Write-up: From Essays to CWE-319 — how this fine-tune beats stock Gemma 4 at CWE classification

• Input: a free-text vulnerability description (text only).
• Output: the CWE ID(s) it maps to, comma-separated — e.g. CWE-79 or CWE-89, CWE-352 .
• Label space: MITRE CWE View-1003 (~117 weakness classes). Multi-label.

This is the merged 16-bit (bf16) model for transformers / vLLM / TGI. Quantized GGUFs for Ollama and llama.cpp are at exploitintel/cve-cwe-gemma4-12b-GGUF .

Results

Held-out test split ( exploitintel/cve-cwe-consensus , 10,514 examples), greedy decoding, description-only (no CVE-ID or label metadata in the prompt). Rows are split into easy (the weakness is named in the text) vs hard (it must be inferred).

* v1 baseline = a 1-epoch Gemma-4-E4B fine-tune. The headline gain is macro-F1 (the rare-CWE long tail), which improves ~8×; hard (must-infer) exact-match of 0.644 is close to easy (0.805), indicating the model genuinely infers weaknesses rather than only keyword-matching.

Usage

Requires transformers >= 5.10 (Gemma 4 is the gemma4_unified architecture).

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "exploitintel/cve-cwe-gemma4-12b"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto", device_map="auto").eval()

cve = ("A vulnerability in the login form allows remote attackers to execute "
       "arbitrary SQL commands via the username parameter.")
messages = [
    {"role": "system", "content": "You are a vulnerability analyst. Given a CVE "
     "description, reply with only the CWE ID(s) it maps to, comma-separated."},
    {"role": "user", "content": cve},
]
prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tok(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=32, do_sample=False)
print(tok.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))
# -> CWE-89

Training

• Base: unsloth/gemma-4-12b-it (4-bit QLoRA, bitsandbytes nf4).
• Method: LoRA (r=16), 3 epochs, context length 512, full-sequence SFT.
• Data: exploitintel/cve-cwe-consensus (train split, 50,074 examples).
• Hardware: single NVIDIA RTX 5090; ~7.1 h wall, ~17 GB peak VRAM.
• Trained with Unsloth .

Intended use & limitations

• Intended use: triage assistance — suggesting candidate CWE mappings for a CVE description.
• It is description-only : quality depends on how well the text describes the weakness. Vague descriptions yield weaker predictions (see the hard split).
• It can predict CWEs outside the true set; treat outputs as suggestions, not authoritative classifications, and keep a human in the loop for security-relevant decisions.
• Scope is MITRE View-1003; CWEs outside that view are not modeled.

License

Apache-2.0, inherited from the Gemma 4 base model.