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