llm-semantic-router / mmbert-safety-classifier-level1

mmBERT Safety Classifier - Level 1 (Binary)

A binary safety classifier for detecting unsafe content in LLM inputs. Part of a hierarchical MLCommons-aligned safety classification system.

Model Description

This is Level 1 of a 2-level hierarchical safety classifier:

• Level 1 (this model) : Binary classification (safe/unsafe) - high recall for catching threats
• Level 2 : 9-class hazard taxonomy (MLCommons AI Safety aligned) - for categorizing unsafe content

Performance

Labels

Training Hyperparameters

Usage

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel

# Load model
tokenizer = AutoTokenizer.from_pretrained("llm-semantic-router/mmbert-safety-classifier-level1")
base_model = AutoModelForSequenceClassification.from_pretrained(
    "jhu-clsp/mmBERT-base", 
    num_labels=2,
    torch_dtype=torch.float32
)
model = PeftModel.from_pretrained(base_model, "llm-semantic-router/mmbert-safety-classifier-level1")
model.eval()

# Inference
text = "How do I bake a chocolate cake?"
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)

with torch.no_grad():
    outputs = model(**inputs)
    pred = outputs.logits.argmax(-1).item()
    
labels = {0: "safe", 1: "unsafe"}
print(f"Prediction: {labels[pred]}")

Hierarchical Classification Pipeline

For complete safety classification, use with Level 2:

# If Level 1 predicts "unsafe", run Level 2 for hazard category
if pred == 1:  # unsafe
    # Load and run Level 2 model for specific hazard category
    # See: llm-semantic-router/mmbert-safety-classifier-level2
    pass

Training Data

• Primary : nvidia/Aegis-AI-Content-Safety-Dataset-2.0
• Enhanced : 110 edge case examples for underrepresented categories (CSE, medical advice, misinformation)
• Balance : 50/50 safe/unsafe split with oversampling

Intended Use

• Content moderation for LLM applications
• Input filtering for AI safety systems
• Guardrail implementation for chatbots and AI assistants

Limitations

• May miss subtle misinformation content
• Trained primarily on English content
• Should be used as part of a defense-in-depth strategy

Citation

@misc{mmbert-safety-classifier,
  title={mmBERT Safety Classifier},
  author={LLM Semantic Router Team},
  year={2026},
  publisher={HuggingFace},
  url={https://huggingface.co/llm-semantic-router/mmbert-safety-classifier-level1}
}