#!/usr/bin/env python3
# Copyright 2026 zn (usezn.com)
# SPDX-License-Identifier: Apache-2.0
"""Reproduce the public benchmark of usezn/Galvanize-2 against open prompt-injection detectors.

Data: Horizon-Labs/prompt-injection-eval-suite at a pinned revision (every parquet is checked against its SHA-256).
Models: public, ungated Hugging Face checkpoints at pinned revisions (weights SHA-256 checked).
Protocol (primary, "serve"): whitespace flattened with re.sub(r"\\s+", " ", text).strip(), each model's own tokenizer,
first 512 tokens, one window, fp32. Optional "hz" protocol: Horizon's evaluate.py sliding windows, max over windows.
Decontamination: the row ids listed in decontam_excluded_ids.json are removed for the primary ("decontaminated") numbers.

Usage:
    python run_benchmark.py                       # every model, serving protocol
    python run_benchmark.py --protocol both       # also Horizon's windowed protocol
    python run_benchmark.py --models galvanize-2,horizon-small --limit 24 --device cpu   # quick smoke

Outputs: <out-dir>/results.json and <out-dir>/results.md. The script never prints benchmark text.
"""
import argparse
import hashlib
import json
import math
import os
import re
import sys
import time

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))

# ---------------------------------------------------------------------------------------------------------------------
# Eval suite (pinned)
# ---------------------------------------------------------------------------------------------------------------------
SUITE_REPO = 'Horizon-Labs/prompt-injection-eval-suite'
SUITE_REV = '27c7bcd6f46e842f4c4d8b114b30eb6b1d892965'
SUITE_SHA = {
    'notinject': '75d8a7bfde6172335d6178dbfba192209570309a19edfdb11d2120a3b6e919d0',
    'xstest': '3ca278fe64310fa5172d34cd17eef9ef8cbe2d20300756c6424d07d41154b90c',
    'orbench_hard': 'ceaa7b010b08377187e637d2f7c61912914fa0b8f3788cabdcfbc1c16b3c85a8',
    'deepset_test': 'e2c1ed4700967312e2e95e87ae48c0e0617d02b0fed4335b0d9a631b66bb47d6',
    'jackhhao_test': '5a1987583bf94e528f8cfc98f21664eb1206af605f44979317e30a8c4a2b5a87',
    'simsonsun_jailbreaks': '967cdd37e16e5377e476706e4047d2871c930ca48dc829b4001abe76ec72a0dc',
    'boundary_pairs_test': '44e5d543bd4bef91249b91a5d996086631e4cae4d499616bed4202fd4b4b774d',
    'bipia': '744ab24ff443589df2d7ea22fc42618f124f5682487938021f35bdf08c8374bd',
    'piarena': '0bbd59a14b622e0ed8630cac365c136537c2ab2b31e4b3fd9587b5edc0b88b9a',
    'llmail_phase2': '707adf5b915a7f59362a50ccce0652f9b48a145c3a8927762e06c5a777552a9f',
    'neuralchemy_test': 'b8228d74858b2e8d2e4e3c9589423e9758fc4b17e1e5a67a00b4bce812a8bcb9',
    'slabs_test': '67fab0e0cf485f6530f53fdf291a89f27c75e01f674e78ae1a61286642b02f90',
    'agentic5k_test': 'a1adbc4052b57594be26bea9a4064d91bc2291297cf452760ea9938fb7aab14a',
}
# (rows, attack rows) per set at SUITE_REV
EXPECT = {'notinject': (339, 0), 'xstest': (450, 0), 'orbench_hard': (1319, 0), 'deepset_test': (116, 60),
          'jackhhao_test': (262, 139), 'simsonsun_jailbreaks': (2426, 2426), 'boundary_pairs_test': (240, 120),
          'bipia': (3178, 3000), 'piarena': (1500, 750), 'llmail_phase2': (2000, 2000), 'neuralchemy_test': (942, 552),
          'slabs_test': (2101, 1051), 'agentic5k_test': (545, 280)}
# Horizon's metric per model-card row: accuracy for all-benign sets, recall for all-attack sets, F1 otherwise.
CARD_METRIC = {'notinject': 'acc', 'xstest': 'acc', 'orbench_hard': 'acc', 'jackhhao_test': 'f1', 'deepset_test': 'f1',
               'simsonsun_jailbreaks': 'recall', 'boundary_pairs_test': 'f1', 'bipia': 'f1', 'piarena': 'f1',
               'llmail_phase2': 'recall'}
SUPP_METRIC = {'neuralchemy_test': 'f1', 'slabs_test': 'f1', 'agentic5k_test': 'f1'}
METRIC_OF = {**CARD_METRIC, **SUPP_METRIC}
SET_NAMES = list(CARD_METRIC) + list(SUPP_METRIC)
SET_LABEL = {'notinject': 'NotInject (benign, trigger words), accuracy', 'xstest': 'XSTest (benign for injection), accuracy',
             'orbench_hard': 'OR-Bench-hard-1k (benign), accuracy', 'deepset_test': 'deepset/prompt-injections test, F1',
             'jackhhao_test': 'jackhhao/jailbreak-classification test, F1', 'simsonsun_jailbreaks': 'Simsonsun jailbreaks, recall',
             'boundary_pairs_test': 'Agentic boundary pairs test, F1', 'bipia': 'BIPIA (derived), F1',
             'piarena': 'PIArena (derived), F1', 'llmail_phase2': 'LLMail-Inject phase 2, recall',
             'neuralchemy_test': 'neuralchemy core test, F1 (suppl.)', 'slabs_test': 'S-Labs test, F1 (suppl.)',
             'agentic5k_test': 'agentic-prompt-injection-5k test, F1 (suppl.)'}

