import gradio as gr
from transformers import BertTokenizer, TFBertForSequenceClassification
import tensorflow as tf
import numpy as np

# Load model and tokenizer from Hugging Face Hub
model_name = "shobrunjb/spiill-fake-review-product-v2"
tokenizer = BertTokenizer.from_pretrained(model_name)
model = TFBertForSequenceClassification.from_pretrained(model_name)

def predict(text):
    # Tokenize input text
    inputs = tokenizer(text, return_tensors="tf", truncation=True, padding=True, max_length=400)
    
    # Predict with the model
    outputs = model(inputs)
    
    # Get predictions for sentiment and review label
    sentiment_logits = outputs.logits[0].numpy()
    review_logits = outputs.logits[1].numpy()
    
    # Convert logits to probabilities using softmax
    sentiment_probs = tf.nn.softmax(sentiment_logits).numpy()
    review_probs = tf.nn.softmax(review_logits).numpy()
    
    # Convert logits to class labels
    sentiment_label_map = {0: 'negative', 1: 'neutral', 2: 'positive'}
    review_label_map = {0: 'fake', 1: 'trusted', 2: 'non'}
    
    sentiment_pred = {sentiment_label_map[i]: f"{sentiment_probs[i]*100:.2f}%" for i in range(3)}
    review_pred = {review_label_map[i]: f"{review_probs[i]*100:.2f}%" for i in range(3)}
    
    return sentiment_pred, review_pred

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Spiill- Deteksi Fake Review Produk")
    gr.Markdown("Model ini memprediksi dan menganalsis review produk dengan multi task learning ")
    
    with gr.Row():
        input_text = gr.Textbox(label="Input", placeholder="Masukkan teks di sini...", lines=4)
        sentiment_output = gr.JSON(label="Sentimen (dengan Presentase)")
        review_output = gr.JSON(label="Label Review (dengan Presentase)")
    
    submit_btn = gr.Button("Submit")
    submit_btn.click(fn=predict, inputs=input_text, outputs=[sentiment_output, review_output])

    gr.Markdown("### Contoh Kalimat:")
    gr.Examples(
        examples=["Bibirku lagi iritasi parah. Kering, gatal, dan mengelupas sampai luka. Diolesin mediheal ini langsung membaik dalam 2 hari. Bibirku jadi gak patchy lagi dan lukanya sembuh.",
                  "Lip balm dengan 100,000ppm Panthenol, shea butter, beragam plant oil yang efektif untuk melembapkan bibir sangat kering, pecah-pecah, dan menghaluskan bibir.",
                  "AKU CINTA BANGET BANGET beneran bantuuu rambut aku pas lagi rontok rontoknya 😭😭🤞🏻 tujuanku make cuma buat ilangin botak tapi surprisingly bikin rambut makin halus juga omg HIDDEN GEM kata aku"],
        inputs=input_text,
    )

# Launch the interface
demo.launch()