Update app.py

app.py

@@ -38,10 +38,14 @@ def infer(prompt):
     """)
     is_safe = safety_check(prompt)
     print(is_safe)
-    words = is_safe.split()
-    # Take the first word and convert it to a string variable
-    status = str(words[0])
-    print(f"FIRST WORD:{status}")
+    
+    match = re.search(r'\bYes\b', is_safe)
+    
+    if match:
+        status = 'Yes'
+    else:
+        status = None
+    
     if status == "Yes" :
         raise gr.Error("Don't ask for such things.")
     else: