from transformers import pipeline
import gradio as gr

classifier = pipeline("text-classification", model="bhadresh-savani/bert-base-uncased-emotion")

label_map = {
    "LABEL_0": "HAM (Not Spam)",
    "LABEL_1": "SPAM"
}

def predict_spam(text):
    result = classifier(text)[0]
    label = label_map.get(result['label'], "Unknown")
    confidence = round(result['score'] * 100, 2)
    return f" Prediction: {label}\n Confidence: {confidence}%"

# Build Gradio UI
gr.Interface(
    fn=predict_spam,
    inputs=gr.Textbox(lines=4, placeholder="Enter a message..."),
    outputs="text",
    title=" Spam Detector",
    description="Detects SPAM vs HAM using a pretrained BERT-tiny model.",
    examples=[
        "Win a free iPhone now!",
        "Hey, are we still on for dinner tonight?",
        "Your OTP is 456123.",
        "Click here to claim your reward."
    ]
).launch()