1#!/usr/bin/env python
2"""
306_split_binary.py
4~~~~~~~~~~~~~~~~~~
5
6Stream-splits a JSONL cybersecurity corpus into *offensive*, *defensive*, and *other* shards
7using **two** fine-tuned SecureBERT heads.
8
9How the two heads work together
10-------------------------------
11We load two independent checkpoints:
12
13* `offensive_vs_rest` → gives **P(offensive | text)**
14* `defensive_vs_rest` → gives **P(defensive | text)**
15
16For every line we:
17
181. run both heads in the same GPU batch;
192. take the positive-class probability from each soft-max;
203. compare against per-head thresholds (from `thresholds.json`, default 0.5);
214. route the text with this truth table
22"""
23
24from __future__ import annotations
25
26import argparse
27import json
28from itertools import islice
29from pathlib import Path
30
31import torch
32from torch.nn.functional import softmax
33from tqdm.auto import tqdm
34from transformers import (
35 AutoModelForSequenceClassification as HFModel,
36 AutoTokenizer,
37)
38
39from config import RAW_JSONL, MODEL_DIR # MODEL_DIR == securebert_finetuned
40
41# ───────────────────────────── GPU SETTINGS ──────────────────────────
42# 1. Use TensorFloat-32 on Ada GPUs (gives a big matmul speed boost).
43torch.backends.cuda.matmul.allow_tf32 = True
44torch.set_float32_matmul_precision("medium")
45
46DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
47
48# ──────────────────────────────── CLI ────────────────────────────────
49cli = argparse.ArgumentParser(description="Split JSONL into offence/defence/other")
50cli.add_argument("--batch_size", type=int, help="override auto batch sizing")
51args = cli.parse_args()
52
53# ───────────────────── BATCH-SIZE HEURISTIC ──────────────────────────
54if args.batch_size: # user override wins
55 BATCH = args.batch_size
56else:
57 try:
58 import pynvml
59
60 pynvml.nvmlInit()
61 free = (
62 pynvml.nvmlDeviceGetMemoryInfo(pynvml.nvmlDeviceGetHandleByIndex(0)).free
63 / 1024**3
64 )
65 pynvml.nvmlShutdown()
66 # ~30 MB per 512-token sequence (bfloat16, two heads) – clamp sensibly
67 BATCH = max(64, min(int(free // 0.03), 1024))
68 except Exception: # any issue → decent default
69 BATCH = 256
70print(f"[split-binary] batch size = {BATCH}")
71
72# ───────────────────────── THRESHOLDS ────────────────────────────────
73thr_path = Path(MODEL_DIR) / "thresholds.json"
74if thr_path.exists():
75 THR = json.loads(thr_path.read_text())
76 print("Loaded thresholds:", THR)
77else:
78 THR = {"off": 0.5, "def": 0.5}
79 print("No thresholds.json → default 0.5 each")
80
81# ─────────────────── MODEL & TOKENISER LOADING ───────────────────────
82def load_model(path: Path):
83 """Load classification head in BF16 (no flash-attention)."""
84 return HFModel.from_pretrained(path, torch_dtype=torch.bfloat16)
85
86
87paths = {
88 "off": Path(MODEL_DIR) / "offensive_vs_rest",
89 "def": Path(MODEL_DIR) / "defensive_vs_rest",
90}
91print("Loading models …")
92m_off = load_model(paths["off"]).to(DEVICE).eval()
93m_def = load_model(paths["def"]).to(DEVICE).eval()
94
95# Optional: compile graphs for a little extra throughput
96try:
97 m_off = torch.compile(m_off, dynamic=True, mode="reduce-overhead")
98 m_def = torch.compile(m_def, dynamic=True, mode="reduce-overhead")
99 print("torch.compile: dynamic=True, reduce-overhead ✓")
100except Exception:
101 pass
102
103tok = AutoTokenizer.from_pretrained(paths["off"])
104ENC = dict(
105 truncation=True,
106 padding="longest",
107 max_length=512,
108 return_tensors="pt",
109)
110
111# ─────────────────────── OUTPUT HANDLES ──────────────────────────────
112outs = {
113 "off": open("offensive.jsonl", "w", encoding="utf-8"),
114 "def": open("defensive.jsonl", "w", encoding="utf-8"),
115 "oth": open("other.jsonl", "w", encoding="utf-8"),
116}
117
118# ───────────────────────── HELPERS ───────────────────────────────────
119def batched(it, n):
120 """Yield `n`-sized chunks from iterator `it`."""
121 while True:
122 chunk = list(islice(it, n))
123 if not chunk:
124 break
125 yield chunk
126
127
128# ───────────────────── MAIN SPLITTING LOOP ───────────────────────────
129with open(RAW_JSONL, "r", encoding="utf-8") as fin, torch.inference_mode():
130 for lines in tqdm(batched(fin, BATCH), desc="Splitting", ncols=110):
131 recs = [json.loads(l) for l in lines]
132 texts = [r.get("content", "") for r in recs]
133
134 # Tokenise → pin CPU mem → async copy to GPU
135 batch = tok(texts, **ENC)
136 batch = {
137 k: v.pin_memory().to(DEVICE, non_blocking=True) for k, v in batch.items()
138 }
139
140 # Positive-class probabilities
141 p_off = softmax(m_off(**batch).logits, dim=-1)[:, 1].cpu()
142 p_def = softmax(m_def(**batch).logits, dim=-1)[:, 1].cpu()
143
144 for r, po, pd in zip(recs, p_off, p_def):
145 txt = r.get("content", "")
146 off, dfn = po >= THR["off"], pd >= THR["def"]
147
148 if off and not dfn:
149 outs["off"].write(json.dumps({"content": txt}) + "\n")
150 elif dfn and not off:
151 outs["def"].write(json.dumps({"content": txt}) + "\n")
152 elif off and dfn: # tie → higher prob wins
153 (outs["off"] if po >= pd else outs["def"]).write(
154 json.dumps({"content": txt}) + "\n"
155 )
156 else:
157 outs["oth"].write(json.dumps({"content": txt}) + "\n")
158
159# ───────────────────────── CLEAN-UP ──────────────────────────────────
160for f in outs.values():
161 f.close()
162print("✅ Done! → offensive.jsonl defensive.jsonl other.jsonl")
