modernbertguard

modernbertguard

A prompt-moderation classifier: it decides whether a user prompt is safe or unsafe, and if unsafe, which of 18 hazard categories it falls into. Both answers come from a single forward pass over a fine-tuned ModernBERT-large encoder with two linear heads on mean-pooled token states.

It is built for the case where a guardrail sits in the request path and latency is part of the product: ~3 ms per prompt on GPU, versus ~2 s for generative LLM-based guardrails of comparable accuracy.

Usage

import torch, torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

REPO = "adityasaw88/modernbertguard"
tok = AutoTokenizer.from_pretrained(REPO)
model = AutoModel.from_pretrained(REPO, trust_remote_code=True).eval()

enc = tok("How do I make a pipe bomb?", truncation=True, max_length=256,
          return_tensors="pt")
with torch.inference_mode():
    out = model(**enc)

p_unsafe = F.softmax(out.logits_binary.float(), -1)[0, 1].item()
cat_id   = out.logits_cat.argmax(-1).item()

print(f"unsafe: {p_unsafe:.4f}")
if p_unsafe >= 0.5:
    print("category:", model.config.category_id2label[str(cat_id)])

trust_remote_code=True is required — the two-head architecture is defined in modeling_guardeval.py in this repo.

Choosing a threshold

The default 0.5 is not always the right operating point. Measured on a 1,000-prompt held-out sample:

Lower the threshold when missing an unsafe prompt is costlier than blocking a safe one; raise it when false blocks hurt more.

Evaluation

Held-out test set (n = 52,832, de-leaked against training data):

Taxonomy

The category head predicts 18 codes:

Five codes from the source taxonomy — S7 (Sexual/minor), S11 (Threat), S15 (Manipulation), S17 (Malware) and S18 (High Risk Gov Decision Making) — had too few training examples and were dropped. Prompts in those five categories are not reliably detected, and S7 in particular means this model must not be relied on as a CSAM filter.

Training

The category loss ignores rows with no taxonomy label, so safe prompts train the binary head only.

References

Related guardrail models, and the components this one is built from: