Automatic prompt moderation

app.py

@@ -56,6 +56,29 @@ BASE_MODEL_DISPLAY: Dict[str, str] = {
 
 FALLBACK_IDS: Dict[str, str] = {}
 
+MODERATION_MODEL_ID = "ifmain/ModerationBERT-En-02"
+MODERATION_CATEGORIES = [
+    "harassment",
+    "harassment_threatening",
+    "hate",
+    "hate_threatening",
+    "self_harm",
+    "self_harm_instructions",
+    "self_harm_intent",
+    "sexual",
+    "sexual_minors",
+    "violence",
+    "violence_graphic",
+    "self-harm",
+    "sexual/minors",
+    "hate/threatening",
+    "violence/graphic",
+    "self-harm/intent",
+    "self-harm/instructions",
+    "harassment/threatening",
+]
+MODERATION_THRESHOLD = 0.35
+
 INIT_RATING = 1000
 K_FACTOR = 32
 SCALE = 400
@@ -723,7 +746,6 @@ def build_inputs(tokenizer, model_id: str, prompt: str):
 def is_diffusion_model(model_id: str) -> bool:
     return "metadiffusion" in model_id.lower()
 
-@spaces.GPU(duration=120)
 def generate_for_model(model_id: str, prompt: str) -> str:
     ensure_models_loaded()
     if model_id not in models or model_id not in tokenizers:
@@ -934,6 +956,58 @@ def create_demo() -> gr.Blocks:
         # -------------------------------------------------------------------
         # Event handlers
         # -------------------------------------------------------------------
+        moderation_model = None
+        moderation_tokenizer = None
+
+        def load_moderation():
+            nonlocal moderation_model, moderation_tokenizer
+            if moderation_model is not None:
+                return
+            from transformers import BertTokenizer, BertForSequenceClassification
+            moderation_tokenizer = BertTokenizer.from_pretrained(MODERATION_MODEL_ID)
+            moderation_model = BertForSequenceClassification.from_pretrained(MODERATION_MODEL_ID, num_labels=18)
+            moderation_model.to(DEVICE)
+            moderation_model.eval()
+
+        def moderate_prompt_impl(prompt: str):
+            try:
+                load_moderation()
+                encoding = moderation_tokenizer(
+                    prompt,
+                    add_special_tokens=True,
+                    max_length=128,
+                    padding="max_length",
+                    truncation=True,
+                    return_attention_mask=True,
+                    return_tensors="pt",
+                )
+                with torch.no_grad():
+                    outputs = moderation_model(
+                        encoding["input_ids"].to(moderation_model.device),
+                        attention_mask=encoding["attention_mask"].to(moderation_model.device),
+                    )
+                scores = torch.sigmoid(outputs.logits)[0]
+                flagged = [MODERATION_CATEGORIES[i] for i in range(len(MODERATION_CATEGORIES)) if scores[i].item() >= MODERATION_THRESHOLD]
+                if not flagged:
+                    return None
+                return ", ".join(flagged)
+            except Exception as e:
+                logger.error(f"Moderation check failed, allowing prompt: {e}\n{traceback.format_exc()}")
+                return None
+
+        @spaces.GPU(duration=120)
+        def battle_gpu(spec: ArenaSpec, user_prompt: str, last_pair_val):
+            reason = moderate_prompt_impl(user_prompt)
+            if reason is not None:
+                return reason, None, None, None, None
+            a, b = pick_random_pair(spec, exclude_pair=last_pair_val)
+            if random.random() < 0.5:
+                a, b = b, a
+            ensure_models_loaded()
+            resp_a = generate_for_model(a, user_prompt)
+            resp_b = generate_for_model(b, user_prompt)
+            return None, resp_a, resp_b, a, b
+
         def on_submit(spec: ArenaSpec, user_prompt: str, last_pair_val):
             user_prompt = (user_prompt or "").strip()
             if not user_prompt:
@@ -951,12 +1025,23 @@ def create_demo() -> gr.Blocks:
                     "", "", False, user_prompt, last_pair_val,
                     leaderboard_dataframe(load_elo(spec), spec)
                 )
-            a, b = pick_random_pair(spec, exclude_pair=last_pair_val)
-            if random.random() < 0.5:
-                a, b = b, a
-            ensure_models_loaded()
-            resp_a = generate_for_model(a, user_prompt)
-            resp_b = generate_for_model(b, user_prompt)
+            reason, resp_a, resp_b, a, b = battle_gpu(spec, user_prompt, last_pair_val)
+            if reason is not None:
+                flag_text = f"This prompt was flagged for: {reason}"
+                return (
+                    gr.update(value=flag_text),
+                    gr.update(value=flag_text),
+                    gr.update(visible=False),
+                    gr.update(visible=False),
+                    gr.update(visible=False, value=""),
+                    gr.update(interactive=False),
+                    gr.update(interactive=False),
+                    gr.update(interactive=False),
+                    gr.update(interactive=False),
+                    gr.update(visible=False),
+                    "", "", False, user_prompt, last_pair_val,
+                    leaderboard_dataframe(load_elo(spec), spec)
+                )
             if not resp_a.strip():
                 resp_a = "[No output... model returned empty]"
             if not resp_b.strip():