sunburst chart

app.py

@@ -8,6 +8,7 @@ from transformers import AutoModelForSequenceClassification
 from transformers import AutoTokenizer
 import gradio as gr
 import matplotlib.pyplot as plt
+import plotly.express as px
 import seaborn as sns
 
 PATH = '/data/' # at least 150GB storage needs to be attached
@@ -99,6 +100,31 @@ def plot_emotion_heatmap(heatmap_data):
     plt.tight_layout()
     return fig
 
+def plot_sunburst_chart(heatmap_data):
+    data = []
+    for item in heatmap_data:
+        sentence = item['sentence']
+        emotions = item['emotions']
+        for i, score in enumerate(emotions):
+            data.append({
+                'root': 'All Sentences',
+                'sentence': sentence,
+                'emotion': id2label[i],
+                'score': float(score)
+            })
+
+    df = pd.DataFrame(data)
+
+    # Plot sunburst
+    fig = px.sunburst(
+        df,
+        path=['root', 'sentence', 'emotion'],
+        values='score',
+        color='emotion',
+        title='Sentence-level Emotion Confidences'
+    )
+
+    return fig
 
 def plot_emotion_barplot(heatmap_data):
     most_probable_emotions = heatmap_data.idxmax(axis=0)
@@ -133,8 +159,9 @@ def predict_wrapper(text, language):
     print(results_heatmap)
 
     figure = plot_emotion_barplot(prepare_heatmap_data(results_heatmap))
+    sunburst_chart = plot_sunburst_chart(results_heatmap)
     output_info = f'Prediction was made using the <a href="https://huggingface.co/{model_id}">{model_id}</a> model.'
-    return results, figure, output_info
+    return results, figure, sunburst_chart, output_info
 
 
 with gr.Blocks() as demo:
@@ -163,7 +190,7 @@ with gr.Blocks() as demo:
     predict_button.click(
         fn=predict_wrapper,
         inputs=[input_text, language_choice],
-        outputs=[result_table, plot, model_info]
+        outputs=[result_table, plot, sunburst_chart, model_info]
     )
 
 if __name__ == "__main__":