|
|
| ---
|
| tags:
|
| - security
|
| - cvss
|
| - cwe
|
| - bert
|
| library_name: transformers
|
| base_model: cisco-ai/SecureBERT2.0-biencoder
|
| ---
|
|
|
| # SecureBERT Multi-Task: CVSS & CWE Predictor
|
|
|
| This model predicts **CVSS v3.1 Metrics** and **CWE Hierarchy** (Pillar, Class, Base, Variant) from vulnerability descriptions.
|
|
|
| ## Usage
|
|
|
| Since this model uses a custom architecture, you need to load the class definition first.
|
|
|
| ```python
|
| import torch
|
| from transformers import AutoTokenizer
|
| # Download model_class.py from this repo or copy the definition
|
| from model_class import SecureBERTMultiHead
|
|
|
| # 1. Load Model
|
| model = SecureBERTMultiHead("cisco-ai/SecureBERT2.0-biencoder")
|
| checkpoint = torch.load(
|
| "pytorch_model.bin", # Downloaded from this repo
|
| map_location="cpu"
|
| )
|
| model.load_state_dict(checkpoint)
|
| model.eval()
|
|
|
| # 2. Tokenizer
|
| tokenizer = AutoTokenizer.from_pretrained("bziemba/SecureBERT-MultiTask-CVSS-CWE")
|
|
|
| # 3. Inference
|
| text = "SQL injection vulnerability in the login form..."
|
| inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
|
|
|
| with torch.no_grad():
|
| outputs = model(inputs['input_ids'], inputs['attention_mask'])
|
| print(outputs['attack_vector']) # CVSS Probabilities
|
| |
