Update app.py

app.py

@@ -10,10 +10,24 @@ df = pd.read_parquet(dataset_path)
 classifier = pipeline("text-classification", model="mrm8488/bert-tiny-finetuned-sms-spam-detection")
 
 def spam_detector(text):
-    """Detect if a message is spam or not."""
+    """Detect if a message is spam or not, with confidence filtering."""
     result = classifier(text)
-    label = result[0]['label'].lower()
-    return "Spam" if label == "spam" else "Not Spam"
+    
+    # Debugging: Print the raw output from the classifier
+    print("Model Output:", result)
+
+    # Extract label and confidence score
+    label = result[0]['label'].lower().strip()
+    confidence = result[0]['score']  # Confidence score (0 to 1)
+
+    # Confidence threshold (adjustable, 0.5 is standard)
+    threshold = 0.7
+
+    # Return based on confidence
+    if label == "spam" and confidence >= threshold:
+        return f"Spam (Confidence: {confidence:.2f})"
+    else:
+        return f"Not Spam (Confidence: {confidence:.2f})"
 
 # Create Gradio UI with enhanced styling
 app = gr.Interface(
@@ -22,11 +36,10 @@ app = gr.Interface(
     outputs=gr.Textbox(label="Prediction"),
     title="AI-Powered Spam Detector",
     description="Enter a message to check if it's spam or not, using a fine-tuned BERT model.",
-    theme="huggingface"
 )
 
 # Run the app
 if __name__ == "__main__":
     print("Loaded dataset preview:")
     print(df.head())
-    app.launch
+    app.launch