app.py

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

# Load model and tokenizer with authentication
model_name = "Anurag3703/bert-spam-classifier"
token = os.environ.get("HF_TOKEN")  # Get token from environment

tokenizer = AutoTokenizer.from_pretrained(model_name, token=token)
model = AutoModelForSequenceClassification.from_pretrained(model_name, token=token)


def classify_text(text):
    """Classify text as spam or not spam"""
    # Tokenize input
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)

    # Get predictions
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)

    # Get probabilities
    ham_prob = predictions[0][0].item()
    spam_prob = predictions[0][1].item()

    return {
        "Not Spam (Ham)": ham_prob,
        "Spam": spam_prob
    }


# Create Gradio interface
demo = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(
        lines=3,
        placeholder="Enter a message to classify...",
        label="Input Text"
    ),
    outputs=gr.Label(num_top_classes=2, label="Classification Results"),
    title="BERT Spam Classifier",
    description="Enter a text message to check if it's spam or not spam.",
    examples=[
        ["Win a free iPhone now! Click here!"],
        ["Hey, can we meet for coffee tomorrow?"],
        ["URGENT: Your account will be closed. Verify now!"],
        ["Thanks for the meeting today, let's catch up next week."]
    ],
    theme="soft"
)

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