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

MODEL_ID = "JustParadis/indobert-sentiment-comment"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
model.eval()

LABELS = model.config.id2label  # uses labels from training config

def predict(text):
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=128
    )

    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=-1)[0]

    return {
        LABELS[i]: float(probs[i])
        for i in range(len(probs))
    }

gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=4, placeholder="Masukkan komentar"),
    outputs=gr.JSON(label="Prediction"),
    title="IndoBERT Comment Sentiment",
    description="2-label comment sentiment classification"
).launch()