MITRE-v16-tactic-bert-case-based

MITRE-v16-tactic-bert-case-based

It's a fine-tuned model from mitre-bert-base-cased on the MITRE ATT&CK version 16 procedure dataset.

Intended uses & limitations

You can use the fine-tuned model for text classification. It aims to identify the tactic that the sentence belongs to in MITRE ATT&CK framework. A sentence or an attack may fall into several tactics.

Note that this model is primarily fine-tuned on text classification for cybersecurity. It may not perform well if the sentence is not related to attacks.

How to use

You can use the model with Tensorflow.

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
model_id = "sarahwei/MITRE-v16-tactic-bert-case-based"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
)
question = 'An attacker performs a SQL injection.'
input_ids = tokenizer(question,return_tensors="pt")
outputs = model(**input_ids)
logits = outputs.logits
sigmoid = torch.nn.Sigmoid()
probs = sigmoid(logits.squeeze().cpu())
predictions = np.zeros(probs.shape)
predictions[np.where(probs >= 0.5)] = 1
predicted_labels = [model.config.id2label[idx] for idx, label in enumerate(predictions) if label == 1.0]

Training procedure

Training parameter

• learning_rate: 2e-5
• train_batch_size: 32
• eval_batch_size: 32
• seed: 0
• num_epochs: 5
• warmup_ratio: 0.01
• weight_decay: 0.001
• optim: adamw_8bit

Training results

• global_step=1755
• train_runtime: 315.2685
• train_samples_per_second: 177.722
• train_steps_per_second: 5.567
• total_flos: 7371850396784640.0
• train_loss: 0.06630994546787013