Spaces:
Sleeping
Sleeping
| 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.") | |
| if "email_body" not in st.session_state: | |
| st.session_state.email_body = "" | |
| email_body = st.text_area("Email Body", value=st.session_state.email_body, height=200, key="email_input") | |
| # Sample emails | |
| sample_spam = """[Spam email content]""" | |
| sample_not_spam_positive = """[Positive email content]""" | |
| sample_not_spam_negative = """[Negative email content]""" | |
| col1, col2, col3 = st.columns(3) | |
| with col1: | |
| if st.button("Spam Email"): | |
| st.session_state.email_body = sample_spam | |
| st.rerun() | |
| 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() | |
| if st.button("Analyze Email"): | |
| if email_body: | |
| spam_result = spam_pipeline(email_body) | |
| spam_label = spam_result[0]["label"] | |
| spam_confidence = spam_result[0]["score"] | |
| if spam_label == "LABEL_1": | |
| st.write(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": | |
| 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: | |
| 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.**") | |
| else: | |
| st.write("Please enter an email body or select a sample to analyze.") | |
| if __name__ == "__main__": | |
| main() |