Update app.py

app.py

@@ -3,57 +3,70 @@ from transformers import pipeline, AutoModelForSequenceClassification, AutoToken
 import torch
 import numpy as np
 
-# Function to analyze email for spam and sentiment
-def analyze_email(email_body):
-    # Load pre-trained models: spam detection and sentiment analysis
-    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)
+# 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)
     tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+    return model, tokenizer
 
-    # Step 1: Check if the email is spam using the spam detection model
-    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":
-        # If spam, return type "spam" and a message indicating no follow-up
-        return "spam", f"This is a spam email (Confidence: {spam_confidence:.2f}). No follow-up needed."
-    else:
-        # Step 2: For non-spam emails, analyze sentiment (positive/negative)
-        inputs = tokenizer(email_body, padding=True, truncation=True, return_tensors='pt')  # Tokenize input
-        outputs = sentiment_model(**inputs)  # Get model predictions
-        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)  # Apply softmax for probabilities
-        predictions = predictions.cpu().detach().numpy()  # Convert to numpy array
-        sentiment_index = np.argmax(predictions)  # Get the predicted sentiment (0 = negative, 1 = positive)
-        sentiment_confidence = predictions[0][sentiment_index]
-        sentiment = "Positive" if sentiment_index == 1 else "Negative"
+# Initialize models
+spam_pipeline = load_spam_pipeline()
+sentiment_model, tokenizer = load_sentiment_model()
 
-        if sentiment == "Positive":
-            # If positive sentiment, no follow-up needed
-            return "positive", (f"This email is not spam (Confidence: {spam_confidence:.2f}).\n"
-                                f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}). No follow-up needed.")
+# 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."
         else:
-            # If negative sentiment, mark as needing follow-up with bolded text
-            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.")
+            # 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)}"
 
 # Main application function
 def main():
-    # Set the app title to the project name
+    # Set title and objective
     st.title("EmailSentry")
-    # Display the project objective
     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 to store email input and analysis results
+    
+    # Initialize session state variables
     if "email_body" not in st.session_state:
-        st.session_state.email_body = ""  # Holds the email text
+        st.session_state.email_body = ""
     if "result" not in st.session_state:
-        st.session_state.result = ""  # Stores the analysis result text
+        st.session_state.result = ""
     if "result_type" not in st.session_state:
