|
|
|
|
|
---
|
|
|
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
|
|
|
|