# ---------------------------------------------------------------------------------------------------------------------
# Models (pinned revision + SHA-256 of model.safetensors; positive label as resolved in the reference run)
# ---------------------------------------------------------------------------------------------------------------------
# Revision of usezn/Galvanize-2 to score. The weights are checked against GALVANIZE2_SHA, the exact file benchmarked.
GALVANIZE2_REVISION = 'b61e50d37c5647788d7b9e301f4e76e246edbfd1'
GALVANIZE2_SHA = '3f474d28f4d39ca3ba95fba760fd5daa4687952de89ee5c1232f66cd99357eee'
GALVANIZE2_THR = 0.661988004118821   # development-set threshold (shown as 0.662)
MODELS = {
    'galvanize-2': {'name': 'Galvanize 2', 'repo': 'usezn/Galvanize-2', 'rev': GALVANIZE2_REVISION, 'sha': GALVANIZE2_SHA,
                    'thr': GALVANIZE2_THR, 'pos_label': 'INJECTION'},
    'horizon-small': {'name': 'Horizon small (141M)', 'repo': 'Horizon-Labs/prompt-injection-guard-small',
                      'rev': 'e0e88c4e940ee7e1b52bdf943d501ec696b83b9b',
                      'sha': 'a23196c117baf414b4b8abd138c5f5656e728877c0bdb2775c361ee89df16deb', 'thr': 0.5, 'pos_label': 'INJECTION'},
    'horizon-base': {'name': 'Horizon base (308M)', 'repo': 'Horizon-Labs/prompt-injection-guard-base',
                     'rev': 'e4fbfc6f603e24c675cdf310c3b8d7abaa825666',
                     'sha': 'd0c0f899f118481855cfe34ba718cf5e44613d266270721c380d54741626d405', 'thr': 0.5, 'pos_label': 'INJECTION'},
    'horizon-large': {'name': 'Horizon large', 'repo': 'Horizon-Labs/prompt-injection-guard-large',
                      'rev': 'bbc67bafa71c791ad802d4df0ae7639d4e880dea',
                      'sha': '81bd637f0964fea52d452ddda3c814cf7d822ca7db6b2a34015db668f4a183bf', 'thr': 0.5, 'pos_label': 'INJECTION'},
    'galvanize-60m': {'name': 'Galvanize-60M', 'repo': 'usezn/Galvanize-60M', 'rev': '7bf8df0d3c9cd724c5773a1c73d7c81ebb0f2082',
                      'sha': '53684cff19857acaa0fd8bb81c9271224d56e5ef1f9b4942a9ddaa846f1e4027', 'thr': 0.95, 'pos_label': 'MALICIOUS'},
    'protectai-v2': {'name': 'ProtectAI v2', 'repo': 'protectai/deberta-v3-base-prompt-injection-v2',
                     'rev': '90c9989b1a342275dd0d1a95aad283c04e075671',
                     'sha': '6521cb8d0ac08148c81464899c424e6148fcc62befa371089fa4061d8b6e0424', 'thr': 0.5, 'pos_label': 'INJECTION'},
    'deepset': {'name': 'deepset', 'repo': 'deepset/deberta-v3-base-injection', 'rev': '80dda00d0b0d9a03917a7685e2ddbcd28e04dbb1',
                'sha': '16862b96671d118a79b8d3a1352a552913797f9fc3cc7db2e06a05002abfc084', 'thr': 0.5, 'pos_label': 'INJECTION'},
    'wolf-defender': {'name': 'Wolf Defender', 'repo': 'patronus-studio/wolf-defender-prompt-injection',
                      'rev': 'a77e7b669f079c3ad7e72e4ad10f3bd173ac4f4d',
                      'sha': 'e8330cc19b083deb45847d873a69bc840cd8913b62a631289e2c034fa5b6cf3b', 'thr': 0.5, 'pos_label': 'injection'},
    'neuraltrust-small': {'name': 'NeuralTrust small', 'repo': 'NeuralTrust/prompt-guard-oss-small',
                          'rev': '40e5c56b68a1b081484c3e56c9bee726e2749138',
                          'sha': '5b1ba9c313447d445146016125ab337a2b5298b2dc2198fcd20e20727a32382f', 'thr': 0.5, 'pos_label': 'jailbreak'},
}
# Galvanize-60M: published head, P(injection) = sigmoid(0.7 * (raw_logit + 17.6063)); its encoder is a 4-layer cut of
# ModernBERT-base (architecture only; every weight comes from the Galvanize-60M checkpoint, loaded strictly).
G60_HEAD = {'scale': 0.7, 'shift': -17.60634142857143}
MBB_REPO, MBB_REV = 'answerdotai/ModernBERT-base', '8949b909ec900327062f0ebf497f51aef5e6f0c8'
NOT_EVALUATED = {
    'leolee99/PIGuard': 'custom architecture that needs trust_remote_code (excluded by policy)',
    'meta-llama/Llama-Prompt-Guard-2-86M': 'gated (manual approval, Llama license)',
    'meta-llama/Llama-Prompt-Guard-2-22M': 'gated (manual approval, Llama license)',
    'qualifire/prompt-injections-benchmark (dataset)': 'gated and non-commercial; not in the public suite',
    'Mindgard/evaded-prompt-injection-and-jailbreak-samples (dataset)': 'gated and non-commercial; not in the public suite',
    'Horizon synthetic_* sets (dataset)': 'generated with the same generator as Horizon\'s training data (in-distribution for '
                                          'Horizon and, through initialisation, for Galvanize 2)',
    'Commercial APIs (e.g. Lakera Guard, Azure Prompt Shields)': 'paid APIs, not open weights',
}
POS_RE = re.compile(r'inject|jailbreak|malicious|unsafe|attack', re.I)

