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")

# Define the function for emotion analysis and visualization
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}
    
    # Generate a bar chart for the emotions
    plt.figure(figsize=(8, 4))
    plt.bar(emotions.keys(), emotions.values(), color='skyblue')
    plt.xlabel("Emotions")
    plt.ylabel("Probability")
    plt.title("Emotion Analysis")
    plt.tight_layout()
    
    # 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()