Update app.py

app.py

@@ -992,7 +992,7 @@ with ui.navset_card_tab(id="tab"):
                     "Select Model Type:",
                     ["pythia", "denseformer"],
                     multiple=True,
-                    selected=None
+                    selected='pythia'
                 )
     
         with ui.card():
@@ -1001,7 +1001,7 @@ with ui.navset_card_tab(id="tab"):
                     "Select Loss Type:",
                     ["compliment", "cross_entropy", "headless"],
                     multiple=True,
-                    selected=None
+                    selected='compliment'
                 )
     
         @output
@@ -1018,8 +1018,49 @@ with ui.navset_card_tab(id="tab"):
                 & (df["loss_type"].isin(loss_type()))
             ]
     
+
             if filtered_df.empty:
                 return None
+            
+            # Define colors for sizes and shapes for loss types
+            size_colors = {
+                "14": "blue",
+                "31": "green",
+                "70": "orange",
+                "160": "red"
+            }
+            
+            loss_markers = {
+                "compliment": "o",
+                "cross_entropy": "^",
+                "headless": "s"
+            }
+            
+            # Create the plot
+            fig, ax = plt.subplots(figsize=(10, 6))
+            
+            # Plot each combination of size and loss type
+            for size in filtered_df["param_type"].unique():
+                for loss_type in filtered_df["loss_type"].unique():
+                    data = filtered_df[(filtered_df["param_type"] == size) & (filtered_df["loss_type"] == loss_type)]
+                    ax.plot(data["epoch"], data["loss"], marker=loss_markers[loss_type], color=size_colors[size], label=f"{size} - {loss_type}")
+            
+            # Customize the plot
+            ax.set_xlabel("Epoch")
+            ax.set_ylabel("Loss")
+            ax.set_title("Training Loss by Size and Loss Type", fontsize=16)
+            
+            # Create a legend for sizes
+            size_legend = ax.legend(title="Size", loc="upper right")
+            ax.add_artist(size_legend)
+            
+            # Create a separate legend for loss types
+            loss_legend_labels = ["Compliment", "Cross Entropy", "Headless"]
+            loss_legend_handles = [plt.Line2D([0], [0], marker=loss_markers[loss_type], color='black', linestyle='None', markersize=8) for loss_type in loss_markers]
+            loss_legend = ax.legend(loss_legend_handles, loss_legend_labels, title="Loss Type", loc="upper right")
+            
+            plt.tight_layout()
+            return fig
     
             # Define colors for sizes and shapes for loss types
             size_colors = {