import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

model_name = "arminfrq/imdb_bert_classifier"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)

# گرفتن مپ لیبل‌ها
id2label = model.config.id2label

def predict_review(user_input):
    result = classifier(user_input)[0]
    label = result['label']
    score = result['score']
    
    # تبدیل LABEL_0 / LABEL_1 به متن قابل فهم
    if label in id2label.values():
        if "POS" in label.upper() or "1" in label:
            return f"Positive ({score*100:.2f}%)"
        else:
            return f"Negative ({score*100:.2f}%)"
    return f"{label} ({score*100:.2f}%)"  # fallback

with gr.Blocks() as demo:
    gr.Markdown("# IMDb Review Sentiment Classifier")
    with gr.Row():
        input_text = gr.Textbox(label="Enter a movie review", placeholder="Write your review here...")
        output_text = gr.Textbox(label="Prediction")
    send_button = gr.Button("Send")
    send_button.click(predict_review, input_text, output_text)

demo.launch(server_name="0.0.0.0", server_port=7860, share=True)