Update app.py

app.py

@@ -7,18 +7,15 @@ def main():
     # Load the spam detection pipeline
     spam_pipeline = pipeline("text-classification", model="cybersectony/phishing-email-detection-distilbert_v2.4.1")
 
-    # Load the sentiment model and tokenizer directly
+    # Load the sentiment model and tokenizer
     sentiment_model = AutoModelForSequenceClassification.from_pretrained("ISOM5240GP4/email_sentiment", num_labels=2)
     tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
 
-    # Title and description
     st.title("Email Analysis Tool")
     st.write("Enter an email body below to check if it's spam and analyze its sentiment.")
 
-    # Text area for email input
     email_body = st.text_area("Email Body", height=200)
 
-    # Button to trigger analysis
     if st.button("Analyze Email"):
         if email_body:
             # Step 1: Check if the email is spam
@@ -26,8 +23,8 @@ def main():
             spam_label = spam_result[0]["label"]
             spam_confidence = spam_result[0]["score"]
 
-            # Assuming "POSITIVE" means spam/phishing (adjust if incorrect)
-            if spam_label == "POSITIVE":
+            # Check if label is 'LABEL_1' for spam
+            if spam_label == "LABEL_1":
                 st.write(f"This is a spam email (Confidence: {spam_confidence:.2f}). No follow-up needed.")
             else:
                 # Step 2: Analyze sentiment for non-spam emails
@@ -38,13 +35,12 @@ def main():
                 sentiment_index = np.argmax(predictions)
                 sentiment_confidence = predictions[0][sentiment_index]
 
-                # Map index to sentiment (1 = positive, 0 = negative)
                 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:  # Negative sentiment
+                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.**")