| """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 |
|
|
|
|
| |
| |
| |
|
|
| 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")) |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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, |
| ) |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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") |
|
|
