Spaces:
Running
Running
| import streamlit as st | |
| from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer | |
| import torch | |
| import numpy as np | |
| # Function to analyze email (no caching, models loaded each time) | |
| def analyze_email(email_body): | |
| """Analyzes an email for spam and sentiment, returning result type and message.""" | |
| 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") | |
| # this is to handle if user does not input anything | |
| 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: | |
| # Step 2: Analyze sentiment for non-spam emails | |
| inputs = tokenizer(email_body, padding=True, truncation=True, return_tensors='pt') | |
| # Pass the tokenized inputs to the sentiment model | |
| # **inputs unpacks the dictionary of inputs to the model | |
| 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] | |
| # Convert numeric prediction to human-readable sentiment | |
| # If index is 1, it's Positive, otherwise Negative | |
| 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 title and objective | |
| st.title("EmailSentry") | |
| 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 | |
| if "email_body" not in st.session_state: | |
| st.session_state.email_body = "" | |
| if "result" not in st.session_state: | |
| st.session_state.result = "" | |
| if "result_type" not in st.session_state: | |
| 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. | |
| - Alternatively, click one of the sample buttons to load a predefined email. | |
| - Press 'Analyze Email' to check if it’s spam and analyze its sentiment. | |
| - Use 'Clear' to reset the input and result. | |
| """) | |
| # Text area for email input | |
| email_body = st.text_area("Email", value=st.session_state.email_body, height=200, key="email_input") | |
| # Define sample emails | |
| 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 Experience with HKTV mall | |
| Dear Sir, | |
| I just received my order and I’m really impressed with the speed of the delivery. Keep up the good work. | |
| 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, | |
| 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..." | |
| # 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 | |
| col_analyze, col_clear = st.columns(2) | |
| with col_analyze: | |
| if st.button("Analyze Email", key="analyze", type="primary"): | |
| 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 = "" | |
| with col_clear: | |
| 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 | |
| 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 with updated result block colors | |
| st.markdown(""" | |
| <style> | |
| /* Sample buttons (light grey, small) */ | |
| div.stButton > button[kind="secondary"] { | |
| font-size: 12px; | |
| padding: 5px 10px; | |
| background-color: #f0f0f0; | |
| color: #333333; | |
| border: 1px solid #cccccc; | |
| border-radius: 3px; | |
| } | |
| /* Analyze Email button (orange, larger) */ | |
| div.stButton > button[kind="primary"] { | |
| background-color: #FF5733; | |
| color: white; | |
| font-size: 18px; | |
| padding: 12px 24px; | |
| border: none; | |
| border-radius: 5px; | |
| margin-right: 10px; | |
| } | |
| div.stButton > button[kind="primary"]:hover { | |
| background-color: #E74C3C; | |
| } | |
| /* Clear button (blue) */ | |
| div.stButton > button[kind="secondary"][key="clear"] { | |
| background-color: #007BFF; | |
| color: white; | |
| font-size: 18px; | |
| padding: 12px 24px; | |
| border: none; | |
| border-radius: 5px; | |
| } | |
| div.stButton > button[kind="secondary"][key="clear"]:hover { | |
| background-color: #0056b3; | |
| } | |
| /* Result boxes: Red for no follow-up, Green for follow-up */ | |
| .spam-result { | |
| background-color: #ff3333; /* Red for no follow-up */ | |
| color: white; | |
| padding: 10px; | |
| border-radius: 5px; | |
| border: 1px solid #cc0000; | |
| } | |
| .positive-result { | |
| background-color: #ff3333; /* Red for no follow-up */ | |
| color: white; | |
| padding: 10px; | |
| border-radius: 5px; | |
| border: 1px solid #cc0000; | |
| } | |
| .negative-result { | |
| background-color: #33cc33; /* Green for follow-up needed */ | |
| color: white; | |
| padding: 10px; | |
| border-radius: 5px; | |
| border: 1px solid #009900; | |
| } | |
| </style> | |
| """, unsafe_allow_html=True) | |
| if __name__ == "__main__": | |
| main() |