Update app.py

app.py

@@ -1,18 +1,37 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-clf = pipeline(
-    "text-classification",
-    model="JustParadis/indobert-sentiment-comment",
-    return_all_scores=True
-)
+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):
-    return clf(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"
+    title="IndoBERT Comment Sentiment",
+    description="2-label comment sentiment classification"
 ).launch()