-        st.session_state.result_type = ""  # Tracks result type (spam, positive, negative)
-
-    # Add collapsible instructions for user guidance
+        st.session_state.result_type = ""
+    
+    # Instructions section
     with st.expander("How to Use", expanded=False):
         st.write("""
         - Type or paste an email into the text box.
@@ -61,11 +74,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 where users input or view the email body
+    
+    # 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 and their snippets for example buttons
+    
+    # Define sample emails
     sample_spam = """
 Subject: Urgent: Verify Your Account Now!
 Dear Customer,
@@ -76,7 +89,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,
@@ -85,7 +98,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,
@@ -94,13 +107,69 @@ Thanks,
 Sarah
     """
     negative_snippet = "Subject: Issue with Recent Delivery Dear Support, I received my package today, but..."
-
-    # Custom CSS for styling buttons and result boxes with higher specificity
-    # Replace your existing st.markdown CSS block with this
+    
+    # Display sample buttons
+    st.subheader("Examples")
+    col1, col2, col3 = st.columns(3)
+    with col1:
+        if st.button(spam_snippet, key="spam_sample"):
+            st.session_state.email_body = sample_spam
+            st.session_state.result = ""
+            st.session_state.result_type = ""
+            st.rerun()
+    with col2:
+        if st.button(positive_snippet, key="positive_sample"):
+            st.session_state.email_body = sample_not_spam_positive
+            st.session_state.result = ""
+            st.session_state.result_type = ""
+            st.rerun()
+    with col3:
+        if st.button(negative_snippet, key="negative_sample"):
+            st.session_state.email_body = sample_not_spam_negative
+            st.session_state.result = ""
+            st.session_state.result_type = ""
+            st.rerun()
+    
+    # Action buttons with custom styling wrappers
+    col_analyze, col_clear = st.columns(2)
+    with col_analyze:
+        st.markdown('<div id="analyze-button">', unsafe_allow_html=True)
+        if st.button("Analyze Email", key="analyze"):
+            if email_body:
+                with st.spinner("Analyzing email..."):
+                    result_type, result = analyze_email(email_body)
+                    st.session_state.result = result
+                    st.session_state.result_type = result_type
+            else:
+                st.session_state.result = "Please enter an email body or select a sample to analyze."
+                st.session_state.result_type = ""
+        st.markdown('</div>', unsafe_allow_html=True)
+    
+    with col_clear:
+        st.markdown('<div id="clear-button">', unsafe_allow_html=True)
+        if st.button("Clear", key="clear"):
+            st.session_state.email_body = ""
+            st.session_state.result = ""
+            st.session_state.result_type = ""
+            st.rerun()
+        st.markdown('</div>', unsafe_allow_html=True)
+    
+    # Display analysis result
+    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)
+        elif st.session_state.result_type == "positive":
+            st.markdown(f'<div class="positive-result">{st.session_state.result}</div>', unsafe_allow_html=True)
+        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 for styling
     st.markdown("""
         <style>
-        /* Ensure high specificity for Analyze Email button (orange) */
-        div[data-testid="stButton"] > button[key="analyze"] {
+        /* Style for Analyze Email button */
+        #analyze-button button {
             background-color: #FF5733 !important; /* Orange */
             color: white !important;
             font-size: 18px !important;
@@ -113,12 +182,12 @@ Sarah
             text-align: center !important;
             margin: 0 !important;
         }
-        div[data-testid="stButton"] > button[key="analyze"]:hover {
-            background-color: #E74C3C !important; /* Darker orange on hover */
+        #analyze-button button:hover {
+            background-color: #E74C3C !important; /* Darker orange */
         }
-    
-        /* Ensure high specificity for Clear button (blue) */
-        div[data-testid="stButton"] > button[key="clear"] {
+        
+        /* Style for Clear button */
+        #clear-button button {
             background-color: #007BFF !important; /* Blue */
             color: white !important;
             font-size: 18px !important;
@@ -131,37 +200,37 @@ Sarah
             text-align: center !important;
             margin: 0 !important;
         }
-        div[data-testid="stButton"] > button[key="clear"]:hover {
-            background-color: #0056b3 !important; /* Darker blue on hover */
+        #clear-button button:hover {
+            background-color: #0056b3 !important; /* Darker blue */
         }
-    
-        /* Style for sample buttons (light gray) */
-        div[data-testid="stButton"] > button[kind="secondary"]:not([key="clear"]):not([key="analyze"]) {
+        
+        /* Style for sample buttons */
+        div[data-testid="stButton"] button:not([key="analyze"]):not([key="clear"]) {
             font-size: 12px !important;
             padding: 5px 10px !important;
-            background-color: #f0f0f0 !important;
+            background-color: #f0f0f0 !important; /* Light gray */
             color: #333333 !important;
             border: 1px solid #cccccc !important;
             border-radius: 3px !important;
         }
-    
-        /* Result boxes (unchanged) */
+        
+        /* Result boxes */
         .spam-result {
-            background-color: #ff3333 !important; /* Red for no follow-up (spam) */
+            background-color: #ff3333 !important; /* Red */
             color: white !important;
             padding: 10px !important;
             border-radius: 5px !important;
             border: 1px solid #cc0000 !important;
         }
         .positive-result {
-            background-color: #ff3333 !important; /* Red for no follow-up (positive) */
+            background-color: #ff3333 !important; /* Red */
             color: white !important;
             padding: 10px !important;
             border-radius: 5px !important;
             border: 1px solid #cc0000 !important;
         }
         .negative-result {
-            background-color: #006633 !important; /* Dark green for follow-up needed */
+            background-color: #006633 !important; /* Dark green */
             color: white !important;
             padding: 10px !important;
             border-radius: 5px !important;
@@ -170,66 +239,5 @@ Sarah
         </style>
     """, unsafe_allow_html=True)
 
-    # Subheading to label the sample email buttons
-    st.subheader("Examples")
-
-    # Layout for sample buttons in 3 columns
-    col1, col2, col3 = st.columns(3)
-    with col1:
-        # Button to load spam sample
-        if st.button(spam_snippet, key="spam_sample"):
-            st.session_state.email_body = sample_spam
-            st.session_state.result = ""
-            st.session_state.result_type = ""
-            st.rerun()
-    with col2:
-        # Button to load positive non-spam sample
-        if st.button(positive_snippet, key="positive_sample"):
-            st.session_state.email_body = sample_not_spam_positive
-            st.session_state.result = ""
-            st.session_state.result_type = ""
-            st.rerun()
-    with col3:
-        # Button to load negative non-spam sample
-        if st.button(negative_snippet, key="negative_sample"):
-            st.session_state.email_body = sample_not_spam_negative
-            st.session_state.result = ""
-            st.session_state.result_type = ""
-            st.rerun()
-
-    # Layout for action buttons (Analyze and Clear) in 2 columns
-    col_analyze, col_clear = st.columns(2)
-    with col_analyze:
-        # Button to trigger email analysis (no type="primary" to rely on CSS)
-        if st.button("Analyze Email", key="analyze"):
-            if email_body:
-                with st.spinner("Analyzing email..."):  # Show spinner during processing
-                    result_type, result = analyze_email(email_body)
-                    st.session_state.result = result
-                    st.session_state.result_type = result_type
-            else:
-                # Error message if no email is provided
-                st.session_state.result = "Please enter an email body or select a sample to analyze."
-                st.session_state.result_type = ""
-    with col_clear:
-        # Button to reset the app state
-        if st.button("Clear", key="clear"):
-            st.session_state.email_body = ""
-            st.session_state.result = ""
-            st.session_state.result_type = ""
-            st.rerun()
-
-    # Display the analysis result in styled boxes based on result type
-    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)
-        elif st.session_state.result_type == "positive":
-            st.markdown(f'<div class="positive-result">{st.session_state.result}</div>', unsafe_allow_html=True)
-        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)  # Display error messages without styling
-
-# Run the app
 if __name__ == "__main__":
     main()