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