Update app.py

app.py

@@ -1,10 +1,15 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
+import torch
+import numpy as np
 
 def main():
-    # Load the models
+    # Load the spam detection pipeline
     spam_pipeline = pipeline("text-classification", model="cybersectony/phishing-email-detection-distilbert_v2.4.1")
-    sentiment_pipeline = pipeline("text-classification", model="ISOM5240GP4/email_sentiment")
+
+    # Load the sentiment model and tokenizer directly
+    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")
@@ -21,21 +26,27 @@ def main():
             spam_label = spam_result[0]["label"]
             spam_confidence = spam_result[0]["score"]
 
-            # If it's spam, display result and stop
-            if spam_label == "POSITIVE":  # Assuming "POSITIVE" means spam/phishing (check model docs)
+            # Assuming "POSITIVE" means spam/phishing (adjust if incorrect)
+            if spam_label == "POSITIVE":
                 st.write(f"This is a spam email (Confidence: {spam_confidence:.2f}). No follow-up needed.")
             else:
-                # Step 2: If not spam, analyze sentiment
-                sentiment_result = sentiment_pipeline(email_body)
-                sentiment_label = sentiment_result[0]["label"]
-                sentiment_confidence = sentiment_result[0]["score"]
+                # 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]
+
+                # Map index to sentiment (1 = positive, 0 = negative)
+                sentiment = "Positive" if sentiment_index == 1 else "Negative"
 
-                if sentiment_label == "POSITIVE":
+                if sentiment == "Positive":
                     st.write(f"This email is not spam (Confidence: {spam_confidence:.2f}).")
-                    st.write(f"Sentiment: Positive (Confidence: {sentiment_confidence:.2f}). No follow-up needed.")
-                else:  # Assuming "NEGATIVE" for negative sentiment
+                    st.write(f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}). No follow-up needed.")
+                else:  # Negative sentiment
                     st.write(f"This email is not spam (Confidence: {spam_confidence:.2f}).")
-                    st.write(f"Sentiment: Negative (Confidence: {sentiment_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 to analyze.")