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

# Load tokenizer and model from Hugging Face Hub
model_name = "TrioF/InSERT2"
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"]  # This will already be mapped using id2label from config.json
    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 6 kategori: hadiah,lowongan/investasi, no spam, program pemerintah/bantuan, promo/penjualan, dan urgensi."
)

demo.launch()