Update app.py

app.py

@@ -6,12 +6,15 @@ from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.linear_model import LogisticRegression
 from sklearn.model_selection import train_test_split
 from sklearn.metrics import accuracy_score
+from collections import Counter
 
 # 1. Load dataset
-dataset = load_dataset("ucirvine/sms_spam", split="train")
+dataset = load_dataset("sms_spam", split="train")
 texts = dataset["sms"]
 labels = [1 if label == "spam" else 0 for label in dataset["label"]]
 
+print("Label distribution:", Counter(labels))  # Debug check
+
 # 2. Clean text
 def clean_text(text):
     text = text.lower()
@@ -20,23 +23,28 @@ def clean_text(text):
 
 texts_cleaned = [clean_text(t) for t in texts]
 
-# 3. Train/test split (use stratified sampling!)
+# 3. Train/test split with stratification
 X_train, X_test, y_train, y_test = train_test_split(
     texts_cleaned, labels, test_size=0.2, random_state=42, stratify=labels
 )
 
-# 4. Build model: TF-IDF + Logistic Regression
+print("Train labels:", Counter(y_train))  # Debug check
+print("Test labels:", Counter(y_test))    # Debug check
+
+# 4. Build model pipeline
 model = make_pipeline(
     TfidfVectorizer(ngram_range=(1, 2), stop_words="english", max_df=0.9),
     LogisticRegression(max_iter=1000, class_weight="balanced")
 )
+
+# 5. Train model
 model.fit(X_train, y_train)
 
-# 5. Show validation accuracy
+# 6. Evaluate
 y_pred = model.predict(X_test)
 print("Validation Accuracy:", accuracy_score(y_test, y_pred))
 
-# 6. Prediction function
+# 7. Predict function
 def predict_spam(message):
     cleaned = clean_text(message)
     pred = model.predict([cleaned])[0]
@@ -44,13 +52,13 @@ def predict_spam(message):
     label = "🚫 Spam" if pred == 1 else "📩 Not Spam (Ham)"
     return f"{label} (Confidence: {prob:.2%})"
 
-# 7. Gradio UI
+# 8. Gradio app
 iface = gr.Interface(
     fn=predict_spam,
     inputs=gr.Textbox(lines=4, label="Enter your SMS message"),
     outputs=gr.Text(label="Prediction"),
-    title="📬 Improved SMS Spam Detector",
-    description="Detects spam in SMS messages using Logistic Regression with TF-IDF bi-grams. Now with higher accuracy!"
+    title="📬 SMS Spam Detector (Improved)",
+    description="Detect spam in SMS messages using Logistic Regression with TF-IDF bi-grams. Trained on a balanced dataset from Hugging Face."
 )
 
 if __name__ == "__main__":