import gradio as gr
from transformers import BertTokenizer, BertForSequenceClassification
import torch

# Load the tokenizer and model
model_name = "AventIQ-AI/bert-spam-detection"
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForSequenceClassification.from_pretrained(model_name)

# Set the model to evaluation mode and move it to the appropriate device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

# Define the prediction function
def predict_spam(text):
    """Predicts whether a given text is spam or not."""
    # Tokenize input text
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
    inputs = {key: value.to(device) for key, value in inputs.items()}

    # Perform inference
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probabilities = torch.softmax(logits, dim=1)
        prediction = torch.argmax(probabilities, dim=1).item()
        confidence = probabilities[0][prediction].item()

    # Map prediction to label
    label_map = {0: "Not Spam", 1: "Spam"}
    result = f"Prediction: {label_map[prediction]}\nConfidence: {confidence:.2f}"
    return result

# Create the Gradio interface
iface = gr.Interface(
    fn=predict_spam,
    inputs=gr.Textbox(label="📧 Input Text", placeholder="Enter the email or message content here...", lines=5),
    outputs=gr.Textbox(label="🔍 Spam Detection Result"),
    title="🛡️ BERT-Based Spam Detector",
    description="Enter the content of an email or message to determine whether it's Spam or Not Spam.",
    examples=[
        ["Congratulations! You've won a $1,000,000 lottery. Click here to claim your prize."],
        ["Hey, are we still meeting for lunch tomorrow?"],
        ["URGENT: Your account has been compromised. Reset your password immediately by clicking this link."],
        ["Don't miss out on our exclusive offer! Buy one, get one free on all items."],
        ["Can you send me the report by end of the day? Thanks!"]
    ],
    theme="compact",
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch()