import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Load the model and tokenizer from Hugging Face
model_name = "vsk21/spam_or_ham_bert-base-uncased"  # Replace with your model's name
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Define the prediction function
def classify_email(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    outputs = model(**inputs)
    predictions = outputs.logits.argmax(axis=-1).item()
    return "Spam" if predictions == 1 else "Ham"

# Create Gradio interface
interface = gr.Interface(fn=classify_email,
                         inputs="text",
                         outputs="text",
                         title="Email Classification",
                         description="Enter an email message to classify it as Spam or Ham.")

# Launch the interface
interface.launch()