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()