A LoRA-finetuned multilingual BERT model for 9-class hazard category classification, following the MLCommons AI Safety Hazard Taxonomy.
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2from peft import PeftModel
3
4# Load base model and tokenizer
5base_model = "jhu-clsp/mmBERT-base"
6tokenizer = AutoTokenizer.from_pretrained("llm-semantic-router/mlcommons-safety-classifier-level2-hazard")
7model = AutoModelForSequenceClassification.from_pretrained(base_model, num_labels=9)
8model = PeftModel.from_pretrained(model, "llm-semantic-router/mlcommons-safety-classifier-level2-hazard")
9
10# Classify
11text = "How to hack into someone's email account"
12inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
13outputs = model(**inputs)
14prediction = outputs.logits.argmax(-1).item()
15
16# Label mapping
17labels = [
18 "S1_violent_crimes", "S2_nonviolent_crimes", "S3_sex_crimes",
19 "S5_weapons_cbrne", "S6_self_harm", "S7_hate",
20 "S8_specialized_advice", "S9_privacy", "S13_misinformation"
21]
22print(f"Hazard Category: {labels[prediction]}")
1{
2 "S1_violent_crimes": 0,
3 "S2_nonviolent_crimes": 1,
4 "S3_sex_crimes": 2,
5 "S5_weapons_cbrne": 3,
6 "S6_self_harm": 4,
7 "S7_hate": 5,
8 "S8_specialized_advice": 6,
9 "S9_privacy": 7,
10 "S13_misinformation": 8
11}
1# Step 1: Binary classification (Level 1)
2level1_pred = level1_model(inputs)
3if level1_pred == "unsafe":
4 # Step 2: Hazard classification (Level 2)
5 hazard_category = level2_model(inputs)
1@misc{mlcommons-safety-classifier,
2 title={MLCommons AI Safety Classifier},
3 author={LLM Semantic Router Team},
4 year={2026},
5 publisher={Hugging Face},
6 url={https://huggingface.co/llm-semantic-router/mlcommons-safety-classifier-level2-hazard}
7}
