#!pip install gradio -q

import gradio as gr
import string
from transformers import pipeline

# 1. Load the model
model_id = "Neo111x/bert_sentiment"
classifier = pipeline("sentiment-analysis", model=model_id)

# 2. Helper to remove punctuation
def remove_punctuation(text):
    return text.translate(str.maketrans('', '', string.punctuation))

# 3. Enhanced Inference function
def predict_sentiment(text, clean_text):
    if not text.strip():
        return "Please enter some text to analyze."
    
    # Toggle punctuation removal
    if clean_text:
        text = remove_punctuation(text)
    
    results = classifier(text)
    label_key = results[0]['label']
    score = results[0]['score']
    
    mapping = {
        "LABEL_0": ("Neutral 😐", "The sentiment is balanced and objective."),
        "LABEL_1": ("Positive 😊", "The sentiment is upbeat and favorable!"),
        "LABEL_2": ("Negative 😡", "The sentiment appears critical or dissatisfied.")
    }
    
    label_display, description = mapping.get(label_key, (label_key, ""))
    return f"### Result: {label_display}\n**Confidence Score:** {score:.2%}\n**Processed Text:** {text}\n\n{description}"

# 4. Create UI with Toggle
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🤖 BERT Sentiment Analyzer
        ### Powered by Fine-tuned BERT-base-uncased
        *Analyze the emotional tone of tweets or messages instantly.*
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Input Message", 
                placeholder="Type a tweet or sentence here...", 
                lines=4
            )
            clean_toggle = gr.Checkbox(label="Remove Punctuation before analysis", value=False)
            btn = gr.Button("Analyze Sentiment ✨", variant="primary")
        
        with gr.Column():
            output_md = gr.Markdown(label="Prediction")

    gr.Examples(
        examples=[
            ["I am so happy with the new service, it's amazing!", False],
            ["The transaction took way too long and the app crashed.", True],
            ["CIBC is a financial institution based in Canada.", False]
        ],
        inputs=[input_text, clean_toggle]
    )

    btn.click(fn=predict_sentiment, inputs=[input_text, clean_toggle], outputs=output_md)

demo.launch(share=True)