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

# Load tokenizer and model from Hugging Face Hub
model_name = "hasnanhaq/indoSPAM"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Create a pipeline
classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)

# Inference function
def classify_text(text):
    result = classifier(text)[0]
    label = result["label"]
    score = round(result["score"], 3)
    return f"{label} ({score})"

# Build Gradio UI
demo = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(lines=4, label="Masukkan pesan SMS/WA"),
    outputs=gr.Textbox(label="Prediksi"),
    title="Klasifikasi Pesan Spam Bahasa Indonesia",
    description=(
        "Model ini mengklasifikasikan pesan menjadi 5 kategori:\n"
        "1. no spam\n"
        "2. penipuan hadiah\n"
        "3. penipuan tawaran kerja\n"
        "4. minta data diri\n"
        "5. minta klik tautan"
    )
)

demo.launch()