Update app.py

app.py

@@ -973,6 +973,92 @@ with ui.navset_card_tab(id="tab"):
             mpl.rcParams.update(mpl.rcParamsDefault)
             fig = plot_loss_rates(df, '14M')
             return fig
+    with ui.nav_panel("Model loss analysis"):
+        ui.page_opts(fillable=True)
+        ui.panel_title("Neurips stuff")
+    
+        with ui.card():
+                param_type = ui.input_selectize(
+                    "param_type",
+                    "Select Param Type:",
+                    ["14", "31", "70", "160"],
+                    multiple=True,
+                    selected=None
+                )
+    
+        with ui.card():
+                model_type = ui.input_selectize(
+                    "model_type",
+                    "Select Model Type:",
+                    ["pythia", "denseformer"],
+                    multiple=True,
+                    selected=None
+                )
+    
+        with ui.card():
+                loss_type = ui.input_selectize(
+                    "loss_type",
+                    "Select Loss Type:",
+                    ["compliment", "cross_entropy", "headless"],
+                    multiple=True,
+                    selected=None
+                )
+    
+        @output
+        @render.plot
+        def plot_training_loss():
+            if csv_file() is None:
+                return None
+            
+            df = pd.read_csv(csv_file())
+    
+            filtered_df = df[
+                (df["param_type"].isin(param_type()))
+                & (df["model_type"].isin(model_type()))
+                & (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 a relplot using Seaborn
+            g = sns.relplot(
+                data=filtered_df,
+                x="epoch",
+                y="loss",
+                hue="param_type",
+                style="loss_type",
+                palette=size_colors,
+                markers=loss_markers,
+                height=6,
+                aspect=1.5
+            )
+    
+            # Customize the plot
+            g.set_xlabels("Epoch")
+            g.set_ylabels("Loss")
+            g.fig.suptitle("Training Loss by Size and Loss Type", fontsize=16)
+            g.add_legend(title="Size")
+    
+            # Create a separate legend for loss types
+            loss_legend = plt.legend(title="Loss Type", loc="upper right", labels=["Compliment", "Cross Entropy", "Headless"])
+            plt.gca().add_artist(loss_legend)
+    
+            plt.tight_layout()
+            return g.fig
 
         
 # @render.image