from transformers import pipeline
import gradio as gr
import matplotlib.pyplot as plt
import io
import base64

# Load the emotion classification model
emotion_pipeline = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")

def analyze_emotions(text):
    # Get the emotion predictions
    results = emotion_pipeline(text, top_k=None)  # Get probabilities for all emotions
    emotions = {res['label']: res['score'] for res in results}

    # Extract emotion labels and probabilities
    emotion_labels = list(emotions.keys())
    emotion_probabilities = list(emotions.values())

    # Generate a pie chart for the emotions
    plt.figure(figsize=(8, 6))
    plt.pie(emotion_probabilities, labels=emotion_labels, autopct='%1.1f%%', startangle=140)
    plt.title("Emotion Distribution")
    plt.axis('equal')

    # Save the plot to a bytes buffer
    buf = io.BytesIO()
    plt.savefig(buf, format="png")
    buf.seek(0)
    img_str = base64.b64encode(buf.read()).decode("utf-8")
    buf.close()

    # Return image data as HTML for Gradio
    img_html = f'<img src="data:image/png;base64,{img_str}" alt="Emotion Analysis"/>'
    return img_html

# Create Gradio interface
interface = gr.Interface(
    fn=analyze_emotions,
    inputs="text",
    outputs="html",
    title="Emotion Detection with Visualization",
    description="Enter text to detect and visualize emotions.")

# Launch the app
if __name__ == "__main__":
    interface.launch()