app.py

import streamlit as st
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
import torch
import numpy as np

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")

    spam_result = spam_pipeline(email_body)
    spam_label = spam_result[0]["label"]
    spam_confidence = spam_result[0]["score"]

    if spam_label == "LABEL_1":
        return 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 (f"This email is not spam (Confidence: {spam_confidence:.2f}).\n"
                    f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}). No follow-up needed.")
        else:
            return (f"This email is not spam (Confidence: {spam_confidence:.2f}).\n"
                    f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}).\n"
                    "**This email needs follow-up as it is not spam and has negative sentiment.**")

def main():
    st.title("Email Analysis Tool")
    st.write("Enter an email body below or click a sample to analyze its spam status and sentiment.")

    # 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 = ""

    # Text area for email input
    email_body = st.text_area("Email Body", value=st.session_state.email_body, height=200, key="email_input")

    # Sample emails (shortened snippets for button labels)
    sample_spam = """
Subject: Urgent: Verify Your Account Now!
Dear Customer,
We have detected unusual activity on your account. To prevent suspension, please verify your login details immediately by clicking the link below:
[Click Here to Verify](http://totally-legit-site.com/verify)
Failure to verify within 24 hours will result in your account being locked. This is for your security.
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 News About Your Project!
Hi Team,
I just wanted to let you know that the project is progressing wonderfully! Everyone’s efforts are paying off, and we’re ahead of schedule. Keep up the fantastic work!
Best,
Alex
    """
    positive_snippet = "Subject: Great News About Your Project! Hi Team, I just wanted to let you know..."

    sample_not_spam_negative = """
Subject: Issue with Recent Delivery
Dear Support,
I received my package today, but it was damaged, and two items were missing. This is really frustrating—please let me know how we can resolve this as soon as possible.
Thanks,
Sarah
    """
    negative_snippet = "Subject: Issue with Recent Delivery Dear Support, I received my package today, but..."

    # Custom CSS for buttons
    st.markdown("""
        <style>
        /* Style for sample buttons (smaller text) */
        div.stButton > button[kind="secondary"] {
            font-size: 12px; /* Smaller text */
            padding: 5px 10px; /* Smaller padding */
            background-color: #f0f0f0; /* Light gray background */
            color: #333333; /* Darker text */
            border: 1px solid #cccccc;
            border-radius: 3px;
        }
        /* Style for Analyze Email button (larger, colored) */
        div.stButton > button[kind="primary"] {
            background-color: #FF5733; /* Orange color */
            color: white;
            font-size: 18px; /* Larger text */
            padding: 12px 24px; /* Larger padding */
            border: none;
            border-radius: 5px;
            display: block;
            margin-top: 15px;
        }
        div.stButton > button[kind="primary"]:hover {
            background-color: #E74C3C; /* Darker orange on hover */
        }
        </style>
    """, unsafe_allow_html=True)

    # Buttons with sample content (in columns)
    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.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.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.rerun()

    # Analyze Email button (distinct style)
    if st.button("Analyze Email", key="analyze", type="primary"):
        if email_body:
            st.session_state.result = analyze_email(email_body)
        else:
            st.session_state.result = "Please enter an email body or select a sample to analyze."

    # Display result
    if st.session_state.result:
        st.write(st.session_state.result)

if __name__ == "__main__":
    main()