id2label

app.py

@@ -74,10 +74,10 @@ def get_most_probable_label(probs):
     return label, probability
 
 def prepare_heatmap_data(data):
-    heatmap_data = pd.DataFrame(0, index=range(len(data)), columns=emotion_mapping.values())
+    heatmap_data = pd.DataFrame(0, index=range(len(data)), columns=id2label.values())
     for idx, item in enumerate(data):
         for idy, confidence in enumerate(item["emotions"]):
-            emotion = emotion_mapping[idy]
+            emotion = id2label[idy]
             heatmap_data.at[idx, emotion] = confidence
     heatmap_data.index = [item["sentence"] for item in data]
     return heatmap_data