File size: 3,165 Bytes
0f27fb6 | 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 | """Run the non-neural baselines on both split schemes.
python scripts/02_run_baselines.py
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
import numpy as np # noqa: E402
from cti_attack import baselines, data, evaluate # noqa: E402
def run_scheme(scheme: str, labels: list[str]) -> dict:
print(f"\n{'=' * 62}\n split scheme: {scheme}\n{'=' * 62}")
tr = data.load_split(scheme, "train")
dv = data.load_split(scheme, "dev")
te = data.load_split(scheme, "test")
Ytr, Ydv, Yte = (evaluate.to_matrix(x, labels) for x in (tr, dv, te))
txt_tr = [r["sentence"] for r in tr]
txt_dv = [r["sentence"] for r in dv]
txt_te = [r["sentence"] for r in te]
results = {}
# ---- frequency prior -------------------------------------------------
dev_s = baselines.frequency_scores(Ytr, len(dv))
test_s = baselines.frequency_scores(Ytr, len(te))
results["frequency"] = _score("frequency", Ydv, dev_s, Yte, test_s, labels)
# ---- ATT&CK keyword match -------------------------------------------
dev_s = baselines.keyword_scores(txt_dv, labels)
test_s = baselines.keyword_scores(txt_te, labels)
results["keyword"] = _score("keyword", Ydv, dev_s, Yte, test_s, labels)
# ---- TF-IDF + one-vs-rest logistic regression ------------------------
print(" fitting tfidf_lr (49 one-vs-rest classifiers) …")
s = baselines.tfidf_lr_scores(txt_tr, Ytr, {"dev": txt_dv, "test": txt_te})
results["tfidf_lr"] = _score("tfidf_lr", Ydv, s["dev"], Yte, s["test"], labels)
return results
def _score(name, Ydv, dev_scores, Yte, test_scores, labels) -> dict:
gt, _ = evaluate.tune_global_threshold(Ydv, dev_scores)
pct = evaluate.tune_per_class_thresholds(Ydv, dev_scores)
rep_g = evaluate.evaluate(Yte, evaluate.apply_thresholds(test_scores, gt), labels)
rep_p = evaluate.evaluate(Yte, evaluate.apply_thresholds(test_scores, pct), labels)
print(f" {name:11} global t={gt:<5} macro-F1={rep_g.macro_f1:.4f} "
f"micro-F1={rep_g.micro_f1:.4f} | per-class macro-F1={rep_p.macro_f1:.4f} "
f"micro-F1={rep_p.micro_f1:.4f}")
return {
"global_threshold": {"threshold": gt, **rep_g.as_dict()},
"per_class_threshold": {
"thresholds": {l: float(t) for l, t in zip(labels, pct)},
**rep_p.as_dict(),
},
}
def main() -> None:
labels = data.load_labels()
all_results = {}
for scheme in ("document", "random"):
all_results[scheme] = run_scheme(scheme, labels)
evaluate.save_report("baselines", scheme, all_results[scheme])
print(f"\n{'=' * 62}\n leakage gap (test macro-F1, per-class thresholds)\n{'=' * 62}")
for name in ("frequency", "keyword", "tfidf_lr"):
d = all_results["document"][name]["per_class_threshold"]["macro_f1"]
r = all_results["random"][name]["per_class_threshold"]["macro_f1"]
print(f" {name:11} document={d:.4f} random={r:.4f} "
f"inflation={r - d:+.4f} ({(r / d - 1) * 100 if d else float('nan'):+.1f}%)")
if __name__ == "__main__":
main()
|