restore heatmap

app.py

@@ -123,7 +123,7 @@ def plot_sunburst_chart(heatmap_data):
         values='score',
         color='emotion',
         hover_data={'score': ':.3f'},
-        title='Sentence-Level Emotion Confidence Sunburst'
+        title='Sentence-Level Emotion Confidence'
     )
 
     fig.update_layout(
@@ -167,9 +167,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)
+    heatmap = plot_emotion_heatmap(prepare_heatmap_data(results_heatmap))
     output_info = f'Prediction was made using the <a href="https://huggingface.co/{model_id}">{model_id}</a> model.'
-    return results, figure, sunburst_chart, output_info
+    return results, figure, sunburst_chart, heatmap, output_info
 
 
 with gr.Blocks() as demo:
@@ -195,13 +195,16 @@ with gr.Blocks() as demo:
     with gr.Row():
         sunburst_chart = gr.Plot()
 
+    with gr.Row():
+        heatmap = gr.Plot()
+
     with gr.Row():
         model_info = gr.Markdown()
 
     predict_button.click(
         fn=predict_wrapper,
         inputs=[input_text, language_choice],
-        outputs=[result_table, plot, sunburst_chart, model_info]
+        outputs=[result_table, plot, sunburst_chart, heatmap, model_info]
     )
 
 if __name__ == "__main__":