teguhsuby's picture
Upload
2fc050b
raw
history blame
1.37 kB
import gradio as gr
import torch
import joblib
from transformers import AutoTokenizer
from dinstilBert import MultiTaskBERT
model = MultiTaskBERT()
model.load_state_dict(torch.load("model.pt", map_location="cpu"))
model.eval()
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-multilingual-cased")
le = joblib.load("label_encoder.pkl")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
def predict(text):
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
with torch.no_grad():
sentiment_logits, lang_logits = model(inputs["input_ids"], inputs["attention_mask"])
pred_sentiment = sentiment_logits.argmax(dim=1).item()
pred_lang = lang_logits.argmax(dim=1).item()
sentiment_label = "positive" if pred_sentiment == 1 else "negative"
lang_label = le.inverse_transform([pred_lang])[0]
return sentiment_label, lang_label
interface = gr.Interface(
fn=predict,
inputs=gr.Textbox(label="Masukkan Teks Dalam Bahasa (Inggris/Belanda/Spanyol/Perancis)"),
outputs=[
gr.Textbox(label="Prediksi Sentiment"),
gr.Textbox(label="Prediksi Bahasa")
],
title="Multitask DistilBERT: Sentiment + Language",
description="Prediksi sentimen dan bahasa dari teks menggunakan model multitask DistilBERT."
)
interface.launch()