bziemba/SecureBERT2.0-final

1---2 3license: mit4datasets:5- bziemba/cve_cwe_cvss6language:7- en8base_model:9- cisco-ai/SecureBERT2.0-biencoder10tags:11- cybersecurity12- vulnerability-classification13- cvss14- cwe15- securebert16- multi-task-learning17---18 19# SecureBERT Vulnerability Classifier (CVSS & CWE Flat Classifier)20 21This 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). 22 23It 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.24 25## 🎯 Intended Use26 27The primary use case is automating the initial **Vulnerability Triage** process. By inputting unstructured threat narratives, security analysts can instantly receive:28* **8 CVSS v3.1 Metrics:** Attack Vector, Attack Complexity, Privileges Required, User Interaction, Scope, Confidentiality, Integrity, and Availability.29* **CWE Classification:** Probabilistic mapping to the MITRE CWE tree across 4 levels of abstraction (Top-K predictions).30 31## 🧠 Model Architecture32 33The model uses a shared `SecureBERT2.0` backbone with 12 distinct classification heads attached to the pooled outputs:34* **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.35* **CWE Heads (4):** Multi-Layer Perceptrons (MLP) consisting of `LayerNorm -> Linear -> GELU -> Dropout -> Linear. These heads utilize the Mean-Pooled token embeddings.36 37## 📂 Repository Structure & Custom Config38 39Unlike standard Hugging Face models, this repository features a highly customized `config.json`. It dynamically dictates the architecture and handles label decoding.40* `cvss_map`: Contains the exact string labels for all 8 CVSS metrics (e.g., `["Network", "Adjacent", "Local", "Physical"]`).41* `cwe_labels`: Contains ID-to-Name mappings for all supported CWEs across `pillar`, `class`, `base`, and `variant` levels.42 43**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.44 45## 💻 Usage & Inference46 47Below is a complete, standalone Python snippet to load the model, tokenizer, and configuration directly from this Hugging Face repository and perform predictions.48 49```python50import json51import torch52import torch.nn as nn53import torch.nn.functional as F54from transformers import AutoConfig, AutoModel, AutoTokenizer55from huggingface_hub import hf_hub_download56 57# 1. Define the Custom Architecture58class SecureBERTFlatClassifier(nn.Module):59    def __init__(self, model_name, cvss_map, class_counts):60        super().__init__()61        config = AutoConfig.from_pretrained(model_name)62        if hasattr(config, "reference_compile"): config.reference_compile = False63        self.bert = AutoModel.from_pretrained(model_name, config=config)64        65        def make_head(out_features, is_cvss=False):66            layers =[67                nn.LayerNorm(768), nn.Dropout(0.1), 68                nn.Linear(768, 768), nn.GELU(), nn.Dropout(0.1),69                nn.Linear(768, 768), nn.GELU(), nn.Dropout(0.1),70                nn.Linear(768, out_features)71            ]72            if is_cvss: layers.append(nn.Softmax(dim=1))73            return nn.Sequential(*layers)74 75        self.cvss_heads = nn.ModuleDict({k: make_head(len(v), True) for k, v in cvss_map.items()})76        self.cwe_heads = nn.ModuleDict({k: make_head(v) for k, v in class_counts.items()})77 78    def forward(self, input_ids, attention_mask):79        out = self.bert(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state80        cls_emb = out[:, 0, :]81        mask = attention_mask.unsqueeze(-1).expand(out.size()).float()82        mean_emb = torch.sum(out * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)83        84        res = {}85        for k, head in self.cvss_heads.items(): res[k] = head(cls_emb)86        for k, head in self.cwe_heads.items(): res[k] = head(mean_emb)87        return res88 89# 2. Inference Wrapper90class VulnPredictor:91    def __init__(self, repo_id):92        self.device = "cuda" if torch.cuda.is_available() else "cpu"93        94        conf_path = hf_hub_download(repo_id=repo_id, filename="config.json")95        model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")96        97        with open(conf_path, "r") as f: self.config = json.load(f)98            99        base_model = self.config.get("base_model", "cisco-ai/SecureBERT2.0-biencoder")100        counts = {k: len(v) for k, v in self.config.get("cwe_labels", {}).items()}101        102        self.tokenizer = AutoTokenizer.from_pretrained(base_model)103        self.model = SecureBERTFlatClassifier(base_model, self.config["cvss_map"], counts)104        self.model.load_state_dict(torch.load(model_path, map_location=self.device), strict=False)105        self.model.to(self.device).eval()106 107    def predict(self, text, top_k=3):108        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(self.device)109        with torch.no_grad():110            out = self.model(inputs['input_ids'], inputs['attention_mask'])111        112        res = {'cvss': {}, 'cwe': {}}113        for task, labels in self.config.get("cvss_map", {}).items():114            score, idx = torch.max(out[task], dim=1)115            res['cvss'][task] = {"value": labels[idx.item()], "confidence": round(score.item(), 4)}116            117        for lv, cwe_data in self.config.get("cwe_labels", {}).items():118            if lv in out:119                probs = F.softmax(out[lv], dim=1)120                scores, idxs = torch.topk(probs, k=min(top_k, probs.size(1)))121                res['cwe'][lv] =[122                    {"id": int(str(cwe_data[i.item()]['id']).replace('CWE-','')), 123                     "name": cwe_data[i.item()]['name'], 124                     "score": round(s.item(), 4)} for s, i in zip(scores[0], idxs[0])125                ]126        return res127 128# 3. Quickstart129if __name__ == "__main__":130    REPO_ID = "bziemba/SecureBERT2.0-final" 131    132    predictor = VulnPredictor(REPO_ID)133    134    sample_cve = "An issue was discovered in the login panel allowing attackers to bypass authentication via crafted SQL queries."135    results = predictor.predict(sample_cve)136    137    print(json.dumps(results, indent=2))138 139