File size: 2,356 Bytes
d104cb5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import streamlit as st
from transformers import pipeline
import matplotlib.pyplot as plt
import io
import json
# Load the emotion classification model
@st.cache_resource
def load_model():
return pipeline(
"text-classification",
model="j-hartmann/emotion-english-distilroberta-base",
return_all_scores=True
)
emotion_classifier = load_model()
# Define a function to predict emotions and generate a bar chart
def predict_emotion_with_chart(text):
if not text.strip():
return None, None
# Get predictions
results = emotion_classifier(text)
emotions = {result['label']: round(result['score'], 2) for result in results[0]}
sorted_emotions = dict(sorted(emotions.items(), key=lambda item: item[1], reverse=True))
# Create a bar chart
fig, ax = plt.subplots(figsize=(6, 4))
ax.bar(sorted_emotions.keys(), sorted_emotions.values(), color="skyblue")
ax.set_title("Emotion Scores")
ax.set_ylabel("Confidence Score")
ax.set_ylim(0, 1)
ax.set_xticklabels(sorted_emotions.keys(), rotation=45, ha="right")
plt.tight_layout()
return sorted_emotions, fig
# Function to generate JSON result
def generate_json(text):
results = emotion_classifier(text)
emotions = {result['label']: round(result['score'], 2) for result in results[0]}
return json.dumps(emotions, indent=2)
# Streamlit UI
st.title("🌟 Enhanced Emotion Detection App")
st.markdown("Analyze the emotions in a sentence and visualize them. Enter text to detect emotions, see a bar chart of scores, and download the results as a JSON file.")
# Input text
text_input = st.text_area("Enter text to analyze emotions:", "")
if st.button("Analyze Emotions"):
if text_input.strip():
emotion_scores, chart = predict_emotion_with_chart(text_input)
if emotion_scores:
st.subheader("Emotion Scores")
st.json(emotion_scores)
st.subheader("Emotion Chart")
st.pyplot(chart)
else:
st.error("No emotions detected. Please enter a valid text.")
else:
st.warning("Please enter some text.")
# Download JSON results
if text_input.strip():
json_data = generate_json(text_input)
st.download_button(label="Download Results as JSON", data=json_data, file_name="emotion_results.json", mime="application/json") |