newVersion app.py (#5)

- newVersion app.py (e37fb3546d643238fd797bd4df6bab6621859c40)


Co-authored-by: kaisexX <kaisex@users.noreply.huggingface.co>

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import json
+import re
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline
 
 # Load Swear Words
@@ -21,41 +22,47 @@ except Exception as e:
     print(f"Error loading model: {e}")
     exit(1)
 
-# Text Classifier Function
+# Text Classification and Censorship Function
 def textclassifier(text):
     if not text.strip():
-        return "Empty input", 0.0
+        return "Empty input", "unknown", 0.0
 
-    # Check for swear words
-    if any(word.lower() in swear_words for word in text.split()):
-        return "swear-word", 1.0
+    # Censor known swear words
+    def censor_word(word):
+        return "***" if word.lower() in swear_words else word
 
-    # Use model
+    words = re.findall(r"\w+|[^\w\s]", text, re.UNICODE)
+    censored_words = [censor_word(word) if re.match(r"\w+", word) else word for word in words]
+    censored_text = " ".join(censored_words)
+
+    # Run model on original input
     try:
         result = text_classifier(text)
         label = result[0]["label"]
         score = result[0]["score"]
 
-        # Threshold logic
+        # Apply threshold for uncertainty
         threshold = 0.994
         if label == "nsfw" and score < threshold:
             label = "uncertain"
 
-        return label, round(score, 4)
+        return censored_text, label, round(score, 4)
 
     except Exception as e:
-        return f"Error: {str(e)}", 0.0
+        return censored_text, f"Error: {str(e)}", 0.0
 
 # Gradio Interface
 interface = gr.Interface(
     fn=textclassifier,
     inputs=gr.Textbox(label="Enter text"),
     outputs=[
-        gr.Label(label="Prediction"),
+        gr.Textbox(label="Censored Text"),
+        gr.Label(label="NSFW Prediction"),
         gr.Number(label="Confidence Score")
     ],
-    title="Text Classifier with Swear Word Filter",
-    # description="First checks for swear words, then uses NSFW text classifier if no swear word is found."
+    title="Text Censorship + NSFW Classifier",
+    description="Censors known swear words using *** and classifies the original text as NSFW, Safe, or Uncertain."
 )
 
-interface.launch()
+if __name__ == "__main__":
+    interface.launch()