import gradio as gr
from transformers import pipeline

# Ensure the model is correctly loaded
model_name = "Abduuu/ArabReview-Sentiment"
sentiment_pipeline = pipeline("text-classification", model=model_name, tokenizer=model_name)

# Define label mapping (ensure it matches actual model outputs)
label_mapping = {
    "LABEL_0": "سلبي",
    "LABEL_1": "إيجابي",
    "Negative": "سلبي",
    "Positive": "إيجابي"
}

# Function to predict sentiment with confidence visualization
def predict_sentiment(review):
    result = sentiment_pipeline(review)[0]
    sentiment_label = label_mapping.get(result["label"], result["label"])  # Handle unexpected labels
    confidence = round(result["score"] * 100, 2)
    
    return {
        "التصنيف": sentiment_label,
        "نسبة الثقة": f"{confidence}%",
    }

# Define Gradio interface with better UI
with gr.Blocks(theme=gr.themes.Default()) as iface:
    gr.Markdown("<h1 style='text-align: center;'>🍽️ تحليل مشاعر مراجعات المطاعم 🚀</h1>")
    gr.Markdown(
        "<p style='text-align: center;'> 🤖 أدخل مراجعة مطعم بالعربية، وسيقوم النموذج بتحليل المشاعر وتحديدها **إيجابية** أو **سلبية**.</p>"
    )

    with gr.Row():
        review_input = gr.Textbox(
            label="✍️ أدخل مراجعتك", 
            placeholder="اكتب مراجعتك هنا...", 
            lines=2
        )

    submit_button = gr.Button("🔍 تحليل المراجعة")

    with gr.Row():
        output_label = gr.Textbox(label="🔹 التصنيف", interactive=False)
        output_confidence = gr.Textbox(label="📊 نسبة الثقة", interactive=False)

    submit_button.click(predict_sentiment, inputs=review_input, outputs=[output_label, output_confidence])

    gr.Examples(
        examples=[
            ["الطعام لذيذ جدًا والخدمة ممتازة"],
            ["التجربة كانت سيئة والطعام غير نظيف"],
            ["الخدمة كانت بطيئة والأسعار مرتفعة جدًا"],  
            ["تجربة رائعة، سأعود مجددًا"],  
            ["أفضل مطعم جربته على الإطلاق"],  
        ],
        inputs=review_input,
    )

# Launch the app with public sharing enabled
if __name__ == "__main__":
    iface.launch(share=True)