added figure logic

app.py

@@ -5,6 +5,8 @@ import numpy as np
 from transformers import AutoModelForSequenceClassification
 from transformers import AutoTokenizer
 import gradio as gr
+import matplotlib.plt as pyplot
+import seaborn as sns
 
 PATH = '/data/' # at least 150GB storage needs to be attached
 os.environ['TRANSFORMERS_CACHE'] = PATH
@@ -70,6 +72,25 @@ def get_most_probable_label(probs):
     probability = f"{round(100 * probs.max(), 2)}%"
     return label, probability
 
+def prepare_heatmap_data(data):
+    heatmap_data = pd.DataFrame(0, index=range(len(data)), columns=emotion_mapping.values())
+    for idx, item in enumerate(data):
+        for idy, confidence in enumerate(item["emotions"]):
+            emotion = emotion_mapping[idy]
+            heatmap_data.at[idx, emotion] = confidence
+    heatmap_data.index = [item["sentence"] for item in data]
+    return heatmap_data
+
+def plot_emotion_heatmap(data):
+    heatmap_data = prepare_heatmap_data(data)
+    fig = plt.figure(figsize=(10, len(data) * 0.5 + 2))
+    sns.heatmap(heatmap_data, annot=True, cmap="coolwarm", cbar=True, linewidths=0.5, linecolor='gray')
+    plt.title("Emotion Confidence Heatmap")
+    plt.xlabel("Emotions")
+    plt.ylabel("Sentences")
+    plt.tight_layout()
+    return fig
+
 def predict_wrapper(text, language):
     model_id = build_huggingface_path(language)
     tokenizer_id = "xlm-roberta-large"
@@ -78,13 +99,17 @@ def predict_wrapper(text, language):
     sentences = split_sentences(text, spacy_model)
 
     results = []
+    results_heatmap = []
     for sentence in sentences:
         probs = predict(sentence, model_id, tokenizer_id)
         label, probability = get_most_probable_label(probs)
         results.append([sentence, label, probability])
+        results_heatmap.append({"sentence":sentence, "emotions":probs})
 
+    figure = 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, output_info
+    return results, figure, output_info
+
 
 with gr.Blocks() as demo:
     with gr.Row():
@@ -108,7 +133,7 @@ with gr.Blocks() as demo:
     predict_button.click(
         fn=predict_wrapper,
         inputs=[input_text, language_choice],
-        outputs=[result_table, model_info]
+        outputs=[result_table, "plot", model_info]
     )
 
 if __name__ == "__main__":