import gradio as gr
from transformers import pipeline

# Load the model
classifier = pipeline("text-classification", model="IreNkweke/HamOrSpamModel")

# Define the function that will be used in the interface
def classify_text(text):
    result = classifier(text)
    label = result[0]['label']
    score = result[0]['score']
    
    # Assuming the model outputs scores for both classes
    if label == 'LABEL_1':
        spam_percentage = round(score * 100, 2)
        not_spam_percentage = round((1 - score) * 100, 2)
    else:
        spam_percentage = round((1 - score) * 100, 2)
        not_spam_percentage = round(score * 100, 2)
    
    label_mapping = {
        'LABEL_0': 'Ham',  # Non-spam
        'LABEL_1': 'Spam'  # Spam
    }
    
    return label_mapping[label], spam_percentage, not_spam_percentage

# Create the Gradio interface
iface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(
        label="Input Text",
        placeholder="Enter your message here...",
        lines=4
    ),
    outputs=[
        gr.Textbox(label="Class", placeholder="Classification result", lines=1),
        gr.Number(label="Spam (%)"),
        gr.Number(label="Not Spam (%)")
    ],
    title="Spam Classifier",
    description="Classify messages as Spam or Ham with percentage confidence. Enter a message below to see the classification and confidence percentages.",
    examples=[
        ["Congratulations! You've won a free gift card!"],
        ["Hi! I wanted to check in and see how you’ve been doing."],
        ["Urgent: Your account has been compromised. Please provide your login details to verify your identity."],
        ["The meeting is scheduled for 10 AM tomorrow. Please let me know if you need any changes."],
        ["Limited Time Offer: Buy one, get one free on all items! Shop now and save big!"]
    ],
    theme="default"
)

# Launch the interface
iface.launch()