Spaces:
Sleeping
Sleeping
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() |