app.py

#!/usr/bin/env python3
"""
Naive Bayes Spam Detector with Gradio UI
- Trains a simple sklearn Pipeline (CountVectorizer + MultinomialNB)
- Provides a text box to classify a single message
- Shows predicted label and spam probability
- Works offline (ships with a tiny embedded toy dataset) but can also read a local CSV if available

Optional local dataset:
If a file named 'sms_spam.csv' exists in the same folder with columns:
    text,label
and labels in {'spam','ham'}
the app will train on that instead of the tiny toy set.

Author: ChatGPT (for Mike)
"""
from __future__ import annotations

import os
import csv
from typing import Tuple, List

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.metrics import accuracy_score
import gradio as gr

# ----------------------------
# 1) Load data (local CSV if present; else tiny toy dataset)
# ----------------------------
def load_dataset() -> Tuple[List[str], List[str]]:
    csv_path = "sms_spam.csv"
    if os.path.isfile(csv_path):
        texts, labels = [], []
        with open(csv_path, encoding="utf-8") as f:
            reader = csv.DictReader(f)
            required = {"text", "label"}
            if not required.issubset(set(reader.fieldnames or [])):
                raise ValueError(
                    "sms_spam.csv must contain columns: text,label"
                )
            for row in reader:
                texts.append((row["text"] or "").strip())
                labels.append((row["label"] or "").strip().lower())
        # basic cleanup
        labels = ["spam" if l == "spam" else "ham" for l in labels]
        return texts, labels

    # Fallback: tiny embedded toy dataset (for offline demo)
    toy = [
        ("Free entry in 2 a wkly comp to win FA Cup final tkts", "spam"),
        ("URGENT! You have won a 1 week FREE membership", "spam"),
        ("Call now to claim your prize", "spam"),
        ("WINNER!! As a valued network customer you have been selected", "spam"),
        ("Congrats! You won a lottery. Reply to claim.", "spam"),
        ("See you at practice tonight.", "ham"),
        ("Are we still on for lunch today?", "ham"),
        ("I'll pick you up at 7.", "ham"),
        ("Don't forget the meeting tomorrow morning.", "ham"),
        ("Can you send the report by EOD?", "ham"),
        ("This is not a drill. Limited time offer!!!", "spam"),
        ("Happy birthday! Hope you have a great day!", "ham"),
        ("Reminder: dentist appointment at 3pm", "ham"),
        ("Get cheap meds without prescription now", "spam"),
        ("Double your income fast. Click here", "spam"),
    ]
    texts, labels = zip(*toy)
    return list(texts), list(labels)

X_texts, y_labels = load_dataset()

# ----------------------------
# 2) Train/test split and Pipeline
# ----------------------------
X_train, X_test, y_train, y_test = train_test_split(
    X_texts, y_labels, test_size=0.25, random_state=42, stratify=y_labels
)

model = Pipeline(
    steps=[
        ("vect", CountVectorizer(stop_words="english")),
        ("clf", MultinomialNB()),
    ]
)

model.fit(X_train, y_train)

# Evaluate simple accuracy for display
if len(X_test) > 0:
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
else:
    accuracy = np.nan

# ----------------------------
# 3) Inference function for Gradio
# ----------------------------
def classify_message(message: str) -> Tuple[str, float]:
    """
    Returns:
        predicted_label ('spam' or 'ham'),
        spam_probability (0-100 as percentage)
    """
    if not message or not message.strip():
        return "ham", 0.0

    # predict_proba -> columns are in model.classes_
    proba = model.predict_proba([message])[0]
    # Find index for 'spam' class
    classes = list(model.classes_)
    if "spam" in classes:
        spam_idx = classes.index("spam")
        spam_prob = float(proba[spam_idx])
    else:
        # Shouldn't happen with our labels, but be safe
        spam_prob = 1.0 - float(proba[0])

    label = "spam" if spam_prob >= 0.5 else "ham"
    return label, round(spam_prob * 100.0, 2)

# ----------------------------
# 4) Build Gradio UI
# ----------------------------
title = "📩 Naive Bayes Spam Classifier"
if not np.isnan(accuracy):
    description = (
        "A simple text classifier using CountVectorizer + MultinomialNB. "
        f"Current test accuracy on this dataset: **{accuracy:.2f}**.\n\n"
        "Type a message to see if it's predicted as **spam** or **ham**, "
        "and view the spam probability."
    )
else:
    description = (
        "A simple text classifier using CountVectorizer + MultinomialNB.\n\n"
        "Type a message to see if it's predicted as **spam** or **ham**, "
        "and view the spam probability."
    )

examples = [
    ["Congratulations! You just won a $1,000 gift card. Click to claim."],
    ["Don't forget our meeting at 10am today."],
    ["URGENT! Your account has been compromised. Verify now."],
    ["See you at practice tonight."],
]

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"# {title}")
    gr.Markdown(description)

    with gr.Row():
        with gr.Column(scale=2):
            message = gr.Textbox(
                label="message",
                placeholder="Type or paste an SMS/email message here...",
                lines=4,
            )
            btn = gr.Button("Submit", variant="primary")
            clr = gr.Button("Clear")

        with gr.Column(scale=1):
            pred = gr.Textbox(label="Prediction (spam/ham)", interactive=False)
            prob = gr.Textbox(label="Spam Probability (%)", interactive=False)

    gr.Examples(
        examples=examples,
        inputs=[message],
        label="Try some examples",
    )

    def _clear():
        return gr.update(value=""), gr.update(value=""), gr.update(value="")

    btn.click(classify_message, inputs=message, outputs=[pred, prob])
    clr.click(_clear, outputs=[message, pred, prob])

# ----------------------------
# 5) Launch
# ----------------------------
if __name__ == "__main__":
    # Set share=True if you want a public share link
    demo.launch()