Update app.py

app.py

@@ -43,21 +43,21 @@ def moderate(user_input, assistant_response):
     result = tokenizer.decode(output[0], skip_special_tokens=True)
     result = result.split(assistant_response)[-1].strip()
 
-    lines = result.lower().split('\n')
+    lines = [line.strip().lower() for line in result.split('\n') if line.strip()]
     if not lines:
-        return "Error", "Invalid output", "Empty response from model"
+        return "Error", "Empty output", "No valid output from model"
 
-    first_line = lines[0].strip()
+    first_line = lines[0]
     if first_line == 'safe':
         safety_status = "Safe"
         violated_categories = "None"
     elif first_line == 'unsafe':
         safety_status = "Unsafe"
-        violated_categories = lines[1].strip() if len(lines) > 1 else "Unspecified"
+        violated_categories = lines[1] if len(lines) > 1 else "Unspecified"
     else:
         safety_status = "Error"
-        violated_categories = "Invalid output"
-            
+        violated_categories = f"Invalid output: {first_line}"
+        
     return safety_status, violated_categories, result
 
 iface = gr.Interface(