File size: 3,688 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 | """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)]
|