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

mmBERT Safety Classifier - Level 2 (9-Class Hazard Taxonomy)

A 9-class hazard classifier aligned with MLCommons AI Safety taxonomy. Part of a hierarchical safety classification system.

Model Description

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

• Level 1 : Binary classification (safe/unsafe) - for initial filtering
• Level 2 (this model) : 9-class hazard taxonomy - for categorizing the type of unsafe content

Performance

MLCommons-Aligned 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-level2")
base_model = AutoModelForSequenceClassification.from_pretrained(
    "jhu-clsp/mmBERT-base", 
    num_labels=9,
    torch_dtype=torch.float32
)
model = PeftModel.from_pretrained(base_model, "llm-semantic-router/mmbert-safety-classifier-level2")
model.eval()

# Inference (use only after Level 1 predicts "unsafe")
text = "How to hack into someone's bank account"
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)

with torch.no_grad():
    outputs = model(**inputs)
    pred = outputs.logits.argmax(-1).item()
    conf = torch.softmax(outputs.logits, dim=-1)[0][pred].item()

labels = {
    0: "S1_violent_crimes",
    1: "S2_nonviolent_crimes", 
    2: "S3_sex_crimes",
    3: "S5_weapons_cbrne",
    4: "S6_self_harm",
    5: "S7_hate",
    6: "S8_specialized_advice",
    7: "S9_privacy",
    8: "S13_misinformation"
}
print(f"Hazard Category: {labels[pred]} ({conf:.2%})")

Hierarchical Classification Pipeline

Use with Level 1 for complete safety classification:

# Step 1: Run Level 1 to check if content is unsafe
level1_pred = run_level1(text)  # Returns 0=safe, 1=unsafe

# Step 2: If unsafe, run Level 2 for hazard category
if level1_pred == 1:
    hazard_category = run_level2(text)
    print(f"Content is unsafe: {hazard_category}")
else:
    print("Content is safe")

Training Data

• Primary : nvidia/Aegis-AI-Content-Safety-Dataset-2.0 (18,164 samples)
• Enhanced Edge Cases :
  • S3_sex_crimes: 20 CSE examples
  • S8_specialized_advice: 25 medical/legal/financial advice examples
  • S13_misinformation: 25 vaccine/election/conspiracy examples
  • S2_nonviolent_crimes: 20 hacking/fraud examples
  • S5_weapons_cbrne: 20 CBRNE examples
• Balance : Oversampling to ~2,000 samples per class

Class Distribution (Raw Data)

Intended Use

• Detailed hazard categorization for content moderation
• Safety analytics and reporting
• Policy-specific content handling
• MLCommons-compliant safety systems

Limitations

• Should only be used after Level 1 classifies content as "unsafe"
• Some category overlap (e.g., weapons + violent crimes)
• Misinformation detection is challenging for subtle cases
• Trained primarily on English content

Citation

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

References

• MLCommons AI Safety Taxonomy
• NVIDIA AEGIS Dataset