File size: 3,477 Bytes
468c4c2 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """Inference API — the thing an analyst actually calls.
from cti_attack.predict import AttackMapper
mapper = AttackMapper("models/modernbert__document")
mapper.predict("The malware base64-encodes its configuration before writing it to disk.")
# [Prediction(technique_id='T1027', name='Obfuscated Files or Information', score=0.91)]
Thresholds are loaded from the model directory, not hard-coded, so predictions
match the numbers reported in the model card.
"""
from __future__ import annotations
import json
from dataclasses import dataclass
from pathlib import Path
import numpy as np
import torch
from . import attack_meta
@dataclass
class Prediction:
technique_id: str
name: str
score: float
def __repr__(self) -> str:
return f"Prediction({self.technique_id}, {self.name!r}, {self.score:.3f})"
class AttackMapper:
"""Maps threat-report sentences to MITRE ATT&CK technique IDs."""
def __init__(self, model_dir: str | Path, threshold_mode: str = "per_class", device: str | None = None):
from .modeling import load_for_inference
self.model_dir = Path(model_dir)
self.model, self.tokenizer = load_for_inference(self.model_dir)
self.device = torch.device(
device or ("cuda" if torch.cuda.is_available() else "cpu"))
self.model.to(self.device)
self.labels = json.loads((self.model_dir / "labels.json").read_text(encoding="utf-8"))
thr_path = self.model_dir / "thresholds.json"
if thr_path.exists():
thr = json.loads(thr_path.read_text(encoding="utf-8"))
if threshold_mode == "per_class":
self.thresholds = np.array([thr["per_class"][l] for l in self.labels])
else:
self.thresholds = np.full(len(self.labels), float(thr["global"]))
else:
self.thresholds = np.full(len(self.labels), 0.5)
try:
self.names = attack_meta.build_technique_names()
except Exception: # offline: degrade to bare IDs rather than failing
self.names = {}
@torch.no_grad()
def scores(self, sentences: list[str], batch_size: int = 32) -> np.ndarray:
if not sentences:
return np.zeros((0, len(self.labels)), dtype=np.float32)
out = []
for i in range(0, len(sentences), batch_size):
enc = self.tokenizer(
sentences[i:i + batch_size],
truncation=True, max_length=256, padding=True, return_tensors="pt",
).to(self.device)
logits = self.model(**enc).logits
out.append(torch.sigmoid(logits.float()).cpu().numpy())
return np.concatenate(out, axis=0)
def predict(self, sentence: str, top_k: int | None = None) -> list[Prediction]:
return self.predict_batch([sentence], top_k=top_k)[0]
def predict_batch(
self, sentences: list[str], top_k: int | None = None
) -> list[list[Prediction]]:
S = self.scores(sentences)
results = []
for row in S:
hits = np.where(row >= self.thresholds)[0]
if top_k is not None:
hits = hits[np.argsort(-row[hits])][:top_k]
else:
hits = hits[np.argsort(-row[hits])]
results.append([
Prediction(self.labels[j], self.names.get(self.labels[j], self.labels[j]), float(row[j]))
for j in hits
])
return results
|