transposed heatmap

app.py

@@ -75,31 +75,23 @@ def get_most_probable_label(probs):
 
 
 def prepare_heatmap_data(data):
-    print("DEBUG: Data type received in prepare_heatmap_data:", type(data))
-    print("DEBUG: Data content:", data)
-
-    heatmap_data = pd.DataFrame(0.0, index=range(len(data)), columns=id2label.values())
+    heatmap_data = pd.DataFrame(0.0, index=id2label.values(), columns=range(len(data)))
 
     for idx, row in enumerate(data):
-        print(f"DEBUG: Item at index {idx}: {row} (type: {type(row)})")
-
-        if not isinstance(row, dict):
-            raise TypeError(f"Expected dict, but got {type(row)} at index {idx}: {row}")
-
         confidences = row["emotions"].tolist()
         for idy, confidence in enumerate(confidences):
             emotion = id2label[idy]
-            heatmap_data.at[idx, emotion] = round(confidence, 4)
+            heatmap_data.at[emotion, idx] = round(confidence, 4)
 
-    heatmap_data.index = [item["sentence"][:18]+"..." for item in data]
+    heatmap_data.columns = [item["sentence"][:18]+"..." for item in data]
     return heatmap_data
 
 def plot_emotion_heatmap(heatmap_data):
-    fig = plt.figure(figsize=(10, len(heatmap_data) * 0.5 + 2))
-    sns.heatmap(heatmap_data, annot=True, cmap="coolwarm", cbar=True, linewidths=0.5, linecolor='gray')
+    fig = plt.figure(figsize=(len(heatmap_data.columns) * 0.5 + 2, 10))
+    sns.heatmap(heatmap_data, annot=False, cmap="coolwarm", cbar=True, linewidths=0.5, linecolor='gray')
     plt.title("Sentence-level emotion confidences")
-    plt.xlabel("Emotions")
-    plt.ylabel("Sentences")
+    plt.xlabel("Sentences")
+    plt.ylabel("Emotions")
     plt.subplots_adjust(left=0.2, right=0.95, top=0.9, bottom=0.2)
     plt.tight_layout()
     return fig