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

# Load model and tokenizer
model_dir = "fine_tuned_bert"
tokenizer = BertTokenizer.from_pretrained(model_dir)
model = BertForSequenceClassification.from_pretrained(model_dir).to("cuda")

# Set model to evaluation mode
model.eval()

# Define inference function
@spaces.GPU
def classify_text(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to("cuda")
    
    with torch.no_grad():
        outputs = model(**inputs)
    
    predicted_class = torch.argmax(outputs.logits, dim=1).item()

    label_map = {0: "Not Spam", 1: "Spam"}
    predicted_label = label_map.get(predicted_class, "Unknown")

    return f"Predicted Class : {predicted_label} "

# Gradio Interface
gr.Interface(
    fn=classify_text,
    inputs="text",
    outputs="text",
    title="BERT Text Classifier"
).launch()