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 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") | |
| st.write("Enter an email body below to check if it's spam and analyze its sentiment.") | |
| # Text area for email input | |
| email_body = st.text_area("Email Body", height=200) | |
| # 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"] | |
| # 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: 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 == "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: # Negative sentiment | |
| 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 to analyze.") | |
| if __name__ == "__main__": | |
| main() |