Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForSequenceClassification
 from scipy.special import softmax
 import torch
+import re
 
 # Load sentiment model
 @st.cache_resource
@@ -17,7 +18,7 @@ def load_emotion_model():
     tokenizer = AutoTokenizer.from_pretrained("j-hartmann/emotion-english-distilroberta-base")
     return model, tokenizer
 
-# Load T5 paraphrasing model
+# Load paraphrasing model
 @st.cache_resource
 def load_paraphrase_model():
     model = AutoModelForSeq2SeqLM.from_pretrained("Vamsi/T5_Paraphrase_Paws")
@@ -42,7 +43,7 @@ def get_emotion(text, model, tokenizer):
     labels = ['anger', 'disgust', 'fear', 'joy', 'neutral', 'sadness', 'surprise']
     return labels[probs.argmax()], float(probs.max()) * 100
 
-# Generate feedback
+# Feedback generation
 def generate_feedback(sentiment, emotion):
     if sentiment == "Negative":
         if emotion in ["anger", "disgust", "sadness"]:
@@ -59,8 +60,16 @@ def generate_feedback(sentiment, emotion):
         else:
             return "🙂 Your message is positive, but think about whether it’s being fully understood."
 
-# Paraphrase / rewrite message
-def rewrite_message(text, model, tokenizer):
+# Profanity detection
+def contains_profanity(text):
+    profane_words = ['fuck', 'shit', 'bitch', 'stupid', 'idiot', 'dumb', 'asshole']
+    return any(re.search(rf"\b{word}\b", text.lower()) for word in profane_words)
+
+# Smart rewrite logic
+def smart_rewrite_message(text, model, tokenizer):
+    if contains_profanity(text):
+        return ["⚠️ Your message may contain harmful language. Please rephrase it with respect and calm."]
+    
     text = "paraphrase: " + text + " </s>"
     encoding = tokenizer.encode_plus(text, return_tensors="pt", max_length=128, truncation=True)
     with torch.no_grad():
@@ -73,9 +82,9 @@ def rewrite_message(text, model, tokenizer):
             temperature=1.5
         )
     rewrites = [tokenizer.decode(o, skip_special_tokens=True) for o in output]
-    return list(set(rewrites))  # remove duplicates
+    return list(set(rewrites))
 
-# UI
+# Streamlit App UI
 st.title("🗣️ Message Tone & Rewrite Checker (Phase 2)")
 st.write("Before you send that message, check how it might be received — and improve it if needed.")
 
@@ -101,7 +110,7 @@ if st.button("Analyze"):
         st.markdown("---")
         st.markdown("### ✨ Try Rewriting Your Message")
         para_model, para_token = load_paraphrase_model()
-        rewrites = rewrite_message(text, para_model, para_token)
+        rewrites = smart_rewrite_message(text, para_model, para_token)
 
         for i, r in enumerate(rewrites, 1):
             st.write(f"**Version {i}:** {r}")