File size: 1,259 Bytes
5956d68 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
---
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
|