JustParadis's picture
Update app.py
fc42c6d verified
raw
history blame contribute delete
981 Bytes
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()