Update app.py

app.py

@@ -237,3 +237,61 @@ with ui.navset_card_tab(id="tab"):
             )
             if fig:
                 return fig
+    with ui.nav_panel("Logits View"):
+        ui.panel_title("Logits et all")
+        with ui.card():
+            ui.input_selectize(
+                "model_bigness",
+                "Select Model size:",
+                ["14", "31", "70", "160", "410"],
+                multiple=True,
+                selected=["70", "160"],
+            )
+            ui.input_selectize(
+                "loss_loss_loss",
+                "Select Loss Type:",
+                ["compliment", "cross_entropy", "headless", "2d_representation_GaussianPlusCE", "2d_representation_MSEPlusCE"],
+                multiple=True,
+                selected=["cross_entropy"],
+            )
+            ui.input_selectize(
+                "logits_select",
+                "Select logits:",
+                ["1", "2", "3", "4", "5", "6", "7", "8"],
+                multiple=True,
+                selected=["6"],
+            )
+        
+        def plot_logits_representation(model_bigness, loss_type, logits):
+            num_rows = 2  # Number of rows in the subplot grid
+            num_cols = len(logits)  # Number of columns based on the number of selected logits
+            fig, axs = plt.subplots(num_rows, num_cols, figsize=(20, 10))
+            axs = axs.flatten()  # Flatten axs to handle 1D indexing
+        
+            for size in model_bigness:
+                for loss in loss_type:
+                    file_name = f"virus_pythia_{size}_1024_{loss}_logit_cumsums.npy"
+                    if os.path.exists(file_name):
+                        data = np.load(file_name, allow_pickle=True).item()
+                        for k, logit in enumerate(logits):
+                            logit_index = int(logit) - 1
+                            axs[k].plot(data['lm_logits_y_cumsum'][0, :, logit_index], label=f'Generated_{loss}_{size}')
+                            axs[k].plot(data['shift_labels_y_cumsum'][0, :, logit_index], label=f'Expected_{loss}_{size}')
+                            axs[k].set_title(f'Logit: {logit}')
+                            axs[k].legend()
+                    else:
+                        print(f"File not found: {file_name}")
+        
+            for k in range(len(logits), num_cols):
+                fig.delaxes(axs[k])  # Remove any extra subplots if fewer logits are selected
+        
+            plt.tight_layout()
+            return fig
+
+        @render.plot()
+        def plot_logits_representation_ui():
+            fig = plot_logits_representation(
+                input.model_bigness(), input.loss_loss_loss(), input.logits_select()
+            )
+            if fig:
+                return fig