SEQ = 512
TOKEN_BUDGET = 32768
SUBSETS = ('decontaminated', 'full')
PROTOCOL_DESC = {'serve': 're.sub(r"\\s+", " ", text).strip(), model tokenizer, first %d tokens, one window, fp32' % SEQ,
                 'hz': "Horizon's evaluate.py: raw text, window min(2048, max positions) (512 for DeBERTa and Galvanize-60M), "
                       'stride window // 4, document score = max over windows, fp32'}
_WS = re.compile(r'\s+')


# ---------------------------------------------------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------------------------------------------------
def log(msg, **kw):
    print(json.dumps({'t': round(time.time() - T0, 1), 'evt': msg, **kw}, default=float), flush=True)


T0 = time.time()


def sha256_file(path):
    h = hashlib.sha256()
    with open(path, 'rb') as f:
        for chunk in iter(lambda: f.read(1 << 20), b''):
            h.update(chunk)
    return h.hexdigest()


def flat(t):
    # Serving input: identical to the reference run.
    return _WS.sub(' ', str(t)).strip()


def logit(p):
    return math.log(p / (1.0 - p))


def r4(x):
    return None if x is None else round(float(x), 4)


_trap = getattr(np, 'trapezoid', None) or getattr(np, 'trapz')


def pauc(y, z, max_fpr=0.10):
    from sklearn.metrics import roc_curve
    fpr, tpr, _ = roc_curve(y, z)
    stop = int(np.searchsorted(fpr, max_fpr, 'right'))
    x = np.r_[fpr[:stop], max_fpr]
    v = np.r_[tpr[:stop], np.interp(max_fpr, fpr, tpr)]
    return float(_trap(v, x) / max_fpr)


def metrics(y, z, p_thr):
    from sklearn.metrics import roc_auc_score
    y, z = np.asarray(y), np.asarray(z)
    n, pos = len(y), int(y.sum())
    if n == 0:
        return {'n': 0}
    pred = z >= logit(p_thr)
    tp, fp = int((pred & (y == 1)).sum()), int((pred & (y == 0)).sum())
    r = {'n': n, 'pos': pos, 'thr': r4(p_thr), 'acc': r4((pred == (y == 1)).mean())}
    if pos:
        r['recall'] = r4(tp / pos)
    if pos < n:
        r['fpr'] = r4(fp / (n - pos))
    if 0 < pos < n:
        prec = tp / (tp + fp) if tp + fp else 0.0
        rec = tp / pos
        r.update(auroc=r4(roc_auc_score(y, z)), pauc10=r4(pauc(y, z)), precision=r4(prec),
                 f1=r4(2 * prec * rec / (prec + rec) if prec + rec else 0.0))
    return r


# ---------------------------------------------------------------------------------------------------------------------
# Data
# ---------------------------------------------------------------------------------------------------------------------
def load_excluded(path):
    d = json.load(open(path, encoding='utf-8'))
    if d.get('dataset') != SUITE_REPO or d.get('revision') != SUITE_REV:
        raise SystemExit('%s does not belong to %s@%s' % (path, SUITE_REPO, SUITE_REV))
    for s in SET_NAMES:
        if len(d['excluded'][s]) != d['counts'][s]['excluded']:
            raise SystemExit('inconsistent exclusion list for ' + s)
    return {s: np.asarray(d['excluded'][s], dtype=np.int64) for s in SET_NAMES}, d['counts']


