import gradio as gr
from transformers import pipeline

# Load the sentiment analysis pipeline from Hugging Face Hub
classifier = pipeline("sentiment-analysis", model="Arvind111/sentiment_newclassifier")

# Define the label mapping from model output to human-readable emotions with emojis
# Based on previous tests: LABEL_0 -> Neutral, LABEL_1 -> Sad, LABEL_2 -> Happy
label_mapping = {
    'LABEL_0': 'Neutral \ud83d\ude10',
    'LABEL_1': 'Sad \ud83d\ude1e',
    'LABEL_2': 'Happy \ud83d\ude0a'
}

def predict_emotion(text):
    if not text:
        return "Please enter some text."
    
    # Get prediction from the pipeline
    predictions = classifier(text)
    
    if predictions:
        predicted_label_id = predictions[0]['label']
        predicted_score = predictions[0]['score']
        
        # Map to human-readable label with emoji
        emotion_label = label_mapping.get(predicted_label_id, "Unknown \ud83e\udd14")
        
        return f"Prediction: {emotion_label} (Score: {predicted_score:.4f})"
    else:
        return "Could not get prediction."


# Create the Gradio interface
iface = gr.Interface(
    fn=predict_emotion,
    inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
    outputs="text",
    title="Emotion Prediction with BERT",
    description="Enter a sentence and the model will predict its primary emotion (Happy, Sad, or Neutral)."
)

# Launch the Gradio interface
if __name__ == '__main__':
    iface.launch(share=False)