import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_1 = "Dimsralf/indobert_no_ros"
MODEL_2 = "Dimsralf/indobert_ros"

label_map = {0: "NEGATIF", 1: "POSITIF"}

@st.cache_resource
def load_model(model_name):
    # Coba load tokenizer fast terlebih dahulu
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
    except Exception as e_fast:
        st.warning(f"Gagal load tokenizer fast untuk {model_name}: {e_fast}")
        try:
            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
        except Exception as e_slow:
            st.error(f"Gagal load tokenizer slow untuk {model_name}: {e_slow}")
            # fallback ke tokenizer yang lebih umum
            tokenizer = AutoTokenizer.from_pretrained("bert-base-multilingual-uncased", use_fast=False)
            st.info(f"Gunakan fallback tokenizer bert-base-multilingual-uncased untuk {model_name}")

    # Load model klasifikasi (asumsi model kamu untuk classification)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    model.eval()
    return tokenizer, model

tokenizer1, model1 = load_model(MODEL_1)
tokenizer2, model2 = load_model(MODEL_2)

st.title("Analisis Sentimen dengan 2 Model IndoBERT")
st.write("Model 1:", MODEL_1)  
st.write("Model 2:", MODEL_2)

text = st.text_input("Masukkan kalimat:")

def predict(text, tokenizer, model):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.softmax(logits, dim=1)
    pred_id = torch.argmax(probs, dim=1).item()
    label = label_map.get(pred_id, "UNKNOWN")
    prob = probs[0][pred_id].item()
    return label, prob

if text:
    label1, prob1 = predict(text, tokenizer1, model1)
    label2, prob2 = predict(text, tokenizer2, model2)

    st.write("## Hasil Prediksi")
    col1, col2 = st.columns(2)

    with col1:
        st.subheader("Model 1 (no_ros)")
        st.write(f"Prediksi: {label1}")
        st.write(f"Probabilitas: {prob1:.4f}")

    with col2:
        st.subheader("Model 2 (ros)")
        st.write(f"Prediksi: {label2}")
        st.write(f"Probabilitas: {prob2:.4f}")