Update README.md

README.md

@@ -72,57 +72,165 @@ Detects **technical red-team / offensive security** text (English).
 ## Quick start
 
 ```python
-import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-
-def classify_texts(model_name, phrases, threshold=0.515):
-    """
-    Returns a list of (probability_offensive, label) tuples for each phrase
-    given a model_name and threshold.
-    """
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSequenceClassification.from_pretrained(model_name)
-    model.eval()
-
-    inputs = tokenizer(phrases, padding=True, truncation=True, return_tensors="pt")
-
-    with torch.no_grad():
-        logits = model(**inputs).logits  # shape: (batch_size, 2)
-        probs_offensive = torch.softmax(logits, dim=1)[:, 1]  # Probability of the "Offensive" class
-
-    results = []
-    for p_val in probs_offensive:
-        p_val = p_val.item()
-        label = "Offensive (red-team)" if p_val >= threshold else "Not Offensive"
-        results.append((p_val, label))
-    return results
-
-def main():
-    # Example phrases: Offensive (red-team), Defensive (blue-team), Non-technical
-    phrases = [
-        # 1) Cybersecurity Offensive / red-team
-        "To exfiltrate sensitive data, launch a phishing campaign that tricks employees into revealing their VPN credentials.",
-        # 2) Cybersecurity Defensive / blue-team
-        "We should deploy an EDR solution, monitor all endpoints for intrusion attempts, and enforce strict password policies.",
-        # 5) Cybersecruity Marketing
-        "“Our marketing team will unveil the new cybersecurity branding materials at next Tuesday’s antivirus product launch",
-        # 5) Non Cybersecruity  related
-        "I'm excited about the company picnic. There's no cybersecurity topic—just burgers and games."
-    ]
-
-    # Classify with both models
-    threshold = 0.515
-    blue_results = classify_texts("HagalazAI/BlueSecureBERT", phrases, threshold)
-    red_results = classify_texts("HagalazAI/RedSecureBERT", phrases, threshold)
-
-    # Print a Markdown table
-    print("| # | Phrase | Blue Score | Blue Label | Red Score | Red Label |")
-    print("|---|--------|-----------|-----------|----------|----------|")
-    for i, text in enumerate(phrases, start=1):
-        blue_score, blue_label = blue_results[i - 1]
-        red_score, red_label = red_results[i - 1]
-        print(f"| {i} | {text} | {blue_score:.3f} | {blue_label} | {red_score:.3f} | {red_label} |")
-
-if __name__ == "__main__":
-    main()
+#!/usr/bin/env python
+"""
+06_split_binary.py
+~~~~~~~~~~~~~~~~~~
+
+Stream-splits a JSONL cybersecurity corpus into *offensive*, *defensive*, and *other* shards
+using **two** fine-tuned SecureBERT heads.
+
+How the two heads work together
+-------------------------------
+We load two independent checkpoints:
+
+* `offensive_vs_rest` → gives **P(offensive | text)**
+* `defensive_vs_rest` → gives **P(defensive | text)**
+
+For every line we:
 
+1. run both heads in the same GPU batch;
+2. take the positive-class probability from each soft-max;
+3. compare against per-head thresholds (from `thresholds.json`, default 0.5);
+4. route the text with this truth table
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from itertools import islice
+from pathlib import Path
+
+import torch
+from torch.nn.functional import softmax
+from tqdm.auto import tqdm
+from transformers import (
+    AutoModelForSequenceClassification as HFModel,
+    AutoTokenizer,
+)
+
+from config import RAW_JSONL, MODEL_DIR  # MODEL_DIR == securebert_finetuned
+
+# ─────────────────────────────  GPU SETTINGS  ──────────────────────────
+# 1. Use TensorFloat-32 on Ada GPUs (gives a big matmul speed boost).
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.set_float32_matmul_precision("medium")
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ────────────────────────────────  CLI  ────────────────────────────────
+cli = argparse.ArgumentParser(description="Split JSONL into offence/defence/other")
+cli.add_argument("--batch_size", type=int, help="override auto batch sizing")
+args = cli.parse_args()
+
+# ─────────────────────  BATCH-SIZE HEURISTIC  ──────────────────────────
+if args.batch_size:  # user override wins
+    BATCH = args.batch_size
+else:
+    try:
+        import pynvml
+
+        pynvml.nvmlInit()
+        free = (
+            pynvml.nvmlDeviceGetMemoryInfo(pynvml.nvmlDeviceGetHandleByIndex(0)).free
+            / 1024**3
+        )
+        pynvml.nvmlShutdown()
+        # ~30 MB per 512-token sequence (bfloat16, two heads) – clamp sensibly
+        BATCH = max(64, min(int(free // 0.03), 1024))
+    except Exception:  # any issue → decent default
+        BATCH = 256
+print(f"[split-binary] batch size = {BATCH}")
+
+# ─────────────────────────  THRESHOLDS  ────────────────────────────────
+thr_path = Path(MODEL_DIR) / "thresholds.json"
+if thr_path.exists():
+    THR = json.loads(thr_path.read_text())
+    print("Loaded thresholds:", THR)
+else:
+    THR = {"off": 0.5, "def": 0.5}
+    print("No thresholds.json → default 0.5 each")
+
+# ───────────────────  MODEL & TOKENISER LOADING  ───────────────────────
+def load_model(path: Path):
+    """Load classification head in BF16 (no flash-attention)."""
+    return HFModel.from_pretrained(path, torch_dtype=torch.bfloat16)
+
+
+paths = {
+    "off": Path(MODEL_DIR) / "offensive_vs_rest",
+    "def": Path(MODEL_DIR) / "defensive_vs_rest",
+}
+print("Loading models …")
+m_off = load_model(paths["off"]).to(DEVICE).eval()
+m_def = load_model(paths["def"]).to(DEVICE).eval()
+
+# Optional: compile graphs for a little extra throughput
+try:
+    m_off = torch.compile(m_off, dynamic=True, mode="reduce-overhead")
+    m_def = torch.compile(m_def, dynamic=True, mode="reduce-overhead")
+    print("torch.compile: dynamic=True, reduce-overhead ✓")
+except Exception:
+    pass
+
+tok = AutoTokenizer.from_pretrained(paths["off"])
+ENC = dict(
+    truncation=True,
+    padding="longest",
+    max_length=512,
+    return_tensors="pt",
+)
+
+# ───────────────────────  OUTPUT HANDLES  ──────────────────────────────
+outs = {
+    "off": open("offensive.jsonl", "w", encoding="utf-8"),
+    "def": open("defensive.jsonl", "w", encoding="utf-8"),
+    "oth": open("other.jsonl", "w", encoding="utf-8"),
+}
+
+# ─────────────────────────  HELPERS  ───────────────────────────────────
+def batched(it, n):
+    """Yield `n`-sized chunks from iterator `it`."""
+    while True:
+        chunk = list(islice(it, n))
+        if not chunk:
+            break
+        yield chunk
+
+
+# ─────────────────────  MAIN SPLITTING LOOP  ───────────────────────────
+with open(RAW_JSONL, "r", encoding="utf-8") as fin, torch.inference_mode():
+    for lines in tqdm(batched(fin, BATCH), desc="Splitting", ncols=110):
+        recs = [json.loads(l) for l in lines]
+        texts = [r.get("content", "") for r in recs]
+
+        # Tokenise → pin CPU mem → async copy to GPU
+        batch = tok(texts, **ENC)
+        batch = {
+            k: v.pin_memory().to(DEVICE, non_blocking=True) for k, v in batch.items()
+        }
+
+        # Positive-class probabilities
+        p_off = softmax(m_off(**batch).logits, dim=-1)[:, 1].cpu()
+        p_def = softmax(m_def(**batch).logits, dim=-1)[:, 1].cpu()
+
+        for r, po, pd in zip(recs, p_off, p_def):
+            txt = r.get("content", "")
+            off, dfn = po >= THR["off"], pd >= THR["def"]
+
+            if off and not dfn:
+                outs["off"].write(json.dumps({"content": txt}) + "\n")
+            elif dfn and not off:
+                outs["def"].write(json.dumps({"content": txt}) + "\n")
+            elif off and dfn:  # tie → higher prob wins
+                (outs["off"] if po >= pd else outs["def"]).write(
+                    json.dumps({"content": txt}) + "\n"
+                )
+            else:
+                outs["oth"].write(json.dumps({"content": txt}) + "\n")
+
+# ─────────────────────────  CLEAN-UP  ──────────────────────────────────
+for f in outs.values():
+    f.close()
+print("✅ Done! → offensive.jsonl  defensive.jsonl  other.jsonl")