Update app.py

app.py

@@ -5,53 +5,55 @@ import numpy as np
 
 # Function to analyze email for spam and sentiment
 def analyze_email(email_body):
-    # Load pre-trained models for spam detection and sentiment analysis
+    # 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)
     tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
 
-    # Step 1: Check if email is spam
+    # 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"]
+    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: 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)
+        # 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"
 
         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.")
         else:
-            # Using HTML <b> tags to ensure bolding
+            # 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.")
 
 # Main application function
 def main():
-    # Set page title
+    # Set the app title to the project name
     st.title("EmailSentry")
-    # Set project objective
+    # 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 variables
+    # Initialize session state to store email input and analysis results
     if "email_body" not in st.session_state:
-        st.session_state.email_body = ""
+        st.session_state.email_body = ""  # Holds the email text
     if "result" not in st.session_state:
-        st.session_state.result = ""
+        st.session_state.result = ""  # Stores the analysis result text
     if "result_type" not in st.session_state:
-        st.session_state.result_type = ""
+        st.session_state.result_type = ""  # Tracks result type (spam, positive, negative)
 
-    # Collapsible instructions section
+    # Add collapsible instructions for user guidance
     with st.expander("How to Use", expanded=False):
         st.write("""
         - Type or paste an email into the text box.
@@ -60,10 +62,10 @@ def main():
         - Use 'Clear' to reset the input and result.
         """)
 
-    # Input text area for email content
+    # Text area where users input or view the email body
     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 buttons
+    # Define sample emails and their snippets for example buttons
     sample_spam = """
 Subject: Urgent: Verify Your Account Now!
 Dear Customer,
@@ -96,7 +98,7 @@ Sarah
     # Custom CSS for styling buttons and result boxes
     st.markdown("""
         <style>
-        /* Sample buttons (smaller text) */
+        /* Style for sample buttons (smaller, light gray) */
         div.stButton > button[kind="secondary"] {
             font-size: 12px;
             padding: 5px 10px;
@@ -105,21 +107,21 @@ Sarah
             border: 1px solid #cccccc;
             border-radius: 3px;
         }
-        /* Analyze Email button (orange) */
+        /* Analyze Email button (orange, matches Clear button size) */
         div.stButton > button[key="analyze"] {
-            background-color: #FF5733;
+            background-color: #FF5733; /* Original orange color */
             color: white;
             font-size: 18px;
             padding: 12px 24px;
             border: none;
             border-radius: 5px;
             width: 100%;
-            height: 50px; /* Fixed height for consistency */
+            height: 50px;
             box-sizing: border-box;
             text-align: center;
         }
         div.stButton > button[key="analyze"]:hover {
-            background-color: #E74C3C;
+            background-color: #E74C3C; /* Darker orange on hover */
         }
         /* Clear button (gray, aligned with Analyze) */
         div.stButton > button[key="clear"] {
@@ -130,80 +132,88 @@ Sarah
             border: none;
             border-radius: 5px;
             width: 100%;
-            height: 50px; /* Fixed height for consistency */
+            height: 50px;
             box-sizing: border-box;
             text-align: center;
         }
         div.stButton > button[key="clear"]:hover {
-            background-color: #b0b0b0;
+            background-color: #b0b0b0; /* Darker gray on hover */
         }
-        /* Result boxes */
+        /* Result boxes with updated colors */
         .spam-result {
-            background-color: #ffdddd; /* Softer red */
+            background-color: #ff3333; /* Red for no follow-up (spam) */
+            color: white;
             padding: 10px;
             border-radius: 5px;
-            border: 1px solid #ffaaaa;
+            border: 1px solid #cc0000;
         }
         .positive-result {
-            background-color: #d4edda; /* Softer green */
+            background-color: #ff3333; /* Red for no follow-up (positive) */
+            color: white;
             padding: 10px;
             border-radius: 5px;
-            border: 1px solid #a3d7a9;
+            border: 1px solid #cc0000;
         }
         .negative-result {
-            background-color: #fff4e6; /* Softer orange */
+            background-color: #006633; /* Dark green for follow-up needed */
+            color: white;
             padding: 10px;
             border-radius: 5px;
-            border: 1px solid #ffd6a5;
+            border: 1px solid #004d26;
         }
         </style>
     """, unsafe_allow_html=True)
 
-    # Subheading for sample buttons
+    # Subheading to label the sample email buttons
     st.subheader("Examples")
 
-    # Sample buttons layout (3 columns)
+    # 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()
 
-    # Action buttons layout (Analyze and Clear)
+    # Layout for action buttons (Analyze and Clear) in 2 columns
     col_analyze, col_clear = st.columns(2)
     with col_analyze:
-        # Removed type="primary" to rely on custom CSS
+        # Button to trigger email analysis
         if st.button("Analyze Email", key="analyze"):
             if email_body:
-                with st.spinner("Analyzing email..."):
+                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 analysis result in styled boxes
+    # 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)
@@ -212,7 +222,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)  # For error messages
+            st.write(st.session_state.result)  # Display error messages without styling
 
 # Run the app
 if __name__ == "__main__":