import tensorflow as tf
import numpy as np
import gradio as gr
from transformers import AutoTokenizer
from keras.layers import TFSMLayer
from huggingface_hub import snapshot_download

# 1️⃣ Download folder model secara lokal
model_path = snapshot_download("jeanetrixsiee/bert-sentimen-model", repo_type="model")

# 2️⃣ Muat model via TFSMLayer
model = TFSMLayer(model_path, call_endpoint="serving_default")

# 3️⃣ Muat tokenizer dari HF
tokenizer = AutoTokenizer.from_pretrained("jeanetrixsiee/bert-sentimen-model")

# 4️⃣ Label map sesuai urutan output model
label_map = {0: "Negatif", 1: "Netral", 2: "Positif", 3: "Campuran", 4: "Tidak Jelas"}

# 5️⃣ Fungsi prediksi
def predict_sentiment(text):
    tokens = tokenizer(text, return_tensors="tf", padding="max_length", truncation=True, max_length=256)
    input_ids, attention_mask = tokens["input_ids"], tokens["attention_mask"]

    outputs = model({"input_ids": input_ids, "attention_mask": attention_mask})
    logits = outputs["classifier"][0].numpy()
    pred = np.argmax(logits)
    score = logits[pred]

    return f"Prediksi: {label_map[pred]} ({score:.2%})"

# 6️⃣ UI Gradio
demo = gr.Interface(fn=predict_sentiment,
                    inputs=gr.Textbox(lines=3, placeholder="Masukkan komentar..."),
                    outputs=gr.Textbox(),
                    title="✅ Sentimen BERT Demo",
                    description="Model BERT TensorFlow via TFSMLayer + Gradio")

if __name__ == "__main__":
    demo.launch()