Update app.py

app.py

@@ -3,70 +3,47 @@ from transformers import pipeline, AutoModelForSequenceClassification, AutoToken
 import torch
 import numpy as np
 
-# Load models with caching to improve performance
-@st.cache_resource
-def load_spam_pipeline():
-    return pipeline("text-classification", model="cybersectony/phishing-email-detection-distilbert_v2.4.1")
-
-@st.cache_resource
-def load_sentiment_model():
-    model = AutoModelForSequenceClassification.from_pretrained("ISOM5240GP4/email_sentiment", num_labels=2)
+def analyze_email(email_body):
+    spam_pipeline = pipeline("text-classification", model="cybersectony/phishing-email-detection-distilbert_v2.4.1")
+    sentiment_model = AutoModelForSequenceClassification.from_pretrained("ISOM5240GP4/email_sentiment", num_labels=2)
     tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
-    return model, tokenizer
 
-# Initialize models
-spam_pipeline = load_spam_pipeline()
-sentiment_model, tokenizer = load_sentiment_model()
+    spam_result = spam_pipeline(email_body)
+    spam_label = spam_result[0]["label"]
+    spam_confidence = spam_result[0]["score"]
 
-# Function to analyze email
-def analyze_email(email_body):
-    """Analyzes an email for spam and sentiment, returning result type and message."""
-    if not email_body.strip():
-        return "error", "Email body is empty. Please provide an email to analyze."
-    
-    try:
-        # Step 1: Check if the email is spam
-        spam_result = spam_pipeline(email_body)
-        spam_label = spam_result[0]["label"]  # LABEL_1 indicates spam
-        spam_confidence = spam_result[0]["score"]
-        
-        if spam_label == "LABEL_1":
-            return "spam", f"This is a spam email (Confidence: {spam_confidence:.2f}). No follow-up needed."
+    if spam_label == "LABEL_1":
+        return "spam", 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":
+            return "positive", (f"This email is not spam (Confidence: {spam_confidence:.2f}).\n"
+                                f"Sentiment: {sentiment} (Confidence: {sentiment_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 "positive", (f"This email is not spam (Confidence: {spam_confidence:.2f}).\n"
-                                    f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}). No follow-up needed.")
-            else:
-                return "negative", (f"This email is not spam (Confidence: {spam_confidence:.2f}).\n"
-                                    f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}).\n"
-                                    "<b>Need to Follow-Up</b>: This email is not spam and has negative sentiment.")
-    except Exception as e:
-        return "error", f"An error occurred during analysis: {str(e)}"
+            return "negative", (f"This email is not spam (Confidence: {spam_confidence:.2f}).\n"
+                                f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}).\n"
+                                "**Need to Follow-Up**: This email is not spam and has negative sentiment.")
 
-# Main application function
 def main():
-    # Set title and objective
     st.title("EmailSentry")
     st.write("Aims to perform analysis on incoming emails and to determine whether there is urgency or higher priority for the company to follow-up.")
-    
-    # Initialize session state variables
+
+    # 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 = ""
     if "result_type" not in st.session_state:
         st.session_state.result_type = ""
-    
-    # Instructions section
+
+    # Collapsible instructions
     with st.expander("How to Use", expanded=False):
         st.write("""
         - Type or paste an email into the text box.
@@ -74,11 +51,11 @@ def main():
         - Press 'Analyze Email' to check if it’s spam and analyze its sentiment.
         - Use 'Clear' to reset the input and result.
         """)
-    
+
     # Text area for email input
     email_body = st.text_area("Email Body", value=st.session_state.email_body, height=200, key="email_input")
-    
-    # Define sample emails
+
+    # Sample emails (shortened snippets for button labels)
     sample_spam = """
 Subject: Urgent: Verify Your Account Now!
 Dear Customer,
@@ -89,7 +66,7 @@ Best regards,
 The Security Team
     """
     spam_snippet = "Subject: Urgent: Verify Your Account Now! Dear Customer, We have detected unusual activity..."
-    
+
     sample_not_spam_positive = """
 Subject: Great Experience with HKTV mall
 Dear Sir,
@@ -98,7 +75,7 @@ Best regards,
 Emily
     """
     positive_snippet = "Subject: Great Experience with HKTV mall Dear Sir, I just received my order and I’m really..."
-    
+
     sample_not_spam_negative = """
 Subject: Issue with Recent Delivery
 Dear Support,
@@ -107,9 +84,67 @@ Thanks,
 Sarah
     """
     negative_snippet = "Subject: Issue with Recent Delivery Dear Support, I received my package today, but..."
