Hugging Face's logo

ctokx
/
cti-attack-mapper-modernbert

Text Classification
Transformers
Safetensors
English
modernbert
cyber-threat-intelligence
mitre-attack
multi-label-classification
defensive-security
blue-team
threat-intelligence
text-embeddings-inference
Model card Files
xet
 Community
cti-attack-mapper-modernbert
File size: 3,105 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
"""Fine-tune an encoder and evaluate it on the held-out test set.

    python scripts/03_train.py --model modernbert --scheme document
    python scripts/03_train.py --model modernbert --scheme random
"""

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from torch.utils.data import DataLoader  # noqa: E402

from cti_attack import config, data, evaluate, modeling  # noqa: E402


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=config.DEFAULT_MODEL, choices=list(config.BASE_MODELS))
    ap.add_argument("--scheme", default="document", choices=["document", "random"])
    ap.add_argument("--epochs", type=int, default=config.EPOCHS)
    args = ap.parse_args()

    labels = data.load_labels()
    tr = data.load_split(args.scheme, "train")
    dv = data.load_split(args.scheme, "dev")
    te = data.load_split(args.scheme, "test")

    run = f"{args.model}__{args.scheme}"
    out_dir = config.MODELS_DIR / run
    print(f"{'=' * 62}\n  {run}\n{'=' * 62}")
    print(f"  train={len(tr)}  dev={len(dv)}  test={len(te)}  labels={len(labels)}")

    best = modeling.train(args.model, args.scheme, tr, dv, labels, out_dir, epochs=args.epochs)

    # ---- evaluate the best checkpoint on dev (tuning) then test (reporting) --
    model, tok = modeling.load_for_inference(out_dir)
    dev_ = modeling.device()
    model.to(dev_)

    amp = args.model not in config.FP32_ONLY_MODELS
    ds_dv = modeling.SentenceDataset(dv, labels, tok, config.MAX_LENGTH)
    ds_te = modeling.SentenceDataset(te, labels, tok, config.MAX_LENGTH)
    s_dv = modeling.predict_scores(model, DataLoader(ds_dv, batch_size=32), dev_, amp=amp)
    s_te = modeling.predict_scores(model, DataLoader(ds_te, batch_size=32), dev_, amp=amp)

    Ydv = ds_dv.Y.astype("int8")
    Yte = ds_te.Y.astype("int8")

    gt, _ = evaluate.tune_global_threshold(Ydv, s_dv)
    pct = evaluate.tune_per_class_thresholds(Ydv, s_dv)

    rep_g = evaluate.evaluate(Yte, evaluate.apply_thresholds(s_te, gt), labels)
    rep_p = evaluate.evaluate(Yte, evaluate.apply_thresholds(s_te, pct), labels)

    print(f"\n  TEST global t={gt}   macro-F1={rep_g.macro_f1:.4f}  micro-F1={rep_g.micro_f1:.4f}")
    print(f"  TEST per-class       macro-F1={rep_p.macro_f1:.4f}  micro-F1={rep_p.micro_f1:.4f}")

    payload = {
        "run": run,
        "base_model": config.BASE_MODELS[args.model],
        "split_scheme": args.scheme,
        "best_epoch": best["epoch"],
        "dev_macro_f1": round(best["macro_f1"], 4),
        "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(),
        },
    }
    evaluate.save_report(args.model, args.scheme, payload)
    (out_dir / "thresholds.json").write_text(
        json.dumps({"global": gt,
                    "per_class": {l: float(t) for l, t in zip(labels, pct)}}, indent=2),
        encoding="utf-8")


if __name__ == "__main__":
    main()