--- license: mit datasets: - bziemba/cve_cwe_cvss language: - en base_model: - cisco-ai/SecureBERT2.0-biencoder tags: - cybersecurity - vulnerability-classification - cvss - cwe - securebert - multi-task-learning --- # SecureBERT Vulnerability Classifier (CVSS & CWE Flat Classifier) This model automatically analyzes raw vulnerability descriptions (e.g., CVE reports, bug bounty submissions) and predicts **CVSS v3.1 metrics** alongside a 4-level **CWE taxonomy** (Pillar, Class, Base, Variant). It is a fine-tuned version of the domain-specific [`cisco-ai/SecureBERT2.0`](https://huggingface.co/cisco-ai/SecureBERT2.0) utilizing a Multi-Task Learning (MTL) architecture with flat classification heads. ## 🎯 Intended Use The primary use case is automating the initial **Vulnerability Triage** process. By inputting unstructured threat narratives, security analysts can instantly receive: * **8 CVSS v3.1 Metrics:** Attack Vector, Attack Complexity, Privileges Required, User Interaction, Scope, Confidentiality, Integrity, and Availability. * **CWE Classification:** Probabilistic mapping to the MITRE CWE tree across 4 levels of abstraction (Top-K predictions). ## 🧠 Model Architecture The model uses a shared `SecureBERT2.0` backbone with 12 distinct classification heads attached to the pooled outputs: * **CVSS Heads (8):** Multi-Layer Perceptrons (MLP) consisting of `LayerNorm -> Linear -> GELU -> Dropout -> Linear -> Softmax`. They use the `[CLS]` token embedding to predict nominal and ordinal CVSS categories. * **CWE Heads (4):** Multi-Layer Perceptrons (MLP) consisting of `LayerNorm -> Linear -> GELU -> Dropout -> Linear. These heads utilize the Mean-Pooled token embeddings. ## 📂 Repository Structure & Custom Config Unlike standard Hugging Face models, this repository features a highly customized `config.json`. It dynamically dictates the architecture and handles label decoding. * `cvss_map`: Contains the exact string labels for all 8 CVSS metrics (e.g., `["Network", "Adjacent", "Local", "Physical"]`). * `cwe_labels`: Contains ID-to-Name mappings for all supported CWEs across `pillar`, `class`, `base`, and `variant` levels. **Note:** Because of the custom multi-head architecture, you cannot use the default `AutoModelForSequenceClassification`. You must define the custom PyTorch class provided in the usage snippet below. ## 💻 Usage & Inference Below is a complete, standalone Python snippet to load the model, tokenizer, and configuration directly from this Hugging Face repository and perform predictions. ```python import json import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoConfig, AutoModel, AutoTokenizer from huggingface_hub import hf_hub_download # 1. Define the Custom Architecture class SecureBERTFlatClassifier(nn.Module): def __init__(self, model_name, cvss_map, class_counts): super().__init__() config = AutoConfig.from_pretrained(model_name) if hasattr(config, "reference_compile"): config.reference_compile = False self.bert = AutoModel.from_pretrained(model_name, config=config) def make_head(out_features, is_cvss=False): layers =[ nn.LayerNorm(768), nn.Dropout(0.1), nn.Linear(768, 768), nn.GELU(), nn.Dropout(0.1), nn.Linear(768, 768), nn.GELU(), nn.Dropout(0.1), nn.Linear(768, out_features) ] if is_cvss: layers.append(nn.Softmax(dim=1)) return nn.Sequential(*layers) self.cvss_heads = nn.ModuleDict({k: make_head(len(v), True) for k, v in cvss_map.items()}) self.cwe_heads = nn.ModuleDict({k: make_head(v) for k, v in class_counts.items()}) def forward(self, input_ids, attention_mask): out = self.bert(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state cls_emb = out[:, 0, :] mask = attention_mask.unsqueeze(-1).expand(out.size()).float() mean_emb = torch.sum(out * mask, 1) / torch.clamp(mask.sum(1), min=1e-9) res = {} for k, head in self.cvss_heads.items(): res[k] = head(cls_emb) for k, head in self.cwe_heads.items(): res[k] = head(mean_emb) return res # 2. Inference Wrapper class VulnPredictor: def __init__(self, repo_id): self.device = "cuda" if torch.cuda.is_available() else "cpu" conf_path = hf_hub_download(repo_id=repo_id, filename="config.json") model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin") with open(conf_path, "r") as f: self.config = json.load(f) base_model = self.config.get("base_model", "cisco-ai/SecureBERT2.0-biencoder") counts = {k: len(v) for k, v in self.config.get("cwe_labels", {}).items()} self.tokenizer = AutoTokenizer.from_pretrained(base_model) self.model = SecureBERTFlatClassifier(base_model, self.config["cvss_map"], counts) self.model.load_state_dict(torch.load(model_path, map_location=self.device), strict=False) self.model.to(self.device).eval() def predict(self, text, top_k=3): inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(self.device) with torch.no_grad(): out = self.model(inputs['input_ids'], inputs['attention_mask']) res = {'cvss': {}, 'cwe': {}} for task, labels in self.config.get("cvss_map", {}).items(): score, idx = torch.max(out[task], dim=1) res['cvss'][task] = {"value": labels[idx.item()], "confidence": round(score.item(), 4)} for lv, cwe_data in self.config.get("cwe_labels", {}).items(): if lv in out: probs = F.softmax(out[lv], dim=1) scores, idxs = torch.topk(probs, k=min(top_k, probs.size(1))) res['cwe'][lv] =[ {"id": int(str(cwe_data[i.item()]['id']).replace('CWE-','')), "name": cwe_data[i.item()]['name'], "score": round(s.item(), 4)} for s, i in zip(scores[0], idxs[0]) ] return res # 3. Quickstart if __name__ == "__main__": REPO_ID = "bziemba/SecureBERT2.0-final" predictor = VulnPredictor(REPO_ID) sample_cve = "An issue was discovered in the login panel allowing attackers to bypass authentication via crafted SQL queries." results = predictor.predict(sample_cve) print(json.dumps(results, indent=2))