src/cti_attack/predict.py

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