Spaces:
Running
Running
| import gradio as gr | |
| from transformers import pipeline | |
| import re | |
| print("Loading model...") | |
| classifier = pipeline( | |
| "text-classification", | |
| model="Umranz/modernbert-enron-spam", | |
| tokenizer="Umranz/modernbert-enron-spam" | |
| ) | |
| print("✅ Model loaded!") | |
| SPAM_KEYWORDS = [ | |
| "free", "win", "winner", "won", "prize", "claim", "urgent", | |
| "congratulations", "offer", "limited", "act now", "click here", | |
| "buy now", "cheap", "discount", "guaranteed", "money", "cash", | |
| "reward", "selected", "exclusive", "deal", "risk free", "no cost", | |
| "bonus", "gift", "million", "billion", "dollar", "$", "€", "£", | |
| "100%", "password", "verify", "account", "suspended", "compromised", | |
| "medication", "pharmacy", "prescription", "weight loss", "earn", | |
| "income", "investment", "crypto", "bitcoin", "profit" | |
| ] | |
| def highlight_keywords(text): | |
| found = [] | |
| text_lower = text.lower() | |
| for kw in SPAM_KEYWORDS: | |
| if kw in text_lower: | |
| found.append(f"`{kw}`") | |
| return found | |
| def predict_single(subject, body, threshold): | |
| if not body.strip(): | |
| return "⚠️ Please enter email body.", "", "", "" | |
| combined = f"Subject: {subject}\n\n{body}" if subject.strip() else body | |
| result = classifier(combined, truncation=True, max_length=512)[0] | |
| label = result["label"] | |
| score = result["score"] * 100 | |
| if score < threshold: | |
| verdict = ( | |
| f"## ⚠️ Low Confidence ({score:.1f}%)\n\n" | |
| f"Best guess: **{label}** but below your threshold of {threshold}%.\n\n" | |
| f"Try adding more email content for better accuracy." | |
| ) | |
| elif label == "SPAM": | |
| bar = "█" * int(score // 5) + "░" * (20 - int(score // 5)) | |
| verdict = ( | |
| f"## 🚨 SPAM DETECTED\n\n" | |
| f"**Confidence:** {score:.1f}%\n\n" | |
| f"{bar} {score:.1f}%\n\n" | |
| f"⚠️ This email shows strong spam characteristics.\n\n" | |
| f"**Recommendation:** Do not click any links or reply." | |
| ) | |
| else: | |
| bar = "█" * int(score // 5) + "░" * (20 - int(score // 5)) | |
| verdict = ( | |
| f"## ✅ LEGITIMATE EMAIL (HAM)\n\n" | |
| f"**Confidence:** {score:.1f}%\n\n" | |
| f"{bar} {score:.1f}%\n\n" | |
| f"This email appears to be legitimate." | |
| ) | |
| found_keywords = highlight_keywords(f"{subject} {body}") | |
| if found_keywords: | |
| keyword_analysis = ( | |
| f"## 🔍 Spam Keywords Found ({len(found_keywords)})\n\n" | |
| f"{', '.join(found_keywords)}\n\n" | |
| f"{'⚠️ High keyword density — be cautious!' if len(found_keywords) > 5 else '⚡ Low keyword density.'}" | |
| ) | |
| else: | |
| keyword_analysis = "## 🔍 Spam Keywords\n\n✅ No common spam keywords detected." | |
| risk_level = ( | |
| "🔴 HIGH RISK" if label == "SPAM" and score > 80 else | |
| "🟠 MEDIUM RISK" if label == "SPAM" and score > 50 else | |
| "🟡 LOW RISK" if label == "SPAM" else | |
| "🟢 SAFE" | |
| ) | |
| risk_summary = ( | |
| f"## 🛡️ Risk Assessment\n\n" | |
| f"**Risk Level:** {risk_level}\n\n" | |
| f"**Spam Score:** {score:.1f}%\n\n" | |
| f"**Keywords Found:** {len(found_keywords)}\n\n" | |
| f"**Model Verdict:** {label}" | |
| ) | |
| if label == "SPAM": | |
| tips = ( | |
| f"## 💡 Safety Tips\n\n" | |
| f"- 🚫 Don't click any links in this email\n" | |
| f"- 🚫 Don't reply or provide personal info\n" | |
| f"- 🚫 Don't download any attachments\n" | |
| f"- ✅ Mark as spam in your email client\n" | |
| f"- ✅ Block the sender" | |
| ) | |
| else: | |
| tips = ( | |
| f"## 💡 Tips\n\n" | |
| f"- ✅ Email appears safe to read\n" | |
| f"- ✅ Always verify sender identity\n" | |
| f"- ✅ Be cautious with attachments\n" | |
| f"- ⚡ Even legitimate emails can have phishing links" | |
| ) | |
| return verdict, keyword_analysis, risk_summary, tips | |
| def predict_batch(emails_text, threshold): | |
| if not emails_text.strip(): | |
| return "⚠️ Please enter at least one email." | |
| lines = [l.strip() for l in emails_text.strip().split("\n") if l.strip()] | |
| results = classifier(lines, truncation=True, max_length=512) | |
| spam_count = 0 | |
| ham_count = 0 | |
| output = "## 📋 Batch Analysis Results\n\n" | |
| for line, res in zip(lines, results): | |
| label = res["label"] | |
| score = res["score"] * 100 | |
| emoji = "🚨" if label == "SPAM" else "✅" | |
| bar = "█" * int(score // 10) + "░" * (10 - int(score // 10)) | |
| if label == "SPAM": | |
| spam_count += 1 | |
| else: | |
| ham_count += 1 | |
| output += ( | |
| f"{emoji} **{label}** ({score:.1f}%) {bar}\n" | |
| f"> {line[:80]}{'...' if len(line) > 80 else ''}\n\n" | |
| ) | |
| total = len(lines) | |
| output += ( | |
| f"---\n" | |
| f"## 📊 Summary\n\n" | |
| f"- **Total:** {total} emails\n" | |
| f"- **🚨 Spam:** {spam_count} ({spam_count/total*100:.1f}%)\n" | |
| f"- **✅ Ham:** {ham_count} ({ham_count/total*100:.1f}%)" | |
| ) | |
| return output | |
| with gr.Blocks(title="SpamShield AI") as demo: | |
| gr.Markdown(""" | |
| # 🛡️ SpamShield AI — Email Spam Classifier | |
| **ModernBERT (2024) fine-tuned on 31K Enron Emails** | |
| > Paste any email to instantly detect if it's spam or legitimate. | |
| """) | |
| with gr.Tabs(): | |
| with gr.Tab("📧 Single Email"): | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| subject_input = gr.Textbox( | |
| lines=1, | |
| placeholder="e.g. You Won $1,000,000!!!", | |
| label="📌 Email Subject" | |
| ) | |
| body_input = gr.Textbox( | |
| lines=8, | |
| placeholder="Paste the email body here...", | |
| label="📝 Email Body" | |
| ) | |
| threshold_slider = gr.Slider( | |
| minimum=10, | |
| maximum=90, | |
| value=50, | |
| step=5, | |
| label="⚙️ Confidence Threshold (%)", | |
| info="Predictions below this % show a low-confidence warning" | |
| ) | |
| analyze_btn = gr.Button( | |
| "🔍 Analyze Email", | |
| variant="primary", | |
| size="lg" | |
| ) | |
| with gr.Column(scale=3): | |
| verdict_output = gr.Markdown(label="Verdict") | |
| keyword_output = gr.Markdown(label="Keyword Analysis") | |
| with gr.Row(): | |
| risk_output = gr.Markdown(label="Risk Assessment") | |
| tips_output = gr.Markdown(label="Safety Tips") | |
| gr.Examples( | |
| examples=[ | |
| ["You Won $1,000,000!!!", "Congratulations! You have been selected as our lucky winner. Click here to claim your prize now. Limited time offer!", 50], | |
| ["Meeting Notes", "Hi team, please find attached the meeting notes from yesterday's standup. Let me know if I missed anything.", 50], | |
| ["URGENT: Verify Account", "Your account has been compromised. Verify your credentials immediately to avoid suspension. Click the link below.", 50], | |
| ["Quick Question", "Hey, are you free for a call tomorrow at 3pm? I wanted to discuss the project timeline.", 50], | |
| ["Exclusive Pharma Deal", "Buy cheap medications online! No prescription needed! 90% discount. Act now before offer expires!", 50], | |
| ], | |
| inputs=[subject_input, body_input, threshold_slider], | |
| label="💡 Try These Examples" | |
| ) | |
| analyze_btn.click( | |
| fn=predict_single, | |
| inputs=[subject_input, body_input, threshold_slider], | |
| outputs=[verdict_output, keyword_output, risk_output, tips_output] | |
| ) | |
| with gr.Tab("📋 Batch Emails"): | |
| gr.Markdown("Enter one email per line for bulk analysis.") | |
| batch_input = gr.Textbox( | |
| lines=10, | |
| placeholder=( | |
| "Congratulations you won $1M click here now!\n" | |
| "Hi John, meeting is rescheduled to 3pm.\n" | |
| "Buy cheap meds online no prescription needed!\n" | |
| "Please review the attached quarterly report." | |
| ), | |
| label="📝 Emails (one per line)" | |
| ) | |
| batch_threshold = gr.Slider( | |
| minimum=10, | |
| maximum=90, | |
| value=50, | |
| step=5, | |
| label="⚙️ Confidence Threshold (%)" | |
| ) | |
| batch_btn = gr.Button("🔍 Analyze All", variant="primary", size="lg") | |
| batch_output = gr.Markdown(label="Batch Results") | |
| gr.Examples( | |
| examples=[[ | |
| "Congratulations you won $1M click here now!\n" | |
| "Hi John, meeting is rescheduled to 3pm.\n" | |
| "Buy cheap meds online no prescription needed!\n" | |
| "Please review the attached quarterly report.\n" | |
| "URGENT your account will be suspended verify now!\n" | |
| "Are you free for lunch this Friday?" | |
| ]], | |
| inputs=[batch_input], | |
| label="💡 Example" | |
| ) | |
| batch_btn.click( | |
| fn=predict_batch, | |
| inputs=[batch_input, batch_threshold], | |
| outputs=[batch_output] | |
| ) | |
| with gr.Tab("ℹ️ About"): | |
| gr.Markdown(""" | |
| ## 🛡️ SpamShield AI | |
| ### 🧠 Model | |
| - **Architecture:** ModernBERT (2024) — state of the art | |
| - **Fine-tuned on:** Enron Spam Dataset (31,716 emails) | |
| - **Input:** Subject + Body combined for maximum context | |
| - **Context window:** 512 tokens | |
| ### 📊 Performance | |
| | Metric | Score | | |
| |---|---| | |
| | Accuracy | ~98% | | |
| | F1 Score | ~98% | | |
| | Precision | ~98% | | |
| | Recall | ~98% | | |
| ### 🔍 Features | |
| - ✅ Single email analysis with risk assessment | |
| - ✅ Spam keyword highlighting | |
| - ✅ Confidence threshold control | |
| - ✅ Batch analysis for multiple emails | |
| - ✅ Safety tips for spam emails | |
| ### 📧 Dataset | |
| - **Enron Email Dataset** — Real corporate emails | |
| - 31,716 training samples | |
| - Industry standard benchmark | |
| ### 👨💻 Built By | |
| Umranz — [HuggingFace Profile](https://huggingface.co/Umranz) | |
| > ⚠️ This tool is for educational purposes. | |
| > Always use proper email security tools in production. | |
| """) | |
| demo.launch(theme=gr.themes.Soft()) |