[WIP] Update

analysis/plot_results.py

@@ -7,6 +7,7 @@
 import pandas as pd
 import seaborn as sns
 import numpy as np
+import matplotlib.patches as patches
 import matplotlib.pyplot as plt
 
 FONT_SIZES = {"small": 12, "medium": 16, "large": 18}
@@ -78,22 +79,20 @@ def plot_main_heatmap(
     # Remove unnecessary column
     df.pop("eng_Latn")
 
-    if top_ten_only:
-        df = df.sort_values(by="Avg_Multilingual", ascending=False).head(10).reset_index(drop=True)
-        data = df[[col for col in df.columns if col not in ("Model_Type", "Avg_Multilingual")]]
-        data = data.set_index("Model")
-        data = data * 100
+    df = df.sort_values(by="Avg_Multilingual", ascending=False).head(10).reset_index(drop=True)
+    data = df[[col for col in df.columns if col not in ("Model_Type", "Avg_Multilingual")]]
+    data = data.set_index("Model")
+    data = data * 100
 
-        fig, ax = plt.subplots(1, 1, figsize=figsize)
-        sns.heatmap(data, ax=ax, cmap="YlGn", annot=True, annot_kws={"size": 14}, fmt=".2f", cbar=False)
-        ax.xaxis.set_ticks_position("top")
-        ax.tick_params(axis="x", rotation=45)
-        ax.set_ylabel("")
+    fig, ax = plt.subplots(1, 1, figsize=figsize)
+    sns.heatmap(data, ax=ax, cmap="YlGn", annot=True, annot_kws={"size": 14}, fmt=".2f", cbar=False)
+    ax.xaxis.set_ticks_position("top")
+    ax.tick_params(axis="x", rotation=45)
+    ax.set_ylabel("")
+    ax.set_yticklabels([f"{model}     " for model in data.index])
 
-        plt.tight_layout()
-        fig.savefig(output_path, bbox_inches="tight")
-    else:
-        pass
+    plt.tight_layout()
+    fig.savefig(output_path, bbox_inches="tight")
 
 
 if __name__ == "__main__":