Update plot.py

plot.py

@@ -1,33 +1,31 @@
 # plot.py
 import matplotlib.pyplot as plt
 import numpy as np
-import io
+from io import BytesIO
 from PIL import Image
 
 def plot_radar_chart(config):
-    # Traits to plot
-    traits = ["layers", "attention_heads", "ffn_dim", "dropout"]
-    values = [
-        config["layers"],
-        config["attention_heads"],
-        config["ffn_dim"] / 512,  # Normalize
-        config["dropout"] * 10,   # Normalize
-    ]
-    # Repeat first value to close the loop
+    labels = list(config.keys())
+    values = list(config.values())
+
+    # Convert boolean to numeric for plotting
+    values = [1 if v is True else 0 if v is False else v for v in values]
+
+    angles = np.linspace(0, 2 * np.pi, len(labels), endpoint=False).tolist()
     values += values[:1]
-    angles = np.linspace(0, 2 * np.pi, len(traits), endpoint=False).tolist()
     angles += angles[:1]
 
-    fig, ax = plt.subplots(figsize=(5, 5), subplot_kw=dict(polar=True))
-    ax.plot(angles, values, "o-", linewidth=2)
+    fig, ax = plt.subplots(figsize=(6, 6), subplot_kw={'polar': True})
+    ax.plot(angles, values, 'o-', linewidth=2)
     ax.fill(angles, values, alpha=0.25)
-    ax.set_yticklabels([])
-    ax.set_xticks(angles[:-1])
-    ax.set_xticklabels(traits)
-    ax.set_title("Final Generation Trait Radar", fontsize=14)
 
-    buf = io.BytesIO()
+    ax.set_thetagrids(np.degrees(angles[:-1]), labels)
+    ax.set_title("Trait Radar Chart", fontsize=14)
+    ax.grid(True)
+
+    buf = BytesIO()
+    plt.tight_layout()
     plt.savefig(buf, format="png")
+    plt.close(fig)
     buf.seek(0)
-    image = Image.open(buf)
-    return image
+    return Image.open(buf)