Add full mode (paper's end-to-end configuration) and its per-seed thresholds; fix model card; remove __pycache__

README.md

@@ -21,9 +21,8 @@ pipeline_tag: text-classification
 language model and the replies the model produces, and it is trained on data generated from a written constitution of
 eleven harm categories.
 
-Paper: *EuroSafeAI Open Constitutional Classifiers* (arXiv: TODO) · Code: TODO · Data:
-[EuroSafeGuard-prompts](https://huggingface.co/datasets/EuroSafeAI/EuroSafeGuard-prompts),
-[EuroSafeGuard-replies](https://huggingface.co/datasets/EuroSafeAI/EuroSafeGuard-replies) (gated)
+Paper: *EuroSafeGuard: Open Constitutional Classifiers* (arXiv: coming soon) · Code: coming soon ·
+Training prompts: coming soon (gated access) · Project page: [guard.safe.eu](https://guard.safe.eu)
 
 ## Components
 
@@ -32,8 +31,9 @@ Paper: *EuroSafeAI Open Constitutional Classifiers* (arXiv: TODO) · Code: TODO
 | `probe/` | MLP probe (hidden size 256) on the activations of all 32 blocks of the protected model, reading the chat-formatted prompt; per-token logits are smoothed with an exponential moving average and the prompt score is their maximum | Llama-3.1-8B-Instruct |
 | `gate/seed{0..4}/` | Input classifier: LoRA adapter + one sigmoid head per harm category (11 heads); score = highest head | Qwen3.5-9B |
 | `exchange/` | Exchange classifier: reads `User: <prompt>\n\nAssistant: <reply>`; LoRA + 11 heads | Qwen3.5-9B |
-| `thresholds/seed{k}.json` | Score normalisation and thresholds for each input-classifier seed | |
-| `eurosafeguard.py` | Single-file loader | |
+| `thresholds/seed{k}.json` | Alert mode: score normalisation and thresholds for each input-classifier seed | |
+| `thresholds/full_seed{k}.json` | Full mode: normalisation (probe, gate, exchange), probe escalation cut, weights and threshold for each seed | |
+| `eurosafeguard.py` | Single-file loader with both modes | |
 
 The five `gate/seed*` folders are the five training runs (same data, different seeds) that the paper reports as a mean.
 They vary, especially against search attacks such as PAIR, so evaluate the seed you deploy.
@@ -42,7 +42,14 @@ They vary, especially against search attacks such as PAIR, so evaluate the seed
 Integrity, Physical Harm, Privacy, Self-Harm, Sexual Content. `taxonomy_11heads_short.json` lists what each category
 allows and disallows.
 
-## How it decides (alert-unless-CBRN)
+## Two modes
+
+| Mode | Load with | Paper configuration | Paper tables |
+|---|---|---|---|
+| `alert` (default) | `mode="alert"` | EuroSafeGuard (Input) at the 5% WildChat-clean budget, plus alert-unless-CBRN and the exchange check | Held-out attacks (Input columns), alert mode ("Alert, strict exchange") |
+| `full` | `mode="full"` | EuroSafeGuard (Full): probe escalation, weighted sum of probe, input and exchange scores | End-to-end table |
+
+### Alert mode (input/alert, alert-unless-CBRN)
 
 1. **Input stage.** The input score is the mean of the probe and input-classifier scores, each z-scored on
    WildChat-clean. Its threshold (`alert_input_thr`) flags 5% of WildChat-clean prompts.
@@ -64,6 +71,28 @@ if decision == "refuse" or g.check_reply(prompt, reply, decision):
     reply = "Sorry, I can't help with that."
 ```
 
+### Full mode (the paper's end-to-end configuration)
+
+1. **Probe screen.** The probe scores every request. Only requests whose probe z-score exceeds `probe_escalation_z`
+   (1.1644, the 90th percentile on WildChat-clean) are escalated; the rest are never blocked.
+2. **Escalated requests** are scored by the input classifier and, once the protected model has replied, by the exchange
+   classifier. The exchange is blocked if `0.6·probe_z + 0.2·input_z + 0.2·exchange_z > full_thr`.
+3. All three scores are z-scored on WildChat-clean (the exchange on Llama-3-8B-Instruct replies: μ = 0.02560,
+   σ = 0.09032). `full_thr` is the smallest threshold that keeps every benign calibration set under its cap (WildChat-clean,
+   MMLU, GSM8K, Alpaca, Dolly, OASST 0.5%; XSTest-safe, OR-Bench-80k 2%; OR-Bench-hard 5%; benign rows of styled
+   jailbreak sets 40%): 3.382 / 3.571 / 3.132 / 3.097 / 3.318 for seeds 0–4. It was fitted once on the undefended
+   Llama-3-8B-Instruct's traffic.
+
+```python
+g = EuroSafeGuard.from_pretrained("EuroSafeAI/EuroSafeGuard", seed=0, mode="full", devices=("cuda:0", "cuda:0", "cuda:1"))
+reply = ...                                        # your protected model
+if g.check_full(prompt, reply):                    # probe on every request; gate + exchange only if escalated
+    reply = "Sorry, I can't help with that."
+```
+
+`full_block(probe, gate, exchange, cfg)` applies the same rule to precomputed raw scores. On the paper's stored scores it
+reproduces the end-to-end attack success of every seed exactly (mean 3.88%).
+
 **Requirements:** `torch`, `transformers` with Qwen3.5 support, `peft`, `safetensors`, `huggingface_hub`. About 16 GB of
 GPU memory for Llama-3.1-8B (the probe) plus about 19 GB for each Qwen3.5-9B classifier (input and exchange). Use
 `load_exchange=False` to run the input stage only. The probe reads Llama-3.1-8B-Instruct activations. If you protect
@@ -71,13 +100,18 @@ Llama-3.1-8B-Instruct itself, you can reuse the activations of its forward pass.
 
 ## Results (from the paper; every guard calibrated to the same false-positive budget on real chat traffic)
 
-- **BELLS-O:** our input classifiers detect more harmful prompts than any supervisor on the leaderboard at a
-  false-positive rate near 1%, and have the highest recall on its held-out jailbreak sets.
+- **End to end** (Full mode, Llama-3-8B-Instruct, 1,974 Jailbreak Foundry attacks on held-out GuidedBench behaviors):
+  attack success falls from 7.95% to 3.88% (EuroSafe-CC++, our reimplementation of Constitutional Classifiers++, 4.57%),
+  adding at most 3.2 points of benign failure on any of six benign benchmarks. On gpt-oss-20b: 13.17% → 5.28%.
+- **Compute:** only the 10% of requests the probe escalates reach the two classifiers, so the guardrail adds 15% of the
+  protected model's FLOPs.
+- **Held-out jailbreaks**, input stage at the 5% WildChat-clean budget: 95.9% recall on the BELLS-O jailbreak set
+  (88.2% for HaloGuard, the strongest open guard); highest recall on three of four held-out jailbreak sets.
 - **PAIR** (iterative black-box attack), input stage at the 5% WildChat-clean budget: 21% of the behaviors that break
   the unprotected model get through, against 49% for the best other open guard.
-- **Chemical and biological goals** transferred to three target models: at most 1.1% get through in every setting.
-- **Alert mode** (this release's rule) vs refusing every flagged request: refusals of benign sensitive requests fall
-  from 37% to 5–9%, while at most 3.6% of chemical and biological goals get through on five protected models.
+- **Chemical and biological goals** transferred to three protected models: the full guardrail lets through at most 1.4% of goals on any model.
+- **Alert mode** (`mode="alert"`) vs refusing every flagged request: refusals of benign sensitive requests fall
+  from 37% to 5–16% on six protected models, while at most 3.6% of chemical and biological goals get through.
 - **Latency:** the full guardrail adds 1.5% to generation time on the requests it escalates.
 
 Numbers are means over the five input-classifier seeds; see the paper for the full tables, judges and budgets.
@@ -100,15 +134,17 @@ Numbers are means over the five input-classifier seeds; see the paper for the fu
 Our adapters, heads, probe and code are released under Apache-2.0. The base models keep their own licenses:
 [Llama 3.1 Community License](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) for the probe's protected model,
 and the [Qwen license](https://huggingface.co/Qwen/Qwen3.5-9B) for the classifier backbones. Some training data is
-licensed for non-commercial use only; see the dataset cards.
+licensed for non-commercial use only.
 
 ## Citation
 
 ```bibtex
-@article{eurosafeguard2026,
-  title   = {EuroSafeAI Open Constitutional Classifiers},
-  author  = {Waibl, Leonhard and Chen, Yen-Shan and Simko, Samuel and others},
-  journal = {arXiv preprint arXiv:TODO},
+@article{waibl2026eurosafeguard,
+  title   = {EuroSafeGuard: Open Constitutional Classifiers},
+  author  = {Waibl, Leonhard and Chen, Yen-Shan and Simko, Samuel and Michalak, Felix and
+             Diwan, Nirav and Singh, Arth and Thomas, Aditya and
+             Sch{\"o}lkopf, Bernhard and Jin, Zhijing},
+  journal = {arXiv preprint},
   year    = {2026}
 }
 ```

eurosafeguard.py

@@ -7,6 +7,11 @@
         system_message = g.system_note(category)                   # prepend to the protected model's context
     blocked = g.check_reply(prompt, reply, decision)               # True -> withhold the reply
 
+    # Full mode: the configuration of the paper's end-to-end table (EuroSafeGuard (Full)).
+    g = EuroSafeGuard.from_pretrained("EuroSafeAI/EuroSafeGuard", seed=0, mode="full")
+    escalated = g.escalate(prompt)                                 # probe-only check, before generation
+    blocked = g.check_full(prompt, reply)                          # True -> withhold the reply
+
 Components (see README.md):
   probe      MLP probe on Llama-3.1-8B-Instruct activations (all 32 blocks) of the chat-formatted prompt; per-token
              logits are smoothed with an EMA and the prompt score is their maximum.
@@ -16,6 +21,10 @@ Rule (alert-unless-CBRN): input score = mean of the probe and gate z-scores (z-s
 alert_input_thr (flags 5% of WildChat-clean): refuse if the gate's top category is CBRN, otherwise answer under an
 alert system note and block the reply if the exchange score exceeds exch_thr_20pct. Below: answer normally and block the
 reply if the exchange score exceeds exch_thr_1pct.
+Rule (full, thresholds/full_seed*.json; reproduces the paper's end-to-end numbers): z-score probe, gate and exchange
+on WildChat-clean. Only requests whose probe z exceeds probe_escalation_z (the WildChat-clean 90th percentile) are
+escalated; an escalated exchange is blocked if 0.6*probe_z + 0.2*gate_z + 0.2*exchange_z > full_thr. Non-escalated
+requests are never blocked, so the gate and the exchange classifier run only on escalated traffic.
 
 Memory: Llama-3.1-8B (~16 GB bf16) + Qwen3.5-9B for the gate (~19 GB) + Qwen3.5-9B for the exchange classifier (~19 GB).
 Pass devices=("cuda:0", "cuda:1", "cuda:2") to spread them, or load_exchange=False to skip the reply check.
@@ -32,6 +41,22 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+W_FULL = np.array([0.6, 0.2, 0.2])   # probe, gate, exchange
+
+
+def full_block(probe, gate, exch, cfg) -> np.ndarray:
+    """Full-mode decision from raw scores (probe EMA-max logit, gate max-head prob, exchange max-head prob).
+    True -> block. Mirrors scripts/e2e_compose_ext.py: score = W.z if probe_z > probe_escalation_z else -inf;
+    blocked if score > full_thr. Pass NaN (or anything) for gate/exch of non-escalated items; they are not used."""
+    Z = np.c_[(np.asarray(probe, dtype=float) - cfg["probe_mu"]) / cfg["probe_sd"],
+              (np.asarray(gate, dtype=float) - cfg["gate_mu"]) / cfg["gate_sd"],
+              (np.asarray(exch, dtype=float) - cfg["exch_mu"]) / cfg["exch_sd"]]
+    esc = Z[:, 0] > cfg["probe_escalation_z"]
+    score = np.full(len(Z), -np.inf)
+    score[esc] = Z[esc] @ W_FULL
+    return score > cfg["full_thr"]
+
+
 HEADS = ["CBRN", "Cyber", "Harm to Minors", "Harmful Manipulation", "Hate Speech", "Illegal Activities",
          "Information Integrity", "Physical Harm", "Privacy", "Self-Harm", "Sexual Content"]
 
@@ -110,11 +135,16 @@ class LoRAClassifier:
 
 # ----------------------------------------------------------------------------------------------------- guard
 class EuroSafeGuard:
-    def __init__(self, root: str, seed: int = 0, devices=("cuda", "cuda", "cuda"), load_exchange: bool = True):
+    def __init__(self, root: str, seed: int = 0, devices=("cuda", "cuda", "cuda"), load_exchange: bool = True,
+                 mode: str = "alert"):
         from safetensors.torch import load_file
         from transformers import AutoModelForCausalLM, AutoTokenizer
 
+        if mode not in ("alert", "full"):
+            raise ValueError("mode must be 'alert' (input/alert rule) or 'full' (end-to-end rule)")
+        self.mode = mode
         self.c = json.load(open(os.path.join(root, "thresholds", f"seed{seed}.json")))
+        self.cf = json.load(open(os.path.join(root, "thresholds", f"full_seed{seed}.json"))) if mode == "full" else None
         self.tax = json.load(open(os.path.join(root, "taxonomy_11heads_short.json")))
         pc = json.load(open(os.path.join(root, "probe", "probe_config.json")))
         self.layers, self.M, self.max_len = pc["layers"], pc["M"], pc["max_len"]
@@ -133,7 +163,8 @@ class EuroSafeGuard:
             return cls(repo_or_path, seed, **kw)
         from huggingface_hub import snapshot_download
         root = snapshot_download(repo_or_path, allow_patterns=["probe/*", f"gate/seed{seed}/*", "exchange/*",
-                                                               f"thresholds/seed{seed}.json", "*.json", "*.py"])
+                                                               f"thresholds/seed{seed}.json",
+                                                               f"thresholds/full_seed{seed}.json", "*.json", "*.py"])
         return cls(root, seed, **kw)
 
     # -- input stage
@@ -196,6 +227,39 @@ class EuroSafeGuard:
         t = [f"User: {a}\n\nAssistant: {b}" for a, b in zip(prompts, replies)]
         return self.exchange.probs(t).max(1)
 
+    # -- full mode (paper's end-to-end configuration)
+    def _need_full(self):
+        if self.cf is None:
+            raise RuntimeError("load with mode='full'")
+
+    def escalate(self, prompt):
+        """Probe-only escalation check (full mode): True if probe z > probe_escalation_z."""
+        self._need_full()
+        single = isinstance(prompt, str)
+        P = [prompt] if single else list(prompt)
+        e = (self.probe_scores(P) - self.cf["probe_mu"]) / self.cf["probe_sd"] > self.cf["probe_escalation_z"]
+        out = [bool(x) for x in e]
+        return out[0] if single else out
+
+    def full_scores(self, prompts: list[str], replies: list[str]) -> np.ndarray:
+        """-> raw [N, 3] (probe, gate, exchange); gate and exchange are NaN for non-escalated items."""
+        self._need_full()
+        p = self.probe_scores(list(prompts))
+        esc = np.flatnonzero((p - self.cf["probe_mu"]) / self.cf["probe_sd"] > self.cf["probe_escalation_z"])
+        R = np.full((len(p), 3), np.nan); R[:, 0] = p
+        if len(esc):
+            R[esc, 1] = self.gate.probs([prompts[i] for i in esc]).max(1)
+            R[esc, 2] = self.exchange_scores([prompts[i] for i in esc], [replies[i] for i in esc])
+        return R
+
+    def check_full(self, prompt, reply):
+        """Full mode: True if the exchange (prompt, reply) should be blocked."""
+        single = isinstance(prompt, str)
+        P, Rp = ([prompt], [reply]) if single else (list(prompt), list(reply))
+        R = self.full_scores(P, Rp)
+        out = [bool(x) for x in full_block(R[:, 0], R[:, 1], R[:, 2], self.cf)]
+        return out[0] if single else out
+
     def check_reply(self, prompt, reply, decision="pass"):
         """True if the reply should be withheld. Stricter threshold for replies to alerted requests."""
         single = isinstance(prompt, str)

thresholds/full_seed0.json

@@ -0,0 +1,18 @@
+{
+ "probe_mu": -6.815694701503735,
+ "probe_sd": 3.444286874679132,
+ "gate_mu": 0.03333534521305959,
+ "gate_sd": 0.10919123782498313,
+ "exch_mu": 0.02559613844389501,
+ "exch_sd": 0.090324355434096,
+ "probe_escalation_z": 1.164447941601081,
+ "probe_escalation_quantile": 0.9,
+ "weights": {
+  "probe": 0.6,
+  "gate": 0.2,
+  "exchange": 0.2
+ },
+ "full_thr": 3.3815571671026943,
+ "calibration": "EuroSafeGuard (Full), the configuration of the paper's end-to-end table. All z-scores use the mean and standard deviation on WildChat-clean (6,853 conversations); the exchange statistics use the exchange classifier's max-head probability on Llama-3-8B-Instruct replies to them. A request is escalated if its probe z exceeds probe_escalation_z (the WildChat-clean 90th percentile); an escalated exchange is blocked if 0.6*probe_z + 0.2*gate_z + 0.2*exchange_z > full_thr. full_thr is the smallest threshold that keeps every benign calibration set under its cap (fit half): WildChat-clean, MMLU, GSM8K, Alpaca, Dolly, OASST 0.5%; XSTest-safe, OR-Bench-80k 2%; OR-Bench-hard 5%; benign rows of styled jailbreak sets 40%. Reproduces scripts/e2e_compose_ext.py.",
+ "seed": 0
+}

thresholds/full_seed1.json

@@ -0,0 +1,18 @@
+{
+ "probe_mu": -6.815694701503735,
+ "probe_sd": 3.444286874679132,
+ "gate_mu": 0.02970501206685616,
+ "gate_sd": 0.11025653741405626,
+ "exch_mu": 0.02559613844389501,
+ "exch_sd": 0.090324355434096,
+ "probe_escalation_z": 1.164447941601081,
+ "probe_escalation_quantile": 0.9,
+ "weights": {
+  "probe": 0.6,
+  "gate": 0.2,
+  "exchange": 0.2
+ },
+ "full_thr": 3.5707621643975536,
+ "calibration": "EuroSafeGuard (Full), the configuration of the paper's end-to-end table. All z-scores use the mean and standard deviation on WildChat-clean (6,853 conversations); the exchange statistics use the exchange classifier's max-head probability on Llama-3-8B-Instruct replies to them. A request is escalated if its probe z exceeds probe_escalation_z (the WildChat-clean 90th percentile); an escalated exchange is blocked if 0.6*probe_z + 0.2*gate_z + 0.2*exchange_z > full_thr. full_thr is the smallest threshold that keeps every benign calibration set under its cap (fit half): WildChat-clean, MMLU, GSM8K, Alpaca, Dolly, OASST 0.5%; XSTest-safe, OR-Bench-80k 2%; OR-Bench-hard 5%; benign rows of styled jailbreak sets 40%. Reproduces scripts/e2e_compose_ext.py.",
+ "seed": 1
+}

thresholds/full_seed2.json

@@ -0,0 +1,18 @@
+{
+ "probe_mu": -6.815694701503735,
+ "probe_sd": 3.444286874679132,
+ "gate_mu": 0.03268957031742354,
+ "gate_sd": 0.12878089011617874,
+ "exch_mu": 0.02559613844389501,
+ "exch_sd": 0.090324355434096,
+ "probe_escalation_z": 1.164447941601081,
+ "probe_escalation_quantile": 0.9,
+ "weights": {
+  "probe": 0.6,
+  "gate": 0.2,
+  "exchange": 0.2
+ },
+ "full_thr": 3.131896552703142,
+ "calibration": "EuroSafeGuard (Full), the configuration of the paper's end-to-end table. All z-scores use the mean and standard deviation on WildChat-clean (6,853 conversations); the exchange statistics use the exchange classifier's max-head probability on Llama-3-8B-Instruct replies to them. A request is escalated if its probe z exceeds probe_escalation_z (the WildChat-clean 90th percentile); an escalated exchange is blocked if 0.6*probe_z + 0.2*gate_z + 0.2*exchange_z > full_thr. full_thr is the smallest threshold that keeps every benign calibration set under its cap (fit half): WildChat-clean, MMLU, GSM8K, Alpaca, Dolly, OASST 0.5%; XSTest-safe, OR-Bench-80k 2%; OR-Bench-hard 5%; benign rows of styled jailbreak sets 40%. Reproduces scripts/e2e_compose_ext.py.",
+ "seed": 2
+}

thresholds/full_seed3.json

@@ -0,0 +1,18 @@
+{
+ "probe_mu": -6.815694701503735,
+ "probe_sd": 3.444286874679132,
+ "gate_mu": 0.03866907581677466,
+ "gate_sd": 0.13433606451867683,
+ "exch_mu": 0.02559613844389501,
+ "exch_sd": 0.090324355434096,
+ "probe_escalation_z": 1.164447941601081,
+ "probe_escalation_quantile": 0.9,
+ "weights": {
+  "probe": 0.6,
+  "gate": 0.2,
+  "exchange": 0.2
+ },
+ "full_thr": 3.097429574322379,
+ "calibration": "EuroSafeGuard (Full), the configuration of the paper's end-to-end table. All z-scores use the mean and standard deviation on WildChat-clean (6,853 conversations); the exchange statistics use the exchange classifier's max-head probability on Llama-3-8B-Instruct replies to them. A request is escalated if its probe z exceeds probe_escalation_z (the WildChat-clean 90th percentile); an escalated exchange is blocked if 0.6*probe_z + 0.2*gate_z + 0.2*exchange_z > full_thr. full_thr is the smallest threshold that keeps every benign calibration set under its cap (fit half): WildChat-clean, MMLU, GSM8K, Alpaca, Dolly, OASST 0.5%; XSTest-safe, OR-Bench-80k 2%; OR-Bench-hard 5%; benign rows of styled jailbreak sets 40%. Reproduces scripts/e2e_compose_ext.py.",
+ "seed": 3
+}

thresholds/full_seed4.json

@@ -0,0 +1,18 @@
+{
+ "probe_mu": -6.815694701503735,
+ "probe_sd": 3.444286874679132,
+ "gate_mu": 0.020549819204687854,
+ "gate_sd": 0.1115630808512582,
+ "exch_mu": 0.02559613844389501,
+ "exch_sd": 0.090324355434096,
+ "probe_escalation_z": 1.164447941601081,
+ "probe_escalation_quantile": 0.9,
+ "weights": {
+  "probe": 0.6,
+  "gate": 0.2,
+  "exchange": 0.2
+ },
+ "full_thr": 3.318287213890092,
+ "calibration": "EuroSafeGuard (Full), the configuration of the paper's end-to-end table. All z-scores use the mean and standard deviation on WildChat-clean (6,853 conversations); the exchange statistics use the exchange classifier's max-head probability on Llama-3-8B-Instruct replies to them. A request is escalated if its probe z exceeds probe_escalation_z (the WildChat-clean 90th percentile); an escalated exchange is blocked if 0.6*probe_z + 0.2*gate_z + 0.2*exchange_z > full_thr. full_thr is the smallest threshold that keeps every benign calibration set under its cap (fit half): WildChat-clean, MMLU, GSM8K, Alpaca, Dolly, OASST 0.5%; XSTest-safe, OR-Bench-80k 2%; OR-Bench-hard 5%; benign rows of styled jailbreak sets 40%. Reproduces scripts/e2e_compose_ext.py.",
+ "seed": 4
+}