src/cti_attack/config.py

"""Central configuration. Every path and constant used by the pipeline lives here."""

from __future__ import annotations

from pathlib import Path

# --- repo layout -----------------------------------------------------------
REPO_ROOT = Path(__file__).resolve().parents[2]
DATA_DIR = REPO_ROOT / "data"
RAW_DIR = DATA_DIR / "raw"
BUILD_DIR = DATA_DIR / "build"
RESULTS_DIR = REPO_ROOT / "results"
MODELS_DIR = REPO_ROOT / "models"

for _d in (RAW_DIR, BUILD_DIR, RESULTS_DIR, MODELS_DIR):
    _d.mkdir(parents=True, exist_ok=True)

# --- upstream sources (licences verified — see DATASET_CARD.md) -------------
TRAM_MULTILABEL_URL = (
    "https://raw.githubusercontent.com/center-for-threat-informed-defense/"
    "tram/main/data/tram2-data/multi_label.json"
)
TRAM_RAW = RAW_DIR / "tram_multi_label.json"

# The single-label TRAM file annotates one technique per sentence but covers a
# different set of sentences: ~4,300 of its labelled sentences never appear in
# multi_label.json, which more than doubles the labelled pool. Its 149 source
# documents are a subset of multi_label's 151, so the two files can only be
# merged *before* the document-level split — the split then keeps every
# document, from either file, on a single side. Merging after the split would
# leak. This is opt-in (build --include-single) so the pure-multi_label v1
# dataset stays reproducible.
TRAM_SINGLELABEL_URL = (
    "https://raw.githubusercontent.com/center-for-threat-informed-defense/"
    "tram/main/data/tram2-data/single_label.json"
)
TRAM_SINGLE_RAW = RAW_DIR / "tram_single_label.json"

ATTACK_STIX_URL = (
    "https://raw.githubusercontent.com/mitre-attack/attack-stix-data/master/"
    "enterprise-attack/enterprise-attack.json"
)
ATTACK_STIX_RAW = RAW_DIR / "enterprise-attack.json"
ATTACK_NAMES_JSON = BUILD_DIR / "attack_technique_names.json"
ATTACK_STATUS_JSON = BUILD_DIR / "attack_technique_status.json"

# --- dataset construction --------------------------------------------------
# A technique needs at least one document in each of train/dev/test to be
# trainable *and* evaluable, so it must appear in >= 3 distinct documents.
# This drops exactly one label: T1557.001 (1 document). The next-rarest
# technique, T1072, has 5 documents and is retained.
MIN_DOCS_PER_TECHNIQUE = 3

SPLIT_FRACTIONS = {"train": 0.70, "dev": 0.15, "test": 0.15}
SPLIT_SEED = 20260802

# The corpus prefixes each document's first sentence with scraped metadata like
# "title: <headline> url: <link>". It is boilerplate, not threat prose.
# Some rows carry a bare "title:" with no URL, so both forms are stripped: the
# full header first, then any leftover marker.
BOILERPLATE_PREFIX_RE = r"^\s*title:\s*.*?\s+url:\s*\S+\s*"
BOILERPLATE_MARKER_RE = r"^\s*title:\s*"

# --- modelling -------------------------------------------------------------
BASE_MODELS = {
    "modernbert": "answerdotai/ModernBERT-base",
    "deberta": "microsoft/deberta-v3-base",
    "securebert": "ehsanaghaei/SecureBERT",
}
DEFAULT_MODEL = "modernbert"

# DeBERTa-v3's disentangled attention overflows under reduced-precision
# autocast: it produced loss=nan from step 1 on this corpus and never
# recovered. It is trained in fp32 instead. Costs ~2x time and still fits in
# 8 GB at batch 16 / 256 tokens.
FP32_ONLY_MODELS = {"deberta"}

MAX_LENGTH = 256          # p95 sentence length is 45 words; 256 tokens is ample
BATCH_SIZE = 16
GRAD_ACCUM = 2            # effective batch 32, comfortable inside 8 GB
LEARNING_RATE = 3e-5
EPOCHS = 6
WARMUP_RATIO = 0.1
WEIGHT_DECAY = 0.01
SEED = 20260802

# Per-class decision thresholds are tuned on dev over this grid.
THRESHOLD_GRID = [round(0.05 * i, 2) for i in range(1, 20)]