README.md

---
license: apache-2.0
base_model: unsloth/gemma-4-12b-it
datasets:
- exploitintel/cve-cwe-consensus
language:
- en
pipeline_tag: text-generation
library_name: transformers
tags:
- cybersecurity
- cve
- cwe
- vulnerability
- text-classification
- gemma4
- qlora
---

# cve-cwe-gemma4-12b

A [Gemma 4 12B](https://huggingface.co/unsloth/gemma-4-12b-it) 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](https://huggingface.co/exploitintel/cve-cwe-gemma4-12b/blob/main/blog.md)

- **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](https://cwe.mitre.org/data/definitions/1003.html) (~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**](https://huggingface.co/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).

| metric | this model (bf16) | v1 baseline* |
|---|---|---|
| exact-match | **0.714** | 0.29 |
| micro-F1 | **0.756** | 0.32 |
| macro-F1 | **0.538** | 0.067 |
| easy exact-match | 0.805 | — |
| hard exact-match | 0.644 | — |

\* 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).

```python
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](https://github.com/unslothai/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.