Update README.md

README.md

@@ -1,9 +1,139 @@
 ---
-license: apache-2.0
+
+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__":
+    # REPALCE WITH YOUR REPOSITORY ID
+    REPO_ID = "YourUsername/Your-Repo-Name" 
+    
+    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))
+