Spaces:
Running
Running
| import streamlit as st | |
| from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer | |
| import torch | |
| import numpy as np | |
| 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 | |
| sentiment_model = AutoModelForSequenceClassification.from_pretrained("ISOM5240GP4/email_sentiment", num_labels=2) | |
| tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased") | |
| st.title("Email Analysis Tool") | |
| st.write("Enter an email body below or select a sample to analyze its spam status and sentiment.") | |
| # Initialize session state for the text area | |
| if "email_body" not in st.session_state: | |
| st.session_state.email_body = "" | |
| # 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 | |
| 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 | |
| """ | |
| 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 | |
| """ | |
| 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 | |
| """ | |
| # Buttons for sample emails (in columns for better layout) | |
| col1, col2, col3 = st.columns(3) | |
| with col1: | |
| if st.button("Spam Email"): | |
| st.session_state.email_body = sample_spam | |
| st.rerun() # Rerun to update the text area | |
| with col2: | |
| if st.button("Not Spam, Positive"): | |
| st.session_state.email_body = sample_not_spam_positive | |
| st.rerun() | |
| with col3: | |
| if st.button("Not Spam, Negative"): | |
| st.session_state.email_body = sample_not_spam_negative | |
| st.rerun() | |
| # Button to trigger analysis | |
| if st.button("Analyze Email"): | |
| if email_body: | |
| # Step 1: Check if the email is spam | |
| spam_result = spam_pipeline(email_body) | |
| spam_label = spam_result[0]["label"] | |
| spam_confidence = spam_result[0]["score"] | |
| # 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 | |
| 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 |