def load_suite(limit, excluded):
    import pandas as pd
    from huggingface_hub import hf_hub_download
    sets = {}
    for s in SET_NAMES:
        p = hf_hub_download(SUITE_REPO, 'data/%s.parquet' % s, repo_type='dataset', revision=SUITE_REV)
        if sha256_file(p) != SUITE_SHA[s]:
            raise SystemExit('sha256 mismatch for %s: not the pinned eval-suite file' % s)
        df = pd.read_parquet(p)
        y = df['label'].astype(int).to_numpy()
        if not (set(np.unique(y)) <= {0, 1} and (len(y), int(y.sum())) == EXPECT[s]):
            raise SystemExit('unexpected row or label counts in ' + s)
        text = ['' if t is None else str(t) for t in df['text'].tolist()]
        row = np.arange(len(y))
        if limit:
            # stratified, fixed seed: up to limit // 2 attacks and limit // 2 benign rows per set
            rng = np.random.default_rng(7)
            k = max(1, limit // 2)
            keep = np.sort(np.r_[rng.permutation(np.where(y == 1)[0])[:k], rng.permutation(np.where(y == 0)[0])[:k]])
            text, y, row = [text[i] for i in keep], y[keep], keep
        masks = {'full': np.ones(len(y), bool), 'decontaminated': ~np.isin(row, excluded[s])}
        sets[s] = {'text': text, 'y': y, 'row': row, 'mask': masks}
    return sets


# ---------------------------------------------------------------------------------------------------------------------
# Scoring
# ---------------------------------------------------------------------------------------------------------------------
def windows(tok, texts, proto, window):
    ids, owner = [], []
    for k in range(0, len(texts), 256):
        chunk = texts[k:k + 256]
        if proto == 'serve':
            enc = tok([flat(t) for t in chunk], truncation=True, max_length=SEQ)
            ids += enc['input_ids']
            owner += list(range(k, k + len(chunk)))
        else:
            enc = tok(chunk, truncation=True, max_length=window, stride=window // 4, return_overflowing_tokens=True)
            ids += enc['input_ids']
            owner += [k + int(j) for j in enc['overflow_to_sample_mapping']]
    return ids, np.array(owner)


def score_ids(fwd, ids, pad_id, device):
    import torch
    order = sorted(range(len(ids)), key=lambda i: len(ids[i]))
    out = np.zeros(len(ids), dtype=np.float64)
    k = 0
    with torch.no_grad():
        while k < len(order):
            L = len(ids[order[min(len(order) - 1, k + 255)]])
            bs = max(1, min(256, TOKEN_BUDGET // max(1, L)))
            idx = order[k:k + bs]
            L = max(len(ids[i]) for i in idx)
            ii = torch.full((len(idx), L), pad_id, dtype=torch.long)
            am = torch.zeros((len(idx), L), dtype=torch.long)
            for j, i in enumerate(idx):
                ii[j, :len(ids[i])] = torch.tensor(ids[i])
                am[j, :len(ids[i])] = 1
            out[idx] = fwd(ii.to(device), am.to(device)).double().cpu().numpy()
            k += len(idx)
    return out


def doc_scores(z_win, owner, n):
    z = np.full(n, -np.inf)
    np.maximum.at(z, owner, z_win)
    return z


def check_weights(spec):
    from huggingface_hub import hf_hub_download
    p = hf_hub_download(spec['repo'], 'model.safetensors', revision=spec['rev'])
    got = sha256_file(p)
    if got != spec['sha']:
        raise RuntimeError('model.safetensors sha256 %s... differs from the benchmarked file %s...' % (got[:12], spec['sha'][:12]))
    return p


def hf_classifier(spec, device):
    from transformers import AutoModelForSequenceClassification
    m = None
    for attn in ('sdpa', 'eager'):
        try:
            m = AutoModelForSequenceClassification.from_pretrained(spec['repo'], revision=spec['rev'], attn_implementation=attn)
            break
        except Exception as e:   # DeBERTa-v2 has no SDPA path: falls back to eager, as in the reference run
            log('load_retry', model=spec['repo'], attn=attn, err=type(e).__name__)
    if m is None:
        raise RuntimeError('could not load ' + spec['repo'])
    if hasattr(m.config, 'reference_compile'):
        m.config.reference_compile = False   # no torch.compile inside ModernBERT (variable shapes)
    m = m.float().eval().to(device)
    labels = {int(k): v for k, v in m.config.id2label.items()}
    if len(labels) != 2:
        raise RuntimeError('expected a 2-class head')
    hits = [i for i, v in labels.items() if POS_RE.search(v)]
    pi = hits[0] if hits else max(labels)
    ni = 1 - pi
    if labels[pi] != spec['pos_label']:
        raise RuntimeError('positive label %r differs from the reference run (%r)' % (labels[pi], spec['pos_label']))

    def fwd(ii, am):
        lg = m(input_ids=ii, attention_mask=am).logits.float()
        return lg[:, pi] - lg[:, ni]   # margin z, P(injection) = sigmoid(z)
    mx = int(getattr(m.config, 'max_position_embeddings', 512) or 512)
    window = 512 if 'deberta' in m.config.model_type else min(2048, mx if mx < 100000 else 512)
    return m, fwd, window, {'positive_label': labels[pi], 'model_type': m.config.model_type, 'hz_window': window}


def load_galvanize60m(spec, device):
    # FP32 weights with the loader published in the usezn/Galvanize-60M model card.
    import torch
    from safetensors.torch import load_file
    from transformers import AutoModel, AutoTokenizer

    class G60Pooler(torch.nn.Module):
        def __init__(self, hidden=768, dk=128, nq=4):
            super().__init__()
            self.q_proj = torch.nn.Linear(hidden, dk, bias=False)
            self.k_proj = torch.nn.Linear(hidden, dk, bias=False)
            self.queries = torch.nn.Parameter(torch.zeros(nq, hidden))

        def forward(self, hs, am):
            att = (self.q_proj(self.queries) @ self.k_proj(hs).transpose(1, 2)) / 128 ** 0.5
            att = att.masked_fill(am[:, None, :] == 0, torch.finfo(att.dtype).min)
            return (torch.softmax(att, dim=-1) @ hs).reshape(hs.size(0), -1)

    class G60(torch.nn.Module):
        def __init__(self, encoder):
            super().__init__()
            self.encoder, self.pooler = encoder, G60Pooler()
            self.classifier = torch.nn.Sequential(torch.nn.Identity(), torch.nn.Linear(3072, 384), torch.nn.GELU(),
                                                  torch.nn.Linear(384, 1))

        def forward(self, input_ids, attention_mask):
            hs = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
            return self.classifier(self.pooler(hs, attention_mask)).squeeze(-1)

    w = check_weights(spec)
    enc = AutoModel.from_pretrained(MBB_REPO, revision=MBB_REV, attn_implementation='eager')
    enc.layers = torch.nn.ModuleList([enc.layers[i] for i in (0, 7, 14, 21)])
    enc.config.num_hidden_layers = 4
    enc.layers[0].attn_norm = torch.nn.Identity()
    if hasattr(enc.config, 'reference_compile'):
        enc.config.reference_compile = False
    m = G60(enc)
    m.load_state_dict(load_file(w), strict=True)
    m = m.float().eval().to(device)
    tok = AutoTokenizer.from_pretrained(spec['repo'], revision=spec['rev'])

    def fwd(ii, am):
        return G60_HEAD['scale'] * (m(ii, am).float() - G60_HEAD['shift'])
    return m, tok, fwd, 512, {'positive_label': 'MALICIOUS', 'model_type': 'galvanize-60m', 'hz_window': 512}


def load_model(key, device):
    from transformers import AutoTokenizer
    spec = MODELS[key]
    if key == 'galvanize-60m':
        return load_galvanize60m(spec, device)
    check_weights(spec)
    tok = AutoTokenizer.from_pretrained(spec['repo'], revision=spec['rev'])
    m, fwd, window, info = hf_classifier(spec, device)
    return m, tok, fwd, window, info


def tokenizer_guard(tok, sets):
    # Galvanize 2 serving contract: its tokenizer must produce the same ids as Horizon small's (the reference run checked this).
    from transformers import AutoTokenizer
    ref = MODELS['horizon-small']
    ref_tok = AutoTokenizer.from_pretrained(ref['repo'], revision=ref['rev'])
    texts = [t for s in SET_NAMES for t in sets[s]['text'][:3]][:16]
    a, b = windows(tok, texts, 'serve', SEQ)[0], windows(ref_tok, texts, 'serve', SEQ)[0]
    if a != b:
        raise RuntimeError('Galvanize 2 tokenizer ids differ from the reference serving tokenizer')


# ---------------------------------------------------------------------------------------------------------------------
# Aggregation and report
# ---------------------------------------------------------------------------------------------------------------------
def macro(res_kp, sub, tk):
    vals = [res_kp[s][sub][tk].get(CARD_METRIC[s]) for s in CARD_METRIC if s in res_kp]
    vals = [v for v in vals if v is not None]
    return {'value': r4(np.mean(vals)) if vals else None, 'n_sets': len(vals)}


def repro_vs_horizon(res, path):
    if not os.path.exists(path):
        return None
    pub = json.load(open(path, encoding='utf-8'))
    out = {}
    for key, by_set in pub['models'].items():
        if key not in res or 'hz' not in res[key]:
            continue
        d_acc, d_auc = {}, {}
        for s, ref in by_set.items():
            if s in res[key]['hz']:
                ours = res[key]['hz'][s]['full']['p05']
                if ref.get('acc') is not None:
                    d_acc[s] = r4(ours['acc'] - ref['acc'])
                if ref.get('auc') is not None and ours.get('auroc') is not None:
                    d_auc[s] = r4(ours['auroc'] - ref['auc'])
        out[key] = {'delta_acc': d_acc, 'delta_auroc': d_auc,
                    'max_abs_delta_acc': r4(max(map(abs, d_acc.values()))) if d_acc else None,
                    'max_abs_delta_auroc': r4(max(map(abs, d_auc.values()))) if d_auc else None}
    return out


def compare_with_reference(res, path):
    if not os.path.exists(path):
        return None
    ref = json.load(open(path, encoding='utf-8')).get('results', {})
    out = {}
    for key, by_proto in res.items():
        deltas = []
        for proto, by_set in by_proto.items():
            for s, by_sub in by_set.items():
                for sub, by_tk in by_sub.items():
                    for tk, m in by_tk.items():
                        r = ref.get(key, {}).get(proto, {}).get(s, {}).get(sub, {}).get(tk)
                        if not r:
                            continue
                        for f in ('acc', 'recall', 'fpr', 'f1', 'auroc', 'pauc10'):
                            if m.get(f) is not None and r.get(f) is not None:
                                deltas.append(abs(m[f] - r[f]))
        out[key] = {'compared_values': len(deltas), 'max_abs_delta': r4(max(deltas)) if deltas else None}
    return out


def _thr_txt(t):
    return ('%.3f' % t).rstrip('0').rstrip('.')


def render_markdown(doc):
    models, res = doc['models'], doc['results']
    first = [k for k in models if k in res]
    cols = []
    for k in first:
        cols.append((k, 'own'))
        if k == 'galvanize-2':
            cols.append((k, 'p05'))
    protos = sorted({p for k in res for p in res[k]}, key=lambda p: 0 if p == 'serve' else 1)

    def col_name(k, tk, thr_free=False):
        return models[k]['name'] if thr_free else '%s @%s' % (models[k]['name'], _thr_txt(models[k]['threshold']) if tk == 'own' else '0.5')

    def fmt(v):
        return '-' if v is None else '%.3f' % v

    def n_rows(proto, s, sub):
        for k in first:
            if s in res[k].get(proto, {}):
                return res[k][proto][s][sub]['own'].get('n', 0)
        return 0

    def table(proto, sub, what, title):
        thr_free = what != 'card'
        cc = [(k, t) for k, t in cols if not (thr_free and t == 'p05')]
        sets = [s for s in SET_NAMES if any(s in res[k].get(proto, {}) for k, _ in cc)]
        if thr_free:
            sets = [s for s in sets if 0 < EXPECT[s][1] < EXPECT[s][0]]
        if not sets:
            return ''
        lines = ['**%s**' % title, '', '| Eval set | n | ' + ' | '.join(col_name(k, t, thr_free) for k, t in cc) + ' |',
                 '|---|---:|' + '---|' * len(cc)]
        for s in sets:
            cells = []
            for k, t in cc:
                m = res[k].get(proto, {}).get(s, {}).get(sub, {}).get(t)
                if not m:
                    cells.append(('-', None))
                    continue
                if what == 'card':
                    v = m.get(METRIC_OF[s])
                    extra = ' (%.2f)' % m['fpr'] if METRIC_OF[s] == 'f1' and m.get('fpr') is not None else ''
                    cells.append((fmt(v) + extra, v))
                else:
                    cells.append((fmt(m.get(what)), m.get(what)))
            vals = [v for _, v in cells if v is not None]
            best = max(vals) if vals else None
            txt = [('**%s**' % c if v is not None and v == best else c) for c, v in cells]
            lines.append('| %s | %d | %s |' % (SET_LABEL[s] if what == 'card' else s, n_rows(proto, s, sub), ' | '.join(txt)))
        if what == 'card':
            mac = doc['macro_card10']
            lines.append('| Macro average, 10 card sets | | %s |' % ' | '.join(
                fmt(mac.get(k, {}).get(proto, {}).get(sub, {}).get(t, {}).get('value')) for k, t in cc))
        return '\n'.join(lines) + '\n'

    run = doc['run']
    md = ['# Galvanize 2 vs open prompt-injection detectors', '',
          'Run %s (UTC) on %s. Data: [%s](https://huggingface.co/datasets/%s) @ `%s`. Primary numbers are on the '
          '**decontaminated** subset (rows listed in `decontam_excluded_ids.json` removed) under the **serving protocol** '
          '(whitespace flattened, first 512 tokens, one window). F1 cells show the false-positive rate on the benign rows in '
          'parentheses. Best value per row in bold.' % (run['date_utc'], run['hardware'], doc['suite']['repo'], doc['suite']['repo'],
                                                       doc['suite']['revision'][:12]), '']
    if run.get('limit'):
        md += ['> SMOKE RUN (`--limit %d`): a few rows per set. The numbers are not meaningful; this only checks the code path.'
               % run['limit'], '']
    if 'serve' in protos:
        md.append(table('serve', 'decontaminated', 'card', 'Serving protocol, decontaminated subset, each model at its own threshold'))
        md.append(table('serve', 'decontaminated', 'auroc', 'Threshold-free: ROC AUC (serving protocol, decontaminated)'))
        md.append(table('serve', 'decontaminated', 'pauc10',
                        'Threshold-free: normalised partial AUC up to 10% FPR (serving protocol, decontaminated)'))
        md.append(table('serve', 'full', 'card', 'Serving protocol, full sets (secondary; includes the excluded rows)'))
    if 'hz' in protos:
        md.append(table('hz', 'full', 'card', "Horizon's protocol (raw text, sliding windows, max over windows), full sets"))
        md.append(table('hz', 'decontaminated', 'card', "Horizon's protocol, decontaminated subset"))
    rep = doc.get('reproduction_vs_horizon_published')
    if rep:
        md += ["**Reproduction check** (Horizon's protocol on the full sets vs the numbers published in Horizon's repository, "
               'threshold 0.5; Horizon used bf16 autocast, this benchmark runs fp32):', '',
               '| Model | max abs delta accuracy | max abs delta ROC AUC |', '|---|---:|---:|']
        md += ['| %s | %s | %s |' % (models[k]['name'], fmt(v['max_abs_delta_acc']), fmt(v['max_abs_delta_auroc'])) for k, v in rep.items()]
        md.append('')
    dc = doc['decontamination']
    has_pct = any('training_overlap_any_pct' in c for c in dc.values())
    md += ['**Decontamination** (rows removed for the decontaminated subset: their text matched a row of the Galvanize 2 '
           'training data; see README):', '']
    if has_pct:
        md += ['| Eval set | n | exact match % | normalised match % | any match % | excluded | kept | upstream train-split overlap % |',
               '|---|---:|---:|---:|---:|---:|---:|---:|']
        for s, c in dc.items():
            md.append('| %s | %d | %.2f | %.2f | %.2f | %d | %d | %.2f |' % (
                s, c['rows'], c['training_overlap_exact_pct'], c['training_overlap_normalised_pct'],
                c['training_overlap_any_pct'], c['excluded'], c['kept'], c['upstream_train_split_overlap_any_pct']))
        md += ['', 'Galvanize 2 is fine-tuned from Horizon small, so any overlap between the eval sets and Horizon\'s training '
               'data is inherited. Horizon\'s training set is not published; the last column is the overlap with the upstream '
               '*train* splits of the sources Horizon lists (deepset, jackhhao, neuralchemy core, S-Labs, '
               'agentic-prompt-injection-5k, or-bench-80k), an upper bound since Horizon states it deduplicated against every '
               'eval set. It is informational and is not used to filter rows.', '']
    else:
        md += ['| Eval set | n | excluded | kept |', '|---|---:|---:|---:|']
        md += ['| %s | %d | %d | %d |' % (s, c['rows'], c['excluded'], c['kept']) for s, c in dc.items()]
        md.append('')
    cmp_ = doc.get('comparison_with_reference')
    if cmp_:
        md += ['**Comparison with the published reference** (`results.json` shipped with this folder; max abs delta over '
               'accuracy, recall, FPR, F1, ROC AUC and pAUC):', '', '| Model | values compared | max abs delta |', '|---|---:|---:|']
        md += ['| %s | %d | %s |' % (models[k]['name'], v['compared_values'], fmt(v['max_abs_delta'])) for k, v in cmp_.items()]
        md.append('')
    md += ['## Method', '',
           '- Hardware and software: %s; torch %s, transformers %s, Python %s.' % (run['hardware'], run['torch'], run['transformers'], run['python']),
           '- Data: `%s` @ `%s`; every parquet checked against its SHA-256. Labels as published (1 = injection or jailbreak, '
           '0 = benign; plain harmful requests are benign under this suite\'s definition). Horizon\'s synthetic sets are not used.'
           % (doc['suite']['repo'], doc['suite']['revision']),
           '- Models (revision, weights SHA-256 checked): ' + '; '.join(
               '%s `%s@%s`' % (v['name'], v['repo'], str(v['revision'])[:12]) if v.get('revision') else
               '%s `%s` (model.safetensors SHA-256 `%s`)' % (v['name'], v['repo'], v['weights_sha256'][:12])
               for v in models.values()) + '.',
           '- Thresholds: Galvanize 2 %s (its development-set threshold, also shown at 0.5); Galvanize-60M 0.95 (its published '
           'calibrated threshold, on sigmoid(0.7 * (logit + 17.6063))); every other model 0.5 (their default).'
           % _thr_txt(GALVANIZE2_THR),
           '- Serving protocol: %s. Horizon protocol: %s.' % (PROTOCOL_DESC['serve'], PROTOCOL_DESC['hz']),
           '- Metrics: accuracy for all-benign sets, recall for all-attack sets, F1 (with FPR) otherwise, as in Horizon\'s card; '
           'ROC AUC and normalised pAUC up to 10% FPR are threshold-free. Macro average = mean over the 10 card sets (Horizon\'s '
           'card macro also includes Qualifire and the two Mindgard sets, which are gated).', '', '**Not evaluated**', '']
    md += ['- `%s`: %s' % (k, v) for k, v in doc['not_evaluated'].items()]
    md += ['- `%s`: failed in this run (%s)' % (k, v['error']) for k, v in models.items() if v.get('error')]
    return '\n'.join(md) + '\n'


# ---------------------------------------------------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser(description=__doc__.split('\n')[0])
    ap.add_argument('--models', default=','.join(MODELS), help='comma-separated keys: ' + ', '.join(MODELS))
    ap.add_argument('--device', choices=('cpu', 'cuda'), default=None, help='default: cuda if available, else cpu')
    ap.add_argument('--limit', type=int, default=0, help='smoke test: at most LIMIT//2 attack + LIMIT//2 benign rows per set')
    ap.add_argument('--protocol', choices=('serve', 'hz', 'both'), default='serve', help='default: serve (the primary numbers)')
    ap.add_argument('--out-dir', default=os.path.join(HERE, 'run_output'))
    ap.add_argument('--excluded', default=os.path.join(HERE, 'decontam_excluded_ids.json'))
    ap.add_argument('--reference', default=os.path.join(HERE, 'results.json'), help='published results to compare against')
    ap.add_argument('--galvanize2-revision', default=GALVANIZE2_REVISION, help='revision of usezn/Galvanize-2 to load')
    a = ap.parse_args()

    keys = [k.strip() for k in a.models.split(',') if k.strip()]
    unknown = [k for k in keys if k not in MODELS]
    if unknown:
        raise SystemExit('unknown model keys %s; choose from %s' % (unknown, ', '.join(MODELS)))
    keys = [k for k in MODELS if k in keys]
    MODELS['galvanize-2']['rev'] = a.galvanize2_revision
    protos = ('serve', 'hz') if a.protocol == 'both' else (a.protocol,)

    import torch
    import transformers
    v = tuple(int(x) for x in re.findall(r'\d+', transformers.__version__)[:2])
    if v < (4, 48):
        raise SystemExit('ModernBERT needs transformers >= 4.48 (found %s)' % transformers.__version__)
    device = a.device or ('cuda' if torch.cuda.is_available() else 'cpu')
    if device == 'cuda' and not torch.cuda.is_available():
        raise SystemExit('--device cuda requested but CUDA is not available')
    torch.backends.cuda.matmul.allow_tf32 = False   # strict fp32
    torch.backends.cudnn.allow_tf32 = False
    np.random.seed(42)
    torch.manual_seed(42)
    gpu = torch.cuda.get_device_name(0) if device == 'cuda' else None
    hardware = ('1x %s GPU, fp32, TF32 off' % gpu) if gpu else 'CPU, fp32'
    log('env', device=device, gpu=gpu, torch=torch.__version__, transformers=transformers.__version__, models=keys,
        protocols=list(protos), limit=a.limit)

    excluded, counts = load_excluded(a.excluded)
    sets = load_suite(a.limit, excluded)
    log('suite', revision=SUITE_REV, sets={s: [len(d['y']), int(d['y'].sum()), int(d['mask']['decontaminated'].sum())]
                                           for s, d in sets.items()})

    scores, models_out = {}, {}
    for key in keys:
        spec = MODELS[key]
        models_out[key] = {'name': spec['name'], 'repo': spec['repo'], 'revision': spec['rev'], 'weights_sha256': spec['sha'],
                           'threshold': spec['thr'], 'positive_label': spec['pos_label'], 'error': None}
        t_load = time.time()
        try:
            m, tok, fwd, window, info = load_model(key, device)
            if key == 'galvanize-2':
                tokenizer_guard(tok, sets)
            models_out[key].update(info, load_s=round(time.time() - t_load, 1))
            scores[key] = {}
            for proto in protos:
                t1 = time.time()
                scores[key][proto] = {}
                for s in SET_NAMES:
                    ids, owner = windows(tok, sets[s]['text'], proto, window)
                    zw = score_ids(fwd, ids, tok.pad_token_id or 0, device)
                    # float32, as stored and scored in the reference run
                    scores[key][proto][s] = doc_scores(zw, owner, len(sets[s]['y'])).astype(np.float32)
                if device == 'cuda':
                    torch.cuda.synchronize()
                log('scored', model=key, protocol=proto, seconds=round(time.time() - t1, 1))
            del m
        except Exception as e:
            models_out[key]['error'] = '%s: %s' % (type(e).__name__, str(e)[:200])
            scores.pop(key, None)
            log('model_failed', model=key, err=models_out[key]['error'])
        __import__('gc').collect()
        if device == 'cuda':
            torch.cuda.empty_cache()

    res = {}
    for key, by_proto in scores.items():
        res[key] = {}
        for proto, by_set in by_proto.items():
            res[key][proto] = {}
            for s, z in by_set.items():
                y = sets[s]['y']
                res[key][proto][s] = {sub: {'own': metrics(y[sets[s]['mask'][sub]], z[sets[s]['mask'][sub]], MODELS[key]['thr']),
                                            'p05': metrics(y[sets[s]['mask'][sub]], z[sets[s]['mask'][sub]], 0.5)}
                                      for sub in SUBSETS}
    mac = {k: {p: {sub: {tk: macro(res[k][p], sub, tk) for tk in ('own', 'p05')} for sub in SUBSETS} for p in res[k]} for k in res}
    doc = {
        'benchmark': 'Galvanize 2 vs open prompt-injection detectors (public eval suite)',
        'run': {'date_utc': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()), 'hardware': hardware, 'device': device,
                'torch': torch.__version__, 'transformers': transformers.__version__, 'python': sys.version.split()[0],
                'protocols': list(protos), 'limit': a.limit or None},
        'suite': {'repo': SUITE_REPO, 'revision': SUITE_REV, 'sha256': SUITE_SHA, 'card_metric': CARD_METRIC,
                  'supplementary_metric': SUPP_METRIC},
        'protocols': {p: PROTOCOL_DESC[p] for p in protos},
        'subsets': {'decontaminated': 'rows listed in decontam_excluded_ids.json removed', 'full': 'every row'},
        'models': models_out,
        'decontamination': {s: {'rows': counts[s]['rows'], 'excluded': counts[s]['excluded'], 'kept': counts[s]['kept']} for s in SET_NAMES},
        'results': res,
        'macro_card10': mac,
        'not_evaluated': NOT_EVALUATED,
    }
    if 'hz' in protos and not a.limit:
        doc['reproduction_vs_horizon_published'] = repro_vs_horizon(res, os.path.join(HERE, 'horizon_published_metrics.json'))
    if not a.limit and os.path.abspath(a.reference) != os.path.abspath(os.path.join(a.out_dir, 'results.json')):
        doc['comparison_with_reference'] = compare_with_reference(res, a.reference)
    os.makedirs(a.out_dir, exist_ok=True)
    with open(os.path.join(a.out_dir, 'results.json'), 'w', encoding='utf-8') as f:
        json.dump(doc, f, indent=1, default=float)
    with open(os.path.join(a.out_dir, 'results.md'), 'w', encoding='utf-8') as f:
        f.write(render_markdown(doc))
    log('done', out_dir=a.out_dir, evaluated=list(res), failed=[k for k, v in models_out.items() if v['error']],
        macro_serve_decontaminated={k: mac[k].get('serve', {}).get('decontaminated', {}).get('own', {}).get('value') for k in mac})


if __name__ == '__main__':
    main()