Update app.py

app.py

@@ -8,6 +8,7 @@ import pandas as pd
 import tempfile
 import time
 import os
+from valx import detect_profanity, detect_hate_speech
 
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 SEED = 1337
@@ -87,6 +88,14 @@ def generate_names(prompt, temperature, top_k, count, retries, seed, randomize_s
     random.seed(seed)
 
     prompt = prompt.strip()
+    promptx = prompt.lower()
+    if detect_profanity([promptx], language='All'):
+        gr.Warning("Profanity detected in the prompt, using the default prompt.")
+        prompt = 'a kind king'
+    elif (hate_speech_result := detect_hate_speech(promptx)) and hate_speech_result[0] in ['Hate Speech', 'Offensive Speech']:
+        gr.Warning('Harmful speech detected in the prompt, using default prompt.')
+        prompt = 'a kind king'
+
     if not prompt:
         raise gr.Error("Prompt cannot be empty.")
     if len(prompt) > 64:
@@ -100,6 +109,13 @@ def generate_names(prompt, temperature, top_k, count, retries, seed, randomize_s
     for _ in range(count):
         for attempt in range(retries):
             name, t = sample_once(prompt, temperature=temperature, top_k=top_k)
+            namex = name.strip().lower()
+            if detect_profanity([namex], language='All'):
+                gr.Warning("Profanity detected in the generated name, flagging...")
+                rejected.append(name + " (Profanity Detected)")
+            elif (hate_speech_result := detect_hate_speech(namex)) and hate_speech_result[0] in ['Hate Speech', 'Offensive Speech']:
+                gr.Warning('Harmful speech detected in the generated name, flagging...')
+                rejected.append(name + " (Harmful Speech Detected)")
             retry_count += 1
             if len(name) >= 3:
                 results.append({"Generated Name": name, "Time (s)": f"{t:.2f}"})