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

MODEL_NAME = "Dimsralf/indobert"

st.title("Demo Model NLP")

st.write("Memuat model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
model.eval()

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

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

if text:
    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[pred_id]

    st.write("### Hasil Prediksi")
    st.write(f"**Label Prediksi:** {label}")
    st.write(f"**Probabilitas:** {probs[0][pred_id].item():.4f}")