Update app.py

app.py

@@ -3,67 +3,99 @@ from transformers import pipeline, AutoModelForSequenceClassification, AutoToken
 import torch
 import numpy as np
 
-def main():
-    # Load the spam detection pipeline
+def analyze_email(email_body):
+    # Load models (ideally cached, but kept here for simplicity)
     spam_pipeline = pipeline("text-classification", model="cybersectony/phishing-email-detection-distilbert_v2.4.1")
-
-    # Load the sentiment model and tokenizer
     sentiment_model = AutoModelForSequenceClassification.from_pretrained("ISOM5240GP4/email_sentiment", num_labels=2)
     tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
 
+    # Step 1: Check if the email is spam
+    spam_result = spam_pipeline(email_body)
+    spam_label = spam_result[0]["label"]
+    spam_confidence = spam_result[0]["score"]
+
+    if spam_label == "LABEL_1":
+        return f"This is a spam email (Confidence: {spam_confidence:.2f}). No follow-up needed."
+    else:
+        # Step 2: Analyze sentiment for non-spam emails
+        inputs = tokenizer(email_body, padding=True, truncation=True, return_tensors='pt')
+        outputs = sentiment_model(**inputs)
+        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        predictions = predictions.cpu().detach().numpy()
+        sentiment_index = np.argmax(predictions)
+        sentiment_confidence = predictions[0][sentiment_index]
+        sentiment = "Positive" if sentiment_index == 1 else "Negative"
+
+        if sentiment == "Positive":
+            return (f"This email is not spam (Confidence: {spam_confidence:.2f}).\n"
+                    f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}). No follow-up needed.")
+        else:
+            return (f"This email is not spam (Confidence: {spam_confidence:.2f}).\n"
+                    f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}).\n"
+                    "**This email needs follow-up as it is not spam and has negative sentiment.**")
+
+def main():
     st.title("Email Analysis Tool")
-    st.write("Enter an email body below or select a sample to analyze its spam status and sentiment.")
+    st.write("Enter an email body below or click a sample to analyze its spam status and sentiment.")
 
+    # Initialize session state
     if "email_body" not in st.session_state:
         st.session_state.email_body = ""
+    if "result" not in st.session_state:
+        st.session_state.result = ""
+
+    # Text area for email input
+    email_body = st.text_area("Email Body", value=st.session_state.email_body, height=200, key="email_input",
+                              on_change=lambda: st.session_state.update({"result": analyze_email(st.session_state.email_input)}))
+
+    # Sample emails (stored as variables, not button labels)
+    sample_spam = """
+Subject: Urgent: Verify Your Account Now!
+Dear Customer,
+We have detected unusual activity on your account. To prevent suspension, please verify your login details immediately by clicking the link below:
+[Click Here to Verify](http://totally-legit-site.com/verify)
+Failure to verify within 24 hours will result in your account being locked. This is for your security.
+Best regards,
+The Security Team
+    """
 
-    email_body = st.text_area("Email Body", value=st.session_state.email_body, height=200, key="email_input")
+    sample_not_spam_positive = """
+Subject: Great News About Your Project!
+Hi Team,
+I just wanted to let you know that the project is progressing wonderfully! Everyone’s efforts are paying off, and we’re ahead of schedule. Keep up the fantastic work!
+Best,
+Alex
+    """
 
-    # Sample emails
-    sample_spam = """[Spam email content]"""
-    sample_not_spam_positive = """[Positive email content]"""
-    sample_not_spam_negative = """[Negative email content]"""
+    sample_not_spam_negative = """
+Subject: Issue with Recent Delivery
+Dear Support,
+I received my package today, but it was damaged, and two items were missing. This is really frustrating—please let me know how we can resolve this as soon as possible.
+Thanks,
+Sarah
+    """
 
+    # Buttons with concise labels (in columns)
     col1, col2, col3 = st.columns(3)
     with col1:
-        if st.button("Spam Email"):
+        if st.button("Spam Sample"):
             st.session_state.email_body = sample_spam
+            st.session_state.result = analyze_email(sample_spam)
             st.rerun()
     with col2:
-        if st.button("Not Spam, Positive"):
+        if st.button("Positive Sample"):
             st.session_state.email_body = sample_not_spam_positive
+            st.session_state.result = analyze_email(sample_not_spam_positive)
             st.rerun()
     with col3:
-        if st.button("Not Spam, Negative"):
+        if st.button("Negative Sample"):
             st.session_state.email_body = sample_not_spam_negative
+            st.session_state.result = analyze_email(sample_not_spam_negative)
             st.rerun()
 
-    if st.button("Analyze Email"):
-        if email_body:
-            spam_result = spam_pipeline(email_body)
-            spam_label = spam_result[0]["label"]
-            spam_confidence = spam_result[0]["score"]
-
-            if spam_label == "LABEL_1":
-                st.write(f"This is a spam email (Confidence: {spam_confidence:.2f}). No follow-up needed.")
-            else:
-                inputs = tokenizer(email_body, padding=True, truncation=True, return_tensors='pt')
-                outputs = sentiment_model(**inputs)
-                predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
-                predictions = predictions.cpu().detach().numpy()
-                sentiment_index = np.argmax(predictions)
-                sentiment_confidence = predictions[0][sentiment_index]
-
-                sentiment = "Positive" if sentiment_index == 1 else "Negative"
-                if sentiment == "Positive":
-                    st.write(f"This email is not spam (Confidence: {spam_confidence:.2f}).")
-                    st.write(f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}). No follow-up needed.")
-                else:
-                    st.write(f"This email is not spam (Confidence: {spam_confidence:.2f}).")
-                    st.write(f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}).")
-                    st.write("**This email needs follow-up as it is not spam and has negative sentiment.**")
-        else:
-            st.write("Please enter an email body or select a sample to analyze.")
+    # Display result
+    if st.session_state.result:
+        st.write(st.session_state.result)
 
 if __name__ == "__main__":
     main()