-    
-    # Display sample buttons
-    st.subheader("Examples")
+
+    # Custom CSS for buttons and result boxes
+    st.markdown("""
+        <style>
+        /* Sample buttons (smaller text) */
+        div.stButton > button[kind="secondary"] {
+            font-size: 12px;
+            padding: 5px 10px;
+            background-color: #f0f0f0;
+            color: #333333;
+            border: 1px solid #cccccc;
+            border-radius: 3px;
+        }
+        /* Analyze Email button (larger, orange) */
+        div.stButton > button[kind="primary"] {
+            background-color: #FF5733;
+            color: white;
+            font-size: 18px;
+            padding: 12px 24px;
+            border: none;
+            border-radius: 5px;
+            margin-right: 10px;
+        }
+        div.stButton > button[kind="primary"]:hover {
+            background-color: #E74C3C;
+        }
+        /* Clear button (gray) */
+        div.stButton > button[kind="secondary"][key="clear"] {
+            background-color: #d3d3d3;
+            color: #333333;
+            font-size: 16px;
+            padding: 10px 20px;
+            border: none;
+            border-radius: 5px;
+        }
+        div.stButton > button[kind="secondary"][key="clear"]:hover {
+            background-color: #b0b0b0;
+        }
+        /* Result boxes */
+        .spam-result {
+            background-color: #ffcccc;
+            padding: 10px;
+            border-radius: 5px;
+            border: 1px solid #ff9999;
+        }
+        .positive-result {
+            background-color: #ccffcc;
+            padding: 10px;
+            border-radius: 5px;
+            border: 1px solid #99cc99;
+        }
+        .negative-result {
+            background-color: #fff3cc;
+            padding: 10px;
+            border-radius: 5px;
+            border: 1px solid #ffcc66;
+        }
+        </style>
+    """, unsafe_allow_html=True)
+
+    # Sample buttons (in columns)
     col1, col2, col3 = st.columns(3)
     with col1:
         if st.button(spam_snippet, key="spam_sample"):
@@ -129,11 +164,11 @@ Sarah
             st.session_state.result = ""
             st.session_state.result_type = ""
             st.rerun()
-    
-    # Action buttons
-    col_analyze, col_clear = st.columns(2)
+
+    # Analyze and Clear buttons (in a row)
+    col_analyze, col_clear = st.columns([1, 1])
     with col_analyze:
-        if st.button("Analyze Email", key="analyze"):
+        if st.button("Analyze Email", key="analyze", type="primary"):
             if email_body:
                 with st.spinner("Analyzing email..."):
                     result_type, result = analyze_email(email_body)
@@ -142,15 +177,14 @@ Sarah
             else:
                 st.session_state.result = "Please enter an email body or select a sample to analyze."
                 st.session_state.result_type = ""
-    
     with col_clear:
         if st.button("Clear", key="clear"):
             st.session_state.email_body = ""
             st.session_state.result = ""
             st.session_state.result_type = ""
             st.rerun()
-    
-    # Display analysis result
+
+    # Display result with styled box
     if st.session_state.result:
         if st.session_state.result_type == "spam":
             st.markdown(f'<div class="spam-result">{st.session_state.result}</div>', unsafe_allow_html=True)
@@ -159,76 +193,7 @@ Sarah
         elif st.session_state.result_type == "negative":
             st.markdown(f'<div class="negative-result">{st.session_state.result}</div>', unsafe_allow_html=True)
         else:
-            st.write(st.session_state.result)
-    
-    # Inject custom CSS without !important, based on your previous working code
-    st.markdown("""
-        <style>
-        /* Style for sample buttons (light grey) */
-        div.stButton > button[kind="secondary"]:not([key="clear"]) {
-            font-size: 12px;
-            padding: 5px 10px;
-            background-color: #f0f0f0;
-            color: #333333;
-            border: 1px solid #cccccc;
-            border-radius: 3px;
-        }
-        /* Analyze Email button (orange) */
-        div.stButton > button[key="analyze"] {
-            background-color: #FF5733;
-            color: white;
-            font-size: 18px;
-            padding: 12px 24px;
-            border: none;
-            border-radius: 5px;
-            width: 100%;
-            height: 50px;
-            box-sizing: border-box;
-            text-align: center;
-        }
-        div.stButton > button[key="analyze"]:hover {
-            background-color: #E74C3C;
-        }
-        /* Clear button (blue) */
-        div.stButton > button[key="clear"] {
-            background-color: #007BFF;
-            color: white;
-            font-size: 18px;
-            padding: 12px 24px;
-            border: none;
-            border-radius: 5px;
-            width: 100%;
-            height: 50px;
-            box-sizing: border-box;
-            text-align: center;
-        }
-        div.stButton > button[key="clear"]:hover {
-            background-color: #0056b3;
-        }
-        /* Result boxes */
-        .spam-result {
-            background-color: #ff3333;
-            color: white;
-            padding: 10px;
-            border-radius: 5px;
-            border: 1px solid #cc0000;
-        }
-        .positive-result {
-            background-color: #ff3333;
-            color: white;
-            padding: 10px;
-            border-radius: 5px;
-            border: 1px solid #cc0000;
-        }
-        .negative-result {
-            background-color: #006633;
-            color: white;
-            padding: 10px;
-            border-radius: 5px;
-            border: 1px solid #004d26;
-        }
-        </style>
-    """, unsafe_allow_html=True)
+            st.write(st.session_state.result)  # For error messages
 
 if __name__ == "__main__":
     main()