import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load model and tokenizer
model_name = "niru-nny/SMS_Spam_Detection"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

def classify_message(text):
    """
    Classify an SMS message as HAM or SPAM
    
    Args:
        text: Input SMS message text
        
    Returns:
        Dictionary with classification results
    """
    if not text or text.strip() == "":
        return {"Error": "Please enter a message"}
    
    # Tokenize input
    inputs = tokenizer(
        text, 
        return_tensors="pt", 
        truncation=True, 
        padding=True, 
        max_length=128
    )
    
    # Get prediction
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
    
    # Extract probabilities
    ham_prob = predictions[0][0].item()
    spam_prob = predictions[0][1].item()
    
    return {
        "HAM (Legitimate)": ham_prob,
        "SPAM": spam_prob
    }

# Example messages
examples = [
    ["Congratulations! You've won a $1000 gift card. Click here to claim now!"],
    ["Hey, are we still meeting for lunch tomorrow at 12?"],
    ["URGENT! Your account has been suspended. Verify now to restore access."],
    ["Thanks for your help today. I really appreciate it!"],
    ["FREE entry in 2 a wkly comp to win FA Cup final tkts 21st May 2005. Text FA to 87121"],
    ["I'll call you later tonight after work."],
    ["WINNER!! As a valued customer, you have been selected to receive £900 prize reward!"],
    ["Can you pick up some milk on your way home?"],
]

# Create Gradio interface
demo = gr.Interface(
    fn=classify_message,
    inputs=gr.Textbox(
        lines=5,
        placeholder="Enter SMS message here...",
        label="SMS Message"
    ),
    outputs=gr.Label(num_top_classes=2, label="Classification Results"),
    title="📱 SMS Spam Detection",
    description="""
    This classifier uses a fine-tuned BERT model to detect spam in SMS messages.
    
    **Performance:** 99.16% Accuracy | 97.30% Precision | 96.43% Recall
    
    Simply enter an SMS message below and the model will classify it as either legitimate (HAM) or spam (SPAM).
    """,
    examples=examples,
    theme=gr.themes.Soft(),
    article="""
    ### About This Model
    
    This model is trained on the SMS Spam Collection dataset and achieves state-of-the-art performance in spam detection.
    
    **Model:** `niru-nny/SMS_Spam_Detection`  
    **Base Architecture:** BERT (bert-base-uncased)  
    **Dataset:** SMS Spam Collection (5,574 messages)
    
    ### Use Cases
    - 📧 Spam filtering in messaging apps
    - 🛡️ Protection against phishing attempts
    - 🔍 Content moderation
    - 💰 Fraud detection
    
    ### Tips for Best Results
    - The model works best with English text messages
    - Keep messages under 128 words for optimal performance
    - The model is trained on SMS-style language (abbreviations, slang included)
    
    ---
    
    **License:** MIT | [Model Card](https://huggingface.co/niru-nny/SMS_Spam_Detection) | [GitHub](https://github.com/niru-nny)
    """
)

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