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

    if pred_sentiment == 2:
        sentiment_label = "positive"
    elif pred_sentiment == 1:
        sentiment_label = "neutral"
    else:
        sentiment_label = "negative"

    lang_code_map = {
        'de': 'German',
        'es': 'Espanyol',
        'en': 'English',
        'fr': 'French'
    }
        
    lang_code = le.inverse_transform([pred_lang])[0]
    lang_label = lang_code_map.get(lang_code, "Unknown")

    return sentiment_label, lang_label


interface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Masukkan Teks Dalam Bahasa (Inggris/Jerman/Spanyol/Perancis)"),
    outputs=[
        gr.Textbox(label="Prediksi Sentiment (Positif/Neutral/Negatif)"),
        gr.Textbox(label="Prediksi Bahasa")
    ],
    title="Multitask DistilBERT: Sentiment + Language",
    description="Prediksi sentimen dan bahasa dari teks menggunakan model multitask DistilBERT."
)

interface.launch()