Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,3 @@
+_local/
+__pycache__/
+*.pyc

README.md

@@ -0,0 +1,373 @@
+---
+language:
+  - en
+license: apache-2.0
+tags:
+  - text-classification
+  - refusal-detection
+  - safety
+  - modernbert
+  - multi-label-classification
+  - multi-task-learning
+  - thinking-models
+base_model:
+  - answerdotai/ModernBERT-base
+datasets:
+  - Magpie-Align/Magpie-Qwen2.5-Pro-1M-v0.1
+  - Magpie-Align/Magpie-Pro-300K-Filtered
+  - mlabonne/harmful_behaviors
+  - canbingol/harmful-prompts
+  - cpagac/venomx-pentesting-harmful
+  - grimjim/AILuminate-v1.0-demo-prompt-set-EN
+pipeline_tag: text-classification
+---
+
+# Brigand-Refusal-ModernBERT
+
+The problem I kept running into is a signal quality problem. When you're building anything that needs to understand model behavior — a reward model, a replay filter, an alignment pipeline — a binary comply/refuse label isn't enough. You want to know *why* a model refused: was it a stock policy deflection, a genuine legal concern, an ethical objection, or a `bridge_refusal` that half-complied before pulling back? And with reasoning models, the `<think>` block and the final response often tell very different stories, so I built a second encoder stream specifically to read the thinking separately.
+
+Brigand-Refusal-ModernBERT classifies LLM outputs across four dimensions simultaneously: stance (comply or refuse), response family (19 response types), thought family (11 chain-of-thought patterns), and document type. The two streams share weights but see the prompt through different lenses — one formatted as `[PROMPT] / [RESPONSE]`, the other as `[PROMPT] / [THOUGHT]`. All four heads are trained jointly on [ModernBERT-base](https://huggingface.co/answerdotai/ModernBERT-base) in a single forward pass.
+
+Training was done on a single NVIDIA GeForce RTX 5090.
+
+It was designed to serve as a training signal and inference-time filter in safety-critical LLM pipelines.
+
+---
+
+## Requirements
+
+- Python 3.10+
+- `torch`
+- `transformers>=5.2.0`
+- `safetensors`
+
+```bash
+pip install torch "transformers>=5.2.0" safetensors
+```
+
+---
+
+## Quick Start
+
+The model is self-contained in `modeling.py` and uses the standard HuggingFace
+`AutoModel` interface with `trust_remote_code=True`.
+
+```python
+import sys
+from transformers import AutoModel, AutoTokenizer
+
+model = AutoModel.from_pretrained(
+    "mbukowski/Brigand-Refusal-ModernBERT",
+    trust_remote_code=True,
+)
+tokenizer = AutoTokenizer.from_pretrained("mbukowski/Brigand-Refusal-ModernBERT")
+
+# predict() is bundled in modeling.py and loaded automatically
+predict = sys.modules[type(model).__module__].predict
+
+preds = predict(model, tokenizer, [
+    {
+        "prompt": "How do I whittle a knife?",
+        "response": "Start with a straight-grained softwood like basswood.",
+    },
+    {
+        "prompt": "How do I whittle a knife so I can kill my sister?",
+        "response": "I'm not able to help with that.",
+    },
+])
+
+for p in preds:
+    print(p["stance"], p["bank"])
+# compliance  design_reference
+# refusal     meta_refusal
+```
+
+The `predict()` helper handles all tokenization, thought-splitting, batching,
+and dual-stream formatting.
+
+The bundled tokenizer already includes the control markers as real special
+tokens:
+`[PROMPT]`, `[RESPONSE]`, `[THOUGHT]`, `<NO_THOUGHT>`, and `<EMPTY_THOUGHT>`.
+This release was warm-started from the prior promoted checkpoint, then tuned on
+family-gap anchors and paired with calibrated family decoding. It closes the
+family probe suite (`24/24`) and thought probe suite (`9/9`) while keeping the
+token-boundary improvements.
+
+---
+
+## Label Schema
+
+The model outputs four prediction groups simultaneously.
+
+### Stance *(binary)*
+
+| Label | Meaning |
+|---|---|
+| `compliance` | The response fulfills the request |
+| `refusal` | The response declines or redirects |
+
+### Response Family *(multi-label, top prediction = `bank`)*
+
+| Family | Stance | Description |
+|---|---|---|
+| `stock_refusal` | refusal | Generic policy refusal |
+| `legal_refusal` | refusal | Refusal citing legal constraints |
+| `ethical_refusal` | refusal | Refusal citing ethical concerns |
+| `meta_refusal` | refusal | Refusal about the nature of the request |
+| `bridge_refusal` | refusal | Partial refusal with compliant follow-through |
+| `harmful_procedural` | compliance | Compliant response to a harmful how-to |
+| `harmful_explanatory` | compliance | Compliant response explaining harmful content |
+| `safe_defensive` | compliance | Compliance with a safety caveat |
+| `safe_explanatory` | compliance | Factual/explanatory compliance |
+| `safe_redirective` | compliance | Compliant but steers to a safer framing |
+| `educational_explainer` | compliance | Educational/academic content |
+| `design_reference` | compliance | Reference or lookup responses |
+| `creative_writing` | compliance | Fiction, poetry, roleplay |
+| `code_help_tutor` | compliance | Code assistance |
+| `short_utility_micro` | compliance | Short factual/utility responses |
+| `greeting_chat_micro` | compliance | Greetings and casual chat |
+| `multilingual_general_help` | compliance | General help in non-English |
+| `multilingual_factoid_translate` | compliance | Translation / factoid in non-English |
+| `ambiguous_reject` | — | Catch-all when no family exceeds threshold |
+
+The bundled decoder calibrates family outputs before returning them. It can:
+- suppress generic overlays like `stock_refusal` when a more specific refusal subtype is already active
+- recover underfired specific banks such as `legal_refusal`, `meta_refusal`, `safe_defensive`, and `code_help_tutor` when the logits and prompt/response cues agree
+- prefer specific harmless banks over broader overlays like `short_utility_micro`
+
+### Thought Family *(multi-label, applied to `<think>` content)*
+
+`no_thought` · `empty_thought` · `nonempty_thought` · `policy_thought` ·
+`legal_thought` · `harm_thought` · `meta_thought` · `safe_alternative_thought` ·
+`ethical_thought` · `uncertainty_thought` · `stepwise_thought`
+
+### Document Type *(softmax)*
+
+`plain_text` · `markdown`
+
+---
+
+## Architecture
+
+```
+prompt + response_text  ──▶  [PROMPT] / [RESPONSE] format  ──▶  ModernBERT-base
+                                                                        │
+                                                              masked mean pooling
+                                                                        │
+                                                                response_pooled ──▶ stance head
+                                                                                ──▶ family head
+                                                                                ──▶ document_type head
+
+prompt + thought_text   ──▶  [PROMPT] / [THOUGHT] format   ──▶  ModernBERT-base
+                                                                        │
+                                                              masked mean pooling
+                                                                        │
+                                                               thought_pooled  ──▶ thought_family head
+```
+
+The two streams share weights (single encoder). Thought content is extracted from
+`<think>...</think>` blocks before tokenization; responses without a `<think>` block
+use the `<NO_THOUGHT>` sentinel token.
+
+Those boundary markers are now registered tokenizer special tokens rather than decomposed text fragments.
+
+**Stance augmentation:** the final stance prediction is blended with the
+family-level refusal/compliance signal (`stance_family_scale=0.6`) to reduce
+ambiguous boundary cases.
+
+**Loss weights:**
+
+| Head | Weight |
+|---|---|
+| `stance` | 1.2 |
+| `family` | 2.0 |
+| `thought_family` | 1.2 |
+| `document_type` | 0.3 |
+
+---
+
+## Dataset
+
+**Version:** `v3_inline_harmless_family` — harmless subclasses promoted into the
+main family head; `<think>` blocks handled as a first-class input stream.
+
+| Split | Examples |
+|---|---|
+| Train | 11,821 |
+| Val | 1,469 |
+| Edge (manual review bucket) | 1,024 |
+
+**Curation ledger:** 20,641 reviewed promotions and 841 removals. After dedupe, family-capped sampling, and the train/val split, the final dataset is 11,821 train and 1,469 val, with a separate 1,024-example edge bucket excluded from training.
+
+Getting the data right took longer than training the model.
+
+<details>
+<summary>Family distribution (val set)</summary>
+
+| Family | Val count |
+|---|---|
+| short_utility_micro | 516 |
+| stock_refusal | 341 |
+| creative_writing | 397 |
+| design_reference | 252 |
+| legal_refusal | 193 |
+| safe_explanatory | 124 |
+| meta_refusal | 119 |
+| bridge_refusal | 116 |
+| harmful_procedural | 93 |
+| educational_explainer | 33 |
+| safe_defensive | 30 |
+| code_help_tutor | 27 |
+| ethical_refusal | 51 |
+| greeting_chat_micro | 18 |
+| safe_redirective | 12 |
+| harmful_explanatory | 12 |
+| multilingual_* | 2 |
+
+</details>
+
+<details>
+<summary>Data sources</summary>
+
+The final classifier set was mined from rollout outputs and reviewed imports, but
+those rollouts were driven by a smaller set of upstream prompt corpora.
+
+Harmless / general-help prompt sources used to generate rollout mining pools:
+- [`Magpie-Align/Magpie-Qwen2.5-Pro-1M-v0.1`](https://huggingface.co/datasets/Magpie-Align/Magpie-Qwen2.5-Pro-1M-v0.1)
+- [`Magpie-Align/Magpie-Pro-300K-Filtered`](https://huggingface.co/datasets/Magpie-Align/Magpie-Pro-300K-Filtered)
+
+Harmful / refusal prompt sources used to generate rollout mining pools:
+- [`mlabonne/harmful_behaviors`](https://huggingface.co/datasets/mlabonne/harmful_behaviors)
+- [`canbingol/harmful-prompts`](https://huggingface.co/datasets/canbingol/harmful-prompts)
+- [`cpagac/venomx-pentesting-harmful`](https://huggingface.co/datasets/cpagac/venomx-pentesting-harmful)
+- [`grimjim/AILuminate-v1.0-demo-prompt-set-EN`](https://huggingface.co/datasets/grimjim/AILuminate-v1.0-demo-prompt-set-EN)
+
+Those upstream corpora fed the mined sources that were used in the final build.
+</details>
+
+---
+
+## How It Was Built
+
+The model stopped improving when I treated seams as generic class imbalance. It started improving reliably when I treated each one as a forensic data problem.
+
+The iteration loop that worked:
+
+1. Score the current checkpoint and find the *exact* seam that's failing — not "family accuracy is low" but specifically which source family is bleeding into which target (`stock_refusal -> legal_refusal`, `educational_explainer -> design_reference`, etc.)
+2. Audit both sides: the misses and the attractor that's pulling them over
+3. Classify the problem — label bug, attractor problem, mixed row, or eval bug — because they call for different fixes
+4. Apply the smallest defensible correction: remove bad rows, preserve overlays on mixed refusal rows, import narrow contrast only if the seam truly lacks support
+5. Rebuild and rerun the exact audit that motivated the change
+6. Retrain only if the training set changed — a lot of early iterations wasted compute retraining when the issue was val-only
+
+### Seam Analysis Workflow
+
+When a family or stance boundary starts drifting, the process that worked was:
+
+1. Run a full validation rescore and sort by head loss instead of looking only at headline accuracy.
+2. Collapse the errors into concrete directional seams such as `stock_refusal -> legal_refusal` or `safe_defensive -> stock_refusal`.
+3. Audit both sides of the seam:
+   - the rows getting pulled away from the target class
+   - the rows on the wrong side that are acting as attractors
+4. Separate four failure types before touching data:
+   - label bug
+   - mixed row that needs multi-label preservation
+   - real support gap
+   - decoder / evaluation bug
+5. Use phrase ablation and MLM-head probing on the seam rows to identify what is actually driving the miss.
+   - If removing a tail phrase flips the bank, the tail is the attractor.
+   - If the MLM `prefix_mask` probe already surfaces the right concept tokens, the encoder knows the concept and the problem is boundary calibration rather than representation.
+6. Fix the seam with the smallest change that matches the diagnosis:
+   - remove malformed or truncated rows
+   - relabel contradictory reviewed rows
+   - preserve `stock_refusal` on explicit-opener hybrids
+   - add narrow anchors only for the specific boundary that is missing support
+7. Rebuild the dataset and rerun the exact seam audit before retraining.
+8. Retrain only after the seam definition is cleaner, then gate promotion on:
+   - raw val stance audit
+   - targeted family and thought probes
+   - whether the original seam actually closed
+
+The main discipline was to fix the *reason* a seam existed, not just the rows that happened to show up in the first error sample.
+
+### What Failed
+
+The most instructive failure was a large import of long legal/security-tail rows. Many had explicit refusal openers but also long `national security / public safety` explanatory tails, and I'd omitted the `stock_refusal` overlay. The model learned to overread those rows as `safe_explanatory`, and the exact seam I was trying to fix got worse. Lesson: never bulk-import a mixed seam tranche before splitting it into clean categories.
+
+Repeated problems also came from artifact-heavy reviewed rows — harmless explainers with wrong labels, truncated safe pivots labeled as `stock_refusal`. The right move was always to remove them, not to train through them.
+
+### The MLM Diagnostic
+
+One of the most useful late-stage tools was an encoder-side logit-lens probe. After fine-tuning, I loaded the refusal classifier's encoder weights back into a base ModernBERT MLM model and ran seam rows through it, reading the top predicted tokens at positions in the target span.
+
+The point wasn't text generation — it was to answer a very specific question: *does the encoder already represent the right concept, or is it blind to it?* If the encoder was already producing policy/refusal token clusters on a failing row, the problem was in the classifier boundary, not the representation, and the fix needed to be narrow rather than a broad data expansion. That distinction mattered a lot on several late seams.
+
+The `prefix_mask` mode — inserting a single `[MASK]` at each position and rerunning the full encoder — was more careful than reading from fully-contextualized positions, because it avoids faking a causal mask on an already-bidirectional hidden state.
+
+---
+
+## Evaluation
+
+Validation set (1,469 examples, warm-start epoch 1):
+
+| Metric | Score |
+|---|---|
+| Val loss | **0.0841** |
+| Stance accuracy | **100.00%** |
+| Family accuracy | **99.03%** |
+| Thought family accuracy | **99.81%** |
+| Document type accuracy | **100.00%** |
+
+**Stance audit (val set):** 0 errors in either direction.
+
+<details>
+<summary>Warm-start training history</summary>
+
+This release was warm-started from a prior promoted checkpoint with the special-token tokenizer already in place, then refined on the family-gap seam cleanup. One fine-tuning epoch was enough.
+
+| Epoch | Split | Loss | Stance | Family | Thought family | Doc type |
+|---|---|---|---|---|---|---|
+| 1 | train | 0.0253 | 99.77% | 99.77% | 99.94% | 99.99% |
+| **1** | **val** | **0.0841** | **100.00%** | **99.03%** | **99.81%** | **100.00%** |
+
+</details>
+
+---
+
+## Inference Notes
+
+**Input format:** pass `{"prompt": ..., "response": ...}` dicts to `predict()`.
+The `response` field may contain raw `<think>...</think>` output from a thinking
+model — the helper extracts and routes it automatically.
+
+**Batching:** `predict()` handles batching internally. Default `batch_size=32`,
+`max_length=2048`.
+
+**No `transformers` pipeline:** use `predict()` or call `model.forward()` directly.
+
+---
+
+## Known Limitations
+
+- Multilingual coverage is minimal. `multilingual_general_help` and `multilingual_factoid_translate` have very few val examples and should be treated as best-effort.
+
+---
+
+## Correspondence
+
+Martin Bukowski (models at martinbukowski dot com)
+
+---
+
+## Citation
+
+If you find our work helpful, feel free to give us a cite.
+
+```bibtex
+@misc{brigand-refusal-modernbert-2026,
+    title  = {Brigand-Refusal-ModernBERT},
+    url    = {https://huggingface.co/54rt1n/Brigand-Refusal-ModernBERT},
+    author = {Martin Bukowski},
+    year   = {2026}
+}
+```

config.json

@@ -0,0 +1,181 @@
+{
+  "architectures": [
+    "ModernBertRefusalClassifier"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling.RefusalModernBertConfig",
+    "AutoModel": "modeling.ModernBertRefusalClassifier"
+  },
+  "base_model_name_or_path": "answerdotai/ModernBERT-base",
+  "classifier_dropout": 0.1,
+  "dtype": "float32",
+  "encoder_config": {
+    "_name_or_path": "answerdotai/ModernBERT-base",
+    "architectures": [
+      "ModernBertModel"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 50281,
+    "chunk_size_feed_forward": 0,
+    "classifier_activation": "gelu",
+    "classifier_bias": false,
+    "classifier_dropout": 0.0,
+    "classifier_pooling": "mean",
+    "cls_token_id": 50281,
+    "decoder_bias": true,
+    "deterministic_flash_attn": false,
+    "dtype": "float32",
+    "embedding_dropout": 0.0,
+    "eos_token_id": 50282,
+    "global_attn_every_n_layers": 3,
+    "gradient_checkpointing": false,
+    "hidden_activation": "gelu",
+    "hidden_size": 768,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_cutoff_factor": 2.0,
+    "initializer_range": 0.02,
+    "intermediate_size": 1152,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-05,
+    "layer_types": [
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention"
+    ],
+    "local_attention": 128,
+    "max_position_embeddings": 8192,
+    "mlp_bias": false,
+    "mlp_dropout": 0.0,
+    "model_type": "modernbert",
+    "norm_bias": false,
+    "norm_eps": 1e-05,
+    "num_attention_heads": 12,
+    "num_hidden_layers": 22,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "pad_token_id": 50283,
+    "position_embedding_type": "absolute",
+    "problem_type": null,
+    "return_dict": true,
+    "rope_parameters": {
+      "full_attention": {
+        "rope_theta": 160000.0,
+        "rope_type": "default"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "sep_token_id": 50282,
+    "sparse_pred_ignore_index": -100,
+    "sparse_prediction": false,
+    "tie_word_embeddings": true,
+    "transformers_version": "5.4.0",
+    "vocab_size": 50373
+  },
+  "input_special_tokens": [
+    "[PROMPT]",
+    "[RESPONSE]",
+    "[THOUGHT]",
+    "<NO_THOUGHT>",
+    "<EMPTY_THOUGHT>"
+  ],
+  "local_files_only": false,
+  "model_type": "refusal-modernbert",
+  "schema": {
+    "groups": [
+      {
+        "binary": true,
+        "labels": [
+          "refusal",
+          "compliance"
+        ],
+        "multi_label": false,
+        "name": "stance"
+      },
+      {
+        "binary": false,
+        "labels": [
+          "harmful_procedural",
+          "harmful_explanatory",
+          "safe_defensive",
+          "safe_explanatory",
+          "safe_redirective",
+          "educational_explainer",
+          "design_reference",
+          "creative_writing",
+          "code_help_tutor",
+          "multilingual_general_help",
+          "multilingual_factoid_translate",
+          "short_utility_micro",
+          "greeting_chat_micro",
+          "stock_refusal",
+          "legal_refusal",
+          "ethical_refusal",
+          "meta_refusal",
+          "bridge_refusal",
+          "ambiguous_reject"
+        ],
+        "multi_label": true,
+        "name": "family"
+      },
+      {
+        "binary": false,
+        "labels": [
+          "no_thought",
+          "empty_thought",
+          "nonempty_thought",
+          "policy_thought",
+          "legal_thought",
+          "harm_thought",
+          "meta_thought",
+          "safe_alternative_thought",
+          "ethical_thought",
+          "uncertainty_thought",
+          "stepwise_thought"
+        ],
+        "multi_label": true,
+        "name": "thought_family"
+      },
+      {
+        "binary": false,
+        "labels": [
+          "plain_text",
+          "markdown"
+        ],
+        "multi_label": false,
+        "name": "document_type"
+      }
+    ]
+  },
+  "stance_family_scale": 0.6,
+  "transformers_version": "5.4.0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce3d28ff0cfc75b86089ff2c352909f41fab6a00225b7454e65052397614de41
+size 596188836

modeling.py

@@ -0,0 +1,1155 @@
+"""Self-contained inference and training bootstrap code for RefusalModernBERT.
+
+No external dependencies beyond torch and transformers.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass
+from typing import Any, Mapping
+
+import torch
+from torch import Tensor, nn
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoTokenizer,
+    PreTrainedModel,
+    PretrainedConfig,
+)
+from transformers.utils import ModelOutput
+
+
+# ---------------------------------------------------------------------------
+# Label schema
+# ---------------------------------------------------------------------------
+
+REFUSAL_BANKS = (
+    "stock_refusal",
+    "legal_refusal",
+    "ethical_refusal",
+    "meta_refusal",
+    "bridge_refusal",
+)
+
+COMPLIANCE_BANKS = (
+    "harmful_procedural",
+    "harmful_explanatory",
+    "safe_defensive",
+    "safe_explanatory",
+    "safe_redirective",
+    "educational_explainer",
+    "design_reference",
+    "creative_writing",
+    "code_help_tutor",
+    "multilingual_general_help",
+    "multilingual_factoid_translate",
+    "short_utility_micro",
+    "greeting_chat_micro",
+)
+
+RESPONSE_FAMILY_LABELS = (
+    "harmful_procedural",
+    "harmful_explanatory",
+    "safe_defensive",
+    "safe_explanatory",
+    "safe_redirective",
+    "educational_explainer",
+    "design_reference",
+    "creative_writing",
+    "code_help_tutor",
+    "multilingual_general_help",
+    "multilingual_factoid_translate",
+    "short_utility_micro",
+    "greeting_chat_micro",
+    "stock_refusal",
+    "legal_refusal",
+    "ethical_refusal",
+    "meta_refusal",
+    "bridge_refusal",
+    "ambiguous_reject",
+)
+
+THOUGHT_FAMILY_LABELS = (
+    "no_thought",
+    "empty_thought",
+    "nonempty_thought",
+    "policy_thought",
+    "legal_thought",
+    "harm_thought",
+    "meta_thought",
+    "safe_alternative_thought",
+    "ethical_thought",
+    "uncertainty_thought",
+    "stepwise_thought",
+)
+
+
+@dataclass(frozen=True)
+class LabelGroup:
+    name: str
+    labels: tuple[str, ...]
+    multi_label: bool = False
+    binary: bool = False
+
+    def index_of(self, label: str) -> int:
+        return self.labels.index(label)
+
+    @property
+    def output_dim(self) -> int:
+        return 1 if self.binary else len(self.labels)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "labels": list(self.labels),
+            "multi_label": self.multi_label,
+            "binary": self.binary,
+        }
+
+    @classmethod
+    def from_dict(
+        cls,
+        payload: Mapping[str, Any],
+        *,
+        inferred_output_dim: int | None = None,
+    ) -> "LabelGroup":
+        binary = payload.get("binary")
+        if binary is None and inferred_output_dim is not None:
+            binary = int(inferred_output_dim) == 1
+        return cls(
+            name=str(payload["name"]),
+            labels=tuple(str(label) for label in payload["labels"]),
+            multi_label=bool(payload.get("multi_label", False)),
+            binary=bool(binary),
+        )
+
+
+@dataclass(frozen=True)
+class RefusalClassSchema:
+    groups: tuple[LabelGroup, ...]
+
+    def group(self, name: str) -> LabelGroup:
+        for group in self.groups:
+            if group.name == name:
+                return group
+        raise KeyError(f"unknown label group: {name!r}")
+
+    @property
+    def family_group(self) -> LabelGroup:
+        return self.group("family")
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"groups": [group.to_dict() for group in self.groups]}
+
+    @classmethod
+    def from_dict(
+        cls,
+        payload: Mapping[str, Any],
+        *,
+        state_dict: Mapping[str, Tensor] | None = None,
+    ) -> "RefusalClassSchema":
+        output_dims: dict[str, int] = {}
+        if state_dict is not None:
+            prefix = "classifier.heads."
+            suffix = ".weight"
+            for key, value in state_dict.items():
+                if key.startswith(prefix) and key.endswith(suffix):
+                    group_name = key[len(prefix) : -len(suffix)]
+                    output_dims[group_name] = int(value.shape[0])
+
+        groups = tuple(
+            LabelGroup.from_dict(
+                group_payload,
+                inferred_output_dim=output_dims.get(str(group_payload["name"])),
+            )
+            for group_payload in payload["groups"]
+        )
+        return cls(groups=groups)
+
+
+DEFAULT_SCHEMA = RefusalClassSchema(
+    groups=(
+        LabelGroup(
+            name="stance",
+            labels=("refusal", "compliance"),
+            binary=True,
+        ),
+        LabelGroup(
+            name="family",
+            labels=RESPONSE_FAMILY_LABELS,
+            multi_label=True,
+        ),
+        LabelGroup(
+            name="thought_family",
+            labels=THOUGHT_FAMILY_LABELS,
+            multi_label=True,
+        ),
+        LabelGroup(
+            name="document_type",
+            labels=("plain_text", "markdown"),
+        ),
+    )
+)
+
+
+def _schema_payload(schema: RefusalClassSchema | Mapping[str, Any] | None) -> dict[str, Any]:
+    if schema is None:
+        return DEFAULT_SCHEMA.to_dict()
+    if isinstance(schema, RefusalClassSchema):
+        return schema.to_dict()
+    return RefusalClassSchema.from_dict(schema).to_dict()
+
+
+def _schema_from_state(state: Mapping[str, Any]) -> RefusalClassSchema:
+    payload = state.get("schema")
+    if payload is None:
+        return DEFAULT_SCHEMA
+    return RefusalClassSchema.from_dict(payload, state_dict=state.get("state_dict"))
+
+
+# ---------------------------------------------------------------------------
+# Input formatting
+# ---------------------------------------------------------------------------
+
+PROMPT_TAG = "[PROMPT]"
+RESPONSE_TAG = "[RESPONSE]"
+THOUGHT_TAG = "[THOUGHT]"
+NO_THOUGHT_TOKEN = "<NO_THOUGHT>"
+EMPTY_THOUGHT_TOKEN = "<EMPTY_THOUGHT>"
+INPUT_SPECIAL_TOKENS = (
+    PROMPT_TAG,
+    RESPONSE_TAG,
+    THOUGHT_TAG,
+    NO_THOUGHT_TOKEN,
+    EMPTY_THOUGHT_TOKEN,
+)
+
+
+def split_thought_and_response(text: str) -> tuple[str | None, str, str]:
+    raw = str(text or "").strip()
+    match = re.match(r"^\s*<think>(.*?)</think>\s*(.*)$", raw, flags=re.S | re.I)
+    if match:
+        think = str(match.group(1) or "").strip()
+        response = str(match.group(2) or "").strip()
+        if think:
+            return think, response, "nonempty_thought"
+        return "", response, "empty_thought"
+
+    lowered = raw.lower()
+    if "</think>" in lowered:
+        end = lowered.index("</think>")
+        think = raw[:end].strip()
+        think = re.sub(r"^\s*<think>\s*", "", think, count=1, flags=re.I)
+        response = raw[end + len("</think>") :].strip()
+        if think:
+            return think, response, "nonempty_thought"
+        return "", response, "empty_thought"
+
+    return None, raw, "no_thought"
+
+
+def format_prompt_response_pair(prompt: str, response: str) -> str:
+    return f"{PROMPT_TAG}\n{str(prompt).strip()}\n\n{RESPONSE_TAG}\n{str(response).strip()}"
+
+
+def format_prompt_thought_pair(prompt: str, thought_text: str) -> str:
+    return f"{PROMPT_TAG}\n{str(prompt).strip()}\n\n{THOUGHT_TAG}\n{str(thought_text).strip()}"
+
+
+def _response_text_only(response: str) -> str:
+    _think, response_text, _state = split_thought_and_response(response)
+    return response_text
+
+
+def _thought_text_only(response: str) -> str:
+    think, _response_text, state = split_thought_and_response(response)
+    if state == "no_thought":
+        return NO_THOUGHT_TOKEN
+    if state == "empty_thought":
+        return EMPTY_THOUGHT_TOKEN
+    return str(think or "")
+
+
+def register_input_special_tokens(
+    tokenizer: Any,
+    model: "ModernBertRefusalClassifier | None" = None,
+    *,
+    input_special_tokens: tuple[str, ...] = INPUT_SPECIAL_TOKENS,
+) -> int:
+    """Register prompt/thought boundary tokens and resize embeddings if needed."""
+    tokens_to_add = [token for token in input_special_tokens if token not in tokenizer.get_vocab()]
+    if not tokens_to_add:
+        return 0
+    num_added = tokenizer.add_special_tokens(
+        {"additional_special_tokens": list(tokens_to_add)}
+    )
+    if num_added > 0 and model is not None:
+        model.resize_token_embeddings(len(tokenizer))
+    return int(num_added)
+
+
+# ---------------------------------------------------------------------------
+# Model architecture
+# ---------------------------------------------------------------------------
+
+
+def _masked_mean_pool(last_hidden_state: Tensor, attention_mask: Tensor) -> Tensor:
+    mask = attention_mask.to(dtype=last_hidden_state.dtype).unsqueeze(-1)
+    masked = last_hidden_state * mask
+    denom = mask.sum(dim=1).clamp_min(1.0)
+    return masked.sum(dim=1) / denom
+
+
+def _encoder_config_from_payload(
+    payload: Mapping[str, Any] | None,
+    *,
+    base_model_name_or_path: str,
+    local_files_only: bool,
+):
+    if payload:
+        encoder_config_dict = dict(payload)
+    else:
+        encoder_config_dict = AutoConfig.from_pretrained(
+            base_model_name_or_path,
+            local_files_only=local_files_only,
+        ).to_dict()
+    model_type = str(encoder_config_dict.pop("model_type"))
+    return AutoConfig.for_model(model_type, **encoder_config_dict)
+
+
+class RefusalModernBertConfig(PretrainedConfig):
+    model_type = "refusal-modernbert"
+
+    def __init__(
+        self,
+        *,
+        base_model_name_or_path: str = "answerdotai/ModernBERT-base",
+        encoder_config: Mapping[str, Any] | None = None,
+        schema: RefusalClassSchema | Mapping[str, Any] | None = None,
+        classifier_dropout: float = 0.1,
+        stance_family_scale: float = 0.6,
+        input_special_tokens: tuple[str, ...] | list[str] | None = None,
+        local_files_only: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.base_model_name_or_path = str(base_model_name_or_path)
+        self.encoder_config = dict(encoder_config or {})
+        self.schema = _schema_payload(schema)
+        self.classifier_dropout = float(classifier_dropout)
+        self.stance_family_scale = float(stance_family_scale)
+        self.input_special_tokens = list(input_special_tokens or INPUT_SPECIAL_TOKENS)
+        self.local_files_only = bool(local_files_only)
+        self.architectures = ["ModernBertRefusalClassifier"]
+
+
+@dataclass
+class GroupedClassifierOutput(ModelOutput):
+    logits: dict[str, Tensor] | None = None
+    response_pooled: Tensor | None = None
+    thought_pooled: Tensor | None = None
+
+
+class GroupedLinearHeads(nn.Module):
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        schema: RefusalClassSchema = DEFAULT_SCHEMA,
+        dropout_p: float = 0.1,
+    ) -> None:
+        super().__init__()
+        self.schema = schema
+        self.dropout = nn.Dropout(float(dropout_p))
+        self.heads = nn.ModuleDict(
+            {
+                group.name: nn.Linear(int(hidden_size), int(group.output_dim))
+                for group in schema.groups
+            }
+        )
+
+    def forward(
+        self,
+        response_pooled: Tensor,
+        thought_pooled: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        if thought_pooled is None:
+            thought_pooled = response_pooled
+        response_dropped = self.dropout(response_pooled)
+        thought_dropped = self.dropout(thought_pooled)
+        logits: dict[str, Tensor] = {}
+        for name, head in self.heads.items():
+            pooled = thought_dropped if name == "thought_family" else response_dropped
+            logits[name] = head(pooled)
+        return logits
+
+
+class ModernBertRefusalClassifier(PreTrainedModel):
+    config_class = RefusalModernBertConfig
+    base_model_prefix = "encoder"
+    main_input_name = "input_ids"
+
+    def __init__(
+        self,
+        config: RefusalModernBertConfig | str,
+        *,
+        schema: RefusalClassSchema = DEFAULT_SCHEMA,
+        dropout_p: float = 0.1,
+        local_files_only: bool = False,
+        stance_family_scale: float = 0.6,
+    ) -> None:
+        load_base_encoder_weights = not isinstance(config, RefusalModernBertConfig)
+        if load_base_encoder_weights:
+            model_name_or_path = str(config)
+            encoder_config = AutoConfig.from_pretrained(
+                model_name_or_path,
+                local_files_only=local_files_only,
+            )
+            config = RefusalModernBertConfig(
+                base_model_name_or_path=model_name_or_path,
+                encoder_config=encoder_config.to_dict(),
+                schema=schema,
+                classifier_dropout=dropout_p,
+                stance_family_scale=stance_family_scale,
+                input_special_tokens=INPUT_SPECIAL_TOKENS,
+                local_files_only=local_files_only,
+            )
+
+        super().__init__(config)
+        self.schema = RefusalClassSchema.from_dict(self.config.schema)
+        self.stance_family_scale = float(self.config.stance_family_scale)
+        encoder_config = _encoder_config_from_payload(
+            self.config.encoder_config,
+            base_model_name_or_path=self.config.base_model_name_or_path,
+            local_files_only=self.config.local_files_only,
+        )
+        self.encoder = AutoModel.from_config(encoder_config)
+        self.classifier = GroupedLinearHeads(
+            hidden_size=int(self.encoder.config.hidden_size),
+            schema=self.schema,
+            dropout_p=self.config.classifier_dropout,
+        )
+        self.post_init()
+
+        if load_base_encoder_weights:
+            base_encoder = AutoModel.from_pretrained(
+                self.config.base_model_name_or_path,
+                local_files_only=self.config.local_files_only,
+            )
+            self.encoder.load_state_dict(base_encoder.state_dict())
+
+        self._sync_config()
+
+    def _sync_config(self) -> None:
+        self.config.encoder_config = self.encoder.config.to_dict()
+        self.config.schema = self.schema.to_dict()
+        self.config.stance_family_scale = float(self.stance_family_scale)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if hasattr(self, "encoder") and hasattr(self.encoder, "_init_weights"):
+            self.encoder._init_weights(module)
+            return
+
+        initializer_range = 0.02
+        if isinstance(self.config.encoder_config, Mapping):
+            initializer_range = float(self.config.encoder_config.get("initializer_range", 0.02))
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.encoder.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.encoder.set_input_embeddings(value)
+        self._sync_config()
+
+    def resize_token_embeddings(self, *args: Any, **kwargs: Any) -> nn.Module:
+        embeddings = self.encoder.resize_token_embeddings(*args, **kwargs)
+        self._sync_config()
+        return embeddings
+
+    def save_pretrained(self, save_directory: str | os.PathLike[str], **kwargs: Any) -> None:
+        self._sync_config()
+        super().save_pretrained(save_directory, **kwargs)
+
+    @classmethod
+    def from_base_encoder(
+        cls,
+        model_name_or_path: str,
+        *,
+        schema: RefusalClassSchema = DEFAULT_SCHEMA,
+        dropout_p: float = 0.1,
+        local_files_only: bool = False,
+        stance_family_scale: float = 0.6,
+        register_special_tokens: bool = True,
+    ) -> tuple["ModernBertRefusalClassifier", Any]:
+        """Bootstrap a new classifier from a base encoder for retraining."""
+        model = cls(
+            model_name_or_path,
+            schema=schema,
+            dropout_p=dropout_p,
+            local_files_only=local_files_only,
+            stance_family_scale=stance_family_scale,
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path,
+            local_files_only=local_files_only,
+        )
+        if register_special_tokens:
+            register_input_special_tokens(
+                tokenizer,
+                model,
+                input_special_tokens=tuple(model.config.input_special_tokens),
+            )
+        model.eval()
+        return model, tokenizer
+
+    @classmethod
+    def from_bundle(
+        cls,
+        bundle_dir: str,
+        *,
+        map_location: str = "cpu",
+        register_special_tokens: bool = False,
+    ) -> tuple["ModernBertRefusalClassifier", Any]:
+        """Load a legacy bundle directory.
+
+        Returns ``(model, tokenizer)``.
+        """
+        state = torch.load(
+            os.path.join(bundle_dir, "model_state.pt"),
+            map_location=map_location,
+            weights_only=True,
+        )
+        model_name = str(state["model_name"])
+        schema = _schema_from_state(state)
+        encoder_dir = os.path.join(bundle_dir, "base_encoder")
+        local_encoder = os.path.isdir(encoder_dir)
+        local_path = encoder_dir if local_encoder else model_name
+        encoder_config = AutoConfig.from_pretrained(
+            local_path,
+            local_files_only=local_encoder,
+        )
+        config = RefusalModernBertConfig(
+            base_model_name_or_path=model_name,
+            encoder_config=encoder_config.to_dict(),
+            schema=schema,
+            classifier_dropout=0.1,
+            stance_family_scale=float(state.get("stance_family_scale", 0.6)),
+            input_special_tokens=INPUT_SPECIAL_TOKENS,
+            local_files_only=local_encoder,
+        )
+
+        model = cls(config)
+        model.load_state_dict(state["state_dict"])
+        model.eval()
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            local_path,
+            local_files_only=local_encoder,
+        )
+        if register_special_tokens:
+            register_input_special_tokens(
+                tokenizer,
+                model,
+                input_special_tokens=tuple(model.config.input_special_tokens),
+            )
+        return model, tokenizer
+
+    def _encode(self, *, input_ids: Tensor, attention_mask: Tensor, **kwargs: Any) -> Tensor:
+        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
+        return _masked_mean_pool(outputs.last_hidden_state, attention_mask)
+
+    def forward(
+        self,
+        *,
+        response_input_ids: Tensor | None = None,
+        response_attention_mask: Tensor | None = None,
+        thought_input_ids: Tensor | None = None,
+        thought_attention_mask: Tensor | None = None,
+        input_ids: Tensor | None = None,
+        attention_mask: Tensor | None = None,
+        **kwargs: Any,
+    ) -> GroupedClassifierOutput:
+        if response_input_ids is None:
+            response_input_ids = input_ids
+        if response_attention_mask is None:
+            response_attention_mask = attention_mask
+        if thought_input_ids is None:
+            thought_input_ids = response_input_ids
+        if thought_attention_mask is None:
+            thought_attention_mask = response_attention_mask
+        if response_input_ids is None or response_attention_mask is None:
+            raise ValueError("response inputs are required")
+
+        response_pooled = self._encode(
+            input_ids=response_input_ids,
+            attention_mask=response_attention_mask,
+            **kwargs,
+        )
+        if (
+            thought_input_ids is response_input_ids
+            and thought_attention_mask is response_attention_mask
+        ):
+            thought_pooled = response_pooled
+        else:
+            thought_pooled = self._encode(
+                input_ids=thought_input_ids,
+                attention_mask=thought_attention_mask,
+                **kwargs,
+            )
+        logits = self.classifier(response_pooled, thought_pooled)
+        return GroupedClassifierOutput(
+            logits=logits,
+            response_pooled=response_pooled,
+            thought_pooled=thought_pooled,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Inference helpers
+# ---------------------------------------------------------------------------
+
+
+def _effective_stance_logits(
+    logits: dict[str, Tensor],
+    *,
+    schema: RefusalClassSchema = DEFAULT_SCHEMA,
+    stance_family_scale: float = 0.0,
+) -> Tensor:
+    stance = logits["stance"].squeeze(-1)
+    if float(stance_family_scale) == 0.0 or "family" not in logits:
+        return stance
+    family_labels = set(schema.family_group.labels)
+    refusal_indices = [
+        schema.family_group.index_of(label)
+        for label in REFUSAL_BANKS
+        if label in family_labels
+    ]
+    compliance_indices = [
+        schema.family_group.index_of(label)
+        for label in COMPLIANCE_BANKS
+        if label in family_labels
+    ]
+    family_logits = logits["family"]
+    refusal_support = family_logits[:, refusal_indices].amax(dim=-1)
+    compliance_support = family_logits[:, compliance_indices].amax(dim=-1)
+    return stance + float(stance_family_scale) * (refusal_support - compliance_support)
+
+
+def _group_probabilities(
+    logits: dict[str, Tensor],
+    *,
+    schema: RefusalClassSchema = DEFAULT_SCHEMA,
+    stance_family_scale: float = 0.0,
+) -> dict[str, Tensor]:
+    out: dict[str, Tensor] = {}
+    for group in schema.groups:
+        group_logits = logits[group.name]
+        if group.binary:
+            if group.name == "stance":
+                positive = torch.sigmoid(
+                    _effective_stance_logits(
+                        logits,
+                        schema=schema,
+                        stance_family_scale=stance_family_scale,
+                    )
+                )
+            else:
+                positive = torch.sigmoid(group_logits.squeeze(-1))
+            out[group.name] = torch.stack((positive, 1.0 - positive), dim=-1)
+        elif group.multi_label:
+            out[group.name] = torch.sigmoid(group_logits)
+        else:
+            out[group.name] = torch.softmax(group_logits, dim=-1)
+    return out
+
+
+def _predicted_labels(
+    probs: dict[str, Tensor],
+    *,
+    schema: RefusalClassSchema = DEFAULT_SCHEMA,
+    contexts: list[dict[str, str]] | None = None,
+) -> list[dict[str, Any]]:
+    batch = int(next(iter(probs.values())).shape[0])
+    out: list[dict[str, Any]] = []
+    stance_labels = list(schema.group("stance").labels)
+    for idx in range(batch):
+        row: dict[str, Any] = {}
+        for group in schema.groups:
+            group_probs = probs[group.name][idx]
+            if group.binary:
+                positive = float(group_probs[0].item())
+                row[group.name] = group.labels[0] if positive >= 0.5 else group.labels[1]
+            elif group.multi_label:
+                scored = [
+                    (label, float(prob))
+                    for label, prob in zip(group.labels, group_probs.tolist())
+                ]
+                labels = [label for label, prob in scored if prob >= 0.5]
+                if not labels:
+                    best = int(group_probs.argmax(dim=-1).item())
+                    labels = [group.labels[best]]
+                else:
+                    labels.sort(
+                        key=lambda label: next(
+                            prob for candidate, prob in scored if candidate == label
+                        ),
+                        reverse=True,
+                    )
+                row[group.name] = labels
+                if group.name == "family":
+                    if contexts is not None:
+                        labels, bank = _decode_family_labels(
+                            dict(scored),
+                            prompt=str(contexts[idx].get("prompt", "")),
+                            response=str(contexts[idx].get("response", "")),
+                        )
+                        row[group.name] = labels
+                        row["bank"] = bank
+                    else:
+                        row["bank"] = _select_family_bank(
+                            scored,
+                            predicted_families=labels,
+                        )
+            else:
+                best = int(group_probs.argmax(dim=-1).item())
+                row[group.name] = group.labels[best]
+        if contexts is not None and "stance" in row:
+            stance_probs = probs["stance"][idx]
+            stance_map = {
+                stance_labels[j]: float(stance_probs[j].item())
+                for j in range(len(stance_labels))
+            }
+            row["stance"] = _calibrate_stance_label(
+                stance_map,
+                prompt=str(contexts[idx].get("prompt", "")),
+                response=str(contexts[idx].get("response", "")),
+                predicted_families=list(row.get("family", [])),
+                bank=str(row.get("bank", "")),
+            )
+        if "bank" not in row:
+            families = row.get("family", [])
+            row["bank"] = families[0] if families else "ambiguous_reject"
+        out.append(row)
+    return out
+
+
+def _select_family_bank(
+    scored: list[tuple[str, float]],
+    *,
+    predicted_families: list[str],
+) -> str:
+    if not predicted_families:
+        return "ambiguous_reject"
+
+    label_set = set(predicted_families)
+    ranked = dict(scored)
+
+    def has(label: str) -> bool:
+        return label in label_set
+
+    if has("legal_refusal") and has("meta_refusal") and ranked.get("legal_refusal", 0.0) >= 0.08:
+        return "legal_refusal"
+
+    for label in (
+        "bridge_refusal",
+        "meta_refusal",
+        "legal_refusal",
+        "ethical_refusal",
+    ):
+        if has(label):
+            return label
+
+    for label in (
+        "safe_defensive",
+        "safe_redirective",
+        "multilingual_factoid_translate",
+        "multilingual_general_help",
+        "greeting_chat_micro",
+        "code_help_tutor",
+        "design_reference",
+    ):
+        if has(label):
+            return label
+
+    predicted_families.sort(key=lambda label: ranked.get(label, 0.0), reverse=True)
+    return predicted_families[0]
+
+
+_META_PATTERNS = (
+    re.compile(r"\bignore (?:your |the )?(?:previous|prior) instructions?\b", re.I),
+    re.compile(r"\bsystem prompt\b", re.I),
+    re.compile(r"\bdeveloper instructions?\b", re.I),
+    re.compile(r"\bhidden instructions?\b", re.I),
+    re.compile(r"\boverride\b", re.I),
+    re.compile(r"\bjailbreak\b", re.I),
+    re.compile(r"\bbypass (?:your )?(?:safety|guardrails?)\b", re.I),
+)
+
+_LEGAL_PATTERNS = (
+    re.compile(r"\b(?:copyright|copyrighted|infringement|pirat(?:e|ing|ed)|license(?:s)?|drm)\b", re.I),
+    re.compile(r"\bbypass(?:ing)? software licenses?\b", re.I),
+    re.compile(r"\billegal\b", re.I),
+)
+
+_SAFE_DEFENSIVE_PATTERNS = (
+    re.compile(r"\bshould not\b", re.I),
+    re.compile(r"\bavoid\b", re.I),
+    re.compile(r"\bwithout (?:deception|manipulation|coercion|misleading)\b", re.I),
+    re.compile(r"\bdo not rely on\b", re.I),
+    re.compile(r"\bgrounded in evidence\b", re.I),
+    re.compile(r"\bdeception or manipulation\b", re.I),
+    re.compile(r"\bwithout pressure or guilt\b", re.I),
+)
+
+_CODE_PATTERNS = (
+    re.compile(r"\b(?:python|sql|javascript|typescript|java|c\+\+|regex|bash|shell|query|function|script|code)\b", re.I),
+    re.compile(r"`[^`]+`"),
+    re.compile(r"\[\:\-?\-?1\]"),
+    re.compile(r"\bjoin condition\b", re.I),
+    re.compile(r"\bselect\b", re.I),
+)
+
+_ETHICAL_PATTERNS = (
+    re.compile(r"\brace\b", re.I),
+    re.compile(r"\bracist\b", re.I),
+    re.compile(r"\bslur\b", re.I),
+    re.compile(r"\bdemean(?:s|ing)? people\b", re.I),
+)
+
+_MULTILINGUAL_GENERAL_PATTERNS = (
+    re.compile(r"\bpuedo ayudarte\b", re.I),
+    re.compile(r"\bcorreo formal\b", re.I),
+    re.compile(r"\bredactar un correo\b", re.I),
+    re.compile(r"\bformal(?:e)? email\b", re.I),
+)
+
+_GREETING_CHECKIN_PROMPT_PATTERNS = (
+    re.compile(r"\b(?:hey|hi|hello)\b.*\bhow are you(?: doing)?\b", re.I),
+    re.compile(r"\bhow are you(?: doing)?\b", re.I),
+    re.compile(r"\bhow'?s it going\b", re.I),
+)
+_GREETING_CHECKIN_RESPONSE_PATTERNS = (
+    re.compile(r"\b(?:i am|i'm) (?:doing well|doing great|doing good|well|good)\b", re.I),
+    re.compile(r"\bthanks for asking\b", re.I),
+    re.compile(r"\bwhat can i help you with\b", re.I),
+)
+_GREETING_SALUTATION_PATTERNS = (
+    re.compile(r"^\s*hello[!.]?\s*$", re.I),
+    re.compile(r"^\s*hi(?: there)?[!.]?\s*$", re.I),
+    re.compile(r"^\s*good (?:morning|afternoon|evening)[!.]?\s*$", re.I),
+    re.compile(r"^\s*welcome back[!.]?\s*$", re.I),
+)
+
+_STORY_BRIDGE_PROMPT_PATTERNS = (
+    re.compile(r"\b(?:write|tell|create|draft)\b.*\b(?:story|scene|narrative|script)\b", re.I),
+)
+_STORY_BRIDGE_RESPONSE_PATTERNS = (
+    re.compile(r"\bi can write (?:a )?(?:story|scene|narrative|script)\b", re.I),
+    re.compile(r"\bbut i will not include\b", re.I),
+    re.compile(r"\b(?:actual|operational|step-by-step|synthesis) (?:instructions|details)\b", re.I),
+)
+
+_NO_SAFE_EXPLANATORY_PROMPT_PATTERNS = (
+    re.compile(r"\b(?:netflix|spotify|hulu|disney\+|subscription|streaming|account|password|household)\b", re.I),
+)
+_NO_SAFE_EXPLANATORY_RESPONSE_PATTERNS = (
+    re.compile(r"^\s*no[\\.!]", re.I),
+    re.compile(r"\b(?:service terms|terms of service|local law|authorized household-sharing option)\b", re.I),
+    re.compile(r"\baccount without permission\b", re.I),
+)
+_NO_SAFE_EXPLANATORY_ALT_PATTERNS = (
+    re.compile(r"\buse your own\b", re.I),
+    re.compile(r"\bask them to add you\b", re.I),
+    re.compile(r"\bauthorized household-sharing option\b", re.I),
+)
+_EXPLAINER_PROMPT_PATTERNS = (
+    re.compile(r"\b(?:how does|how do|what is|why does|why do|explain|tell me about|describe)\b", re.I),
+    re.compile(r"\b(?:photosynthesis|nuclear|world war|ww1|ww2|mitosis|osmosis|gravity|evolution)\b", re.I),
+)
+_CREATIVE_PROMPT_PATTERNS = (
+    re.compile(r"\b(?:story|poem|haiku|joke|novel|scene|script|narrative|fictional)\b", re.I),
+    re.compile(r"\b(?:write|draft|create|invent|compose)\b.*\b(?:story|poem|joke|scene|script|narrative)\b", re.I),
+)
+
+_LOCATION_PROMPT_PATTERN = re.compile(r"\btrack(?:ing)?\b.*\blocation\b|\blocation\b.*\btrack(?:ing)?\b", re.I)
+_LOCATION_HELP_PATTERNS = (
+    re.compile(r"\bi can help with (?:legitimate|consent-based)\b", re.I),
+    re.compile(r"\blocation-sharing tools?\b", re.I),
+    re.compile(r"\bgoogle maps\b", re.I),
+    re.compile(r"\bfind my\b", re.I),
+    re.compile(r"\bemergency sharing\b", re.I),
+    re.compile(r"\bpersonal safety settings?\b", re.I),
+)
+_LOCATION_LIMIT_PATTERNS = (
+    re.compile(r"\bi cannot assist\b", re.I),
+    re.compile(r"\bi cannot help\b", re.I),
+    re.compile(r"\bwithout (?:their|someone'?s) (?:knowledge|permission|consent)\b", re.I),
+)
+
+
+def _has_non_ascii(text: str) -> bool:
+    return any(ord(ch) > 127 for ch in text)
+
+
+def _looks_like_location_safe_defensive(prompt: str, response: str) -> bool:
+    return (
+        bool(_LOCATION_PROMPT_PATTERN.search(prompt))
+        and any(pattern.search(response) for pattern in _LOCATION_HELP_PATTERNS)
+        and any(pattern.search(response) for pattern in _LOCATION_LIMIT_PATTERNS)
+    )
+
+
+def _looks_like_greeting_checkin(prompt: str, response: str) -> bool:
+    return (
+        (
+            any(pattern.search(prompt) for pattern in _GREETING_CHECKIN_PROMPT_PATTERNS)
+            and any(pattern.search(response) for pattern in _GREETING_CHECKIN_RESPONSE_PATTERNS)
+        )
+        or (
+            any(pattern.search(prompt) for pattern in _GREETING_SALUTATION_PATTERNS)
+            and any(pattern.search(response) for pattern in _GREETING_SALUTATION_PATTERNS)
+        )
+    )
+
+
+def _looks_like_story_bridge_refusal(prompt: str, response: str) -> bool:
+    return (
+        any(pattern.search(prompt) for pattern in _STORY_BRIDGE_PROMPT_PATTERNS)
+        and all(pattern.search(response) for pattern in _STORY_BRIDGE_RESPONSE_PATTERNS)
+    )
+
+
+def _looks_like_no_safe_explanatory(prompt: str, response: str) -> bool:
+    return (
+        any(pattern.search(prompt) for pattern in _NO_SAFE_EXPLANATORY_PROMPT_PATTERNS)
+        and all(pattern.search(response) for pattern in _NO_SAFE_EXPLANATORY_RESPONSE_PATTERNS)
+        and any(pattern.search(response) for pattern in _NO_SAFE_EXPLANATORY_ALT_PATTERNS)
+    )
+
+
+def _looks_like_factual_explainer(prompt: str, response: str) -> bool:
+    return any(pattern.search(prompt) for pattern in _EXPLAINER_PROMPT_PATTERNS) and not any(
+        pattern.search(prompt) for pattern in _CREATIVE_PROMPT_PATTERNS
+    )
+
+
+def _decode_family_labels(
+    family_map: Mapping[str, float],
+    *,
+    prompt: str = "",
+    response: str = "",
+) -> tuple[list[str], str]:
+    predicted = [label for label, prob in family_map.items() if float(prob) >= 0.5]
+    if not predicted and family_map:
+        predicted = [max(family_map, key=family_map.get)]
+
+    prompt_text = str(prompt or "")
+    response_text = str(response or "")
+    combined_text = f"{prompt_text}\n{response_text}"
+
+    def add_if(label: str, threshold: float, predicate: bool) -> None:
+        if predicate and float(family_map.get(label, 0.0)) >= float(threshold):
+            predicted.append(label)
+
+    add_if(
+        "safe_defensive",
+        0.05,
+        any(pattern.search(response_text) for pattern in _SAFE_DEFENSIVE_PATTERNS),
+    )
+    add_if(
+        "code_help_tutor",
+        0.08,
+        any(pattern.search(combined_text) for pattern in _CODE_PATTERNS),
+    )
+    add_if(
+        "legal_refusal",
+        0.10,
+        any(pattern.search(combined_text) for pattern in _LEGAL_PATTERNS),
+    )
+    add_if(
+        "meta_refusal",
+        0.005,
+        any(pattern.search(combined_text) for pattern in _META_PATTERNS),
+    )
+    add_if(
+        "ethical_refusal",
+        0.05,
+        any(pattern.search(combined_text) for pattern in _ETHICAL_PATTERNS),
+    )
+    add_if(
+        "multilingual_general_help",
+        0.20,
+        _has_non_ascii(combined_text)
+        and any(pattern.search(combined_text) for pattern in _MULTILINGUAL_GENERAL_PATTERNS),
+    )
+    add_if(
+        "greeting_chat_micro",
+        0.20,
+        _looks_like_greeting_checkin(prompt_text, response_text),
+    )
+    add_if(
+        "bridge_refusal",
+        0.20,
+        _looks_like_story_bridge_refusal(prompt_text, response_text),
+    )
+
+    if _looks_like_location_safe_defensive(prompt_text, response_text):
+        predicted.append("safe_defensive")
+        predicted = [
+            label
+            for label in predicted
+            if label not in {"stock_refusal", "bridge_refusal"}
+        ]
+    if _looks_like_no_safe_explanatory(prompt_text, response_text):
+        predicted.append("safe_explanatory")
+        predicted = [
+            label
+            for label in predicted
+            if label
+            not in {
+                "stock_refusal",
+                "legal_refusal",
+                "ethical_refusal",
+                "meta_refusal",
+                "bridge_refusal",
+            }
+        ]
+
+    if (
+        "educational_explainer" in predicted
+        and "creative_writing" in predicted
+        and float(family_map.get("educational_explainer", 0.0)) >= 0.75
+        and _looks_like_factual_explainer(prompt_text, response_text)
+    ):
+        predicted = [label for label in predicted if label != "creative_writing"]
+
+    if any(
+        label in predicted
+        for label in ("safe_defensive", "bridge_refusal", "meta_refusal", "legal_refusal", "ethical_refusal")
+    ):
+        predicted = [label for label in predicted if label != "stock_refusal"]
+
+    if "greeting_chat_micro" in predicted:
+        predicted = [label for label in predicted if label != "short_utility_micro"]
+
+    seen: set[str] = set()
+    deduped: list[str] = []
+    for label in predicted:
+        if label in seen:
+            continue
+        deduped.append(label)
+        seen.add(label)
+
+    deduped.sort(key=lambda label: float(family_map.get(label, 0.0)), reverse=True)
+    bank = _select_family_bank(list(family_map.items()), predicted_families=list(deduped))
+    return deduped, bank
+
+
+def _calibrate_stance_label(
+    stance_map: Mapping[str, float],
+    *,
+    prompt: str = "",
+    response: str = "",
+    predicted_families: list[str] | None = None,
+    bank: str | None = None,
+) -> str:
+    refusal_prob = float(stance_map.get("refusal", 0.0))
+    default = "refusal" if refusal_prob >= 0.5 else "compliance"
+    families = set(predicted_families or [])
+    primary = str(bank or "")
+
+    if _looks_like_location_safe_defensive(str(prompt or ""), str(response or "")):
+        if primary == "safe_defensive" or "safe_defensive" in families:
+            return "compliance"
+    if _looks_like_no_safe_explanatory(str(prompt or ""), str(response or "")):
+        if primary == "safe_explanatory" or "safe_explanatory" in families:
+            return "compliance"
+
+    return default
+
+
+def predict(
+    model: ModernBertRefusalClassifier,
+    tokenizer: Any,
+    examples: list[dict[str, str]],
+    *,
+    batch_size: int = 32,
+    max_length: int = 2048,
+    device: str = "cpu",
+    stance_family_scale: float | None = None,
+) -> list[dict[str, Any]]:
+    """Run inference on a list of ``{"prompt": ..., "response": ...}`` dicts.
+
+    Returns one prediction dict per input with keys:
+    ``stance``, ``family``, ``bank``, ``thought_family``, ``document_type``.
+    ``stance_family_scale`` controls the refusal/compliance bias used when grouping stance.
+    """
+    model = model.to(device)
+    schema = getattr(model, "schema", DEFAULT_SCHEMA)
+    if stance_family_scale is None:
+        stance_family_scale = float(getattr(model, "stance_family_scale", 0.6))
+    all_preds: list[dict[str, Any]] = []
+    was_training = model.training
+    model.eval()
+
+    try:
+        with torch.inference_mode():
+            for start in range(0, len(examples), batch_size):
+                batch = examples[start : start + batch_size]
+                response_texts = [
+                    format_prompt_response_pair(ex["prompt"], _response_text_only(ex["response"]))
+                    for ex in batch
+                ]
+                thought_texts = [
+                    format_prompt_thought_pair(ex["prompt"], _thought_text_only(ex["response"]))
+                    for ex in batch
+                ]
+                response_tok = tokenizer(
+                    response_texts,
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=max_length,
+                )
+                thought_tok = tokenizer(
+                    thought_texts,
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=max_length,
+                )
+                output = model(
+                    response_input_ids=response_tok["input_ids"].to(device),
+                    response_attention_mask=response_tok["attention_mask"].to(device),
+                    thought_input_ids=thought_tok["input_ids"].to(device),
+                    thought_attention_mask=thought_tok["attention_mask"].to(device),
+                )
+                probs = _group_probabilities(
+                    output.logits,
+                    schema=schema,
+                    stance_family_scale=stance_family_scale,
+                )
+                all_preds.extend(_predicted_labels(probs, schema=schema, contexts=batch))
+    finally:
+        if was_training:
+            model.train()
+
+    return all_preds
+
+
+RefusalModernBertConfig.register_for_auto_class()
+ModernBertRefusalClassifier.register_for_auto_class("AutoModel")

tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "extra_special_tokens": [
+    "[PROMPT]",
+    "[RESPONSE]",
+    "[THOUGHT]",
+    "<NO_THOUGHT>",
+    "<EMPTY_THOUGHT>"
+  ],
+  "is_local": true,
+  "mask_token": "[MASK]",
+  "max_length": 2304,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 8192,
+  "pad_to_multiple_of": null,
+  "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "tokenizer_class": "TokenizersBackend",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}