Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,18 +1,19 @@
 import streamlit as st
 import pandas as pd
+import string
+import matplotlib.pyplot as plt
+import seaborn as sns
 from sklearn.model_selection import train_test_split
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.naive_bayes import MultinomialNB
-from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix
 
-# Set page config
+# ----------------- STREAMLIT CONFIG -----------------
 st.set_page_config(page_title="📧 Email Spam Detector", layout="centered")
-
-# Title and Description
 st.title("📧 Email Spam Detector")
-st.markdown("This app uses **Machine Learning** to classify emails as **Spam** or **Ham (Not Spam)**. Just type or paste your message below to test!")
+st.markdown("This app uses **Machine Learning** (Naive Bayes + TF-IDF) to classify emails as **Spam** or **Ham (Not Spam)**.")
 
-# Load and preprocess dataset
+# ----------------- DATA LOADING -----------------
 @st.cache_data
 def load_data():
     df = pd.read_csv("spam.csv", encoding='latin-1')[['v1', 'v2']]
@@ -22,44 +23,70 @@ def load_data():
 
 df = load_data()
 
-# Train/test split
+# ----------------- PREPROCESS FUNCTION -----------------
+def clean_text(text):
+    text = text.lower().strip()
+    text = text.translate(str.maketrans("", "", string.punctuation))
+    return text
+
+df['message'] = df['message'].apply(clean_text)
+
+# ----------------- TRAIN / TEST SPLIT -----------------
 X_train, X_test, y_train, y_test = train_test_split(
-    df['message'], df['label'], test_size=0.2, random_state=42)
+    df['message'], df['label'], test_size=0.2, random_state=42
+)
 
-# Vectorize with TF-IDF
+# ----------------- TF-IDF VECTORIZATION -----------------
 vectorizer = TfidfVectorizer(stop_words='english')
 X_train_tfidf = vectorizer.fit_transform(X_train)
 X_test_tfidf = vectorizer.transform(X_test)
 
-# Train model
+# ----------------- MODEL TRAINING -----------------
 model = MultinomialNB()
 model.fit(X_train_tfidf, y_train)
 
-# Accuracy (optional to show)
-accuracy = accuracy_score(y_test, model.predict(X_test_tfidf))
+# ----------------- METRICS -----------------
+y_pred = model.predict(X_test_tfidf)
+accuracy = accuracy_score(y_test, y_pred)
+precision = precision_score(y_test, y_pred)
+recall = recall_score(y_test, y_pred)
+f1 = f1_score(y_test, y_pred)
 
-# Sidebar
-st.sidebar.header("🔍 Model Info")
+# ----------------- SIDEBAR METRICS -----------------
+st.sidebar.header("🔍 Model Performance")
 st.sidebar.write(f"**Accuracy:** {accuracy:.2%}")
+st.sidebar.write(f"**Precision:** {precision:.2%}")
+st.sidebar.write(f"**Recall:** {recall:.2%}")
+st.sidebar.write(f"**F1 Score:** {f1:.2%}")
 st.sidebar.markdown("Model: `Multinomial Naive Bayes`  \nVectorizer: `TF-IDF`")
 
-# Predict function
+# Confusion Matrix
+cm = confusion_matrix(y_test, y_pred)
+fig, ax = plt.subplots()
+sns.heatmap(cm, annot=True, fmt='d', cmap="Blues", xticklabels=["Ham", "Spam"], yticklabels=["Ham", "Spam"])
+plt.ylabel('Actual')
+plt.xlabel('Predicted')
+st.sidebar.pyplot(fig)
+
+# ----------------- PREDICT FUNCTION -----------------
 def predict_message(msg):
-    vect_msg = vectorizer.transform([msg])
+    msg_clean = clean_text(msg)
+    vect_msg = vectorizer.transform([msg_clean])
     pred = model.predict(vect_msg)[0]
-    return "🚫 Spam" if pred == 1 else "✅ Ham (Not Spam)"
+    prob = model.predict_proba(vect_msg)[0][pred]
+    return ("🚫 Spam", prob) if pred == 1 else ("✅ Ham (Not Spam)", prob)
 
-# Input section
+# ----------------- USER INPUT -----------------
 st.subheader("✉️ Test Your Message")
 user_input = st.text_area("Enter your email message here:")
 
 if st.button("Detect"):
     if user_input.strip() == "":
-        st.warning("Please enter a message to classify.")
+        st.warning("⚠️ Please enter a message to classify.")
     else:
-        result = predict_message(user_input)
-        st.success(f"Prediction: **{result}**")
+        result, confidence = predict_message(user_input)
+        st.success(f"Prediction: **{result}**  \nConfidence: **{confidence:.2%}**")
 
-# Optional: Show raw data
+# ----------------- SAMPLE DATA -----------------
 with st.expander("📂 View Sample Dataset"):
     st.dataframe(df.sample(10))