bziemba/SecureBERT-MultiTask-CVSS-CWE

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