import gradio as gr
from transformers import pipeline
import pandas as pd

# Load dataset from Hugging Face Hub
dataset_path = "hf://datasets/ucirvine/sms_spam/plain_text/train-00000-of-00001.parquet"
df = pd.read_parquet(dataset_path)

# Load a spam classification model
classifier = pipeline("text-classification", model="mrm8488/bert-tiny-finetuned-sms-spam-detection")

def spam_detector(text):
    """Detect if a message is spam or not."""
    result = classifier(text)
    label = result[0]['label'].lower()
    return "Spam" if label == "spam" else "Not Spam"

# Create Gradio UI with enhanced styling
app = gr.Interface(
    fn=spam_detector,
    inputs=gr.Textbox(label="Enter a message", placeholder="Type your message here..."),
    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