Update app.py

app.py

@@ -266,7 +266,7 @@ with ui.navset_card_tab(id="tab"):
             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
+           # axs = axs.flatten()  # Flatten axs to handle 1D indexing
         
             for size in model_bigness:
                 for loss in loss_type:
@@ -274,11 +274,16 @@ with ui.navset_card_tab(id="tab"):
                     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()
+                            if len(logits) == 1:
+                                logit_index = int(logit) - 1
+                                axs[0].plot(data['lm_logits_y_cumsum'][0, :, logit_index], label=f'Generated_{loss}_{size}')
+                                axs[0].plot(data['shift_labels_y_cumsum'][0, :, logit_index], label=f'Expected_{loss}_{size}')
+                                axs[0].set_title(f'Logit: {logit}- Single')
+                                axs[0].legend()
+                                axs[1].plot(data['lm_logits_y_full_cumsum'][0, :, logit_index], label=f'Generated_{loss}_{size}')
+                                axs[1].plot(data['shift_labels_y_full_cumsum'][0, :, logit_index], label=f'Expected_{loss}_{size}')
+                                axs[1].set_title(f'Logit: {logit} - Full')
+                                axs[1].legend()
                     else:
                         print(f"File not found: {file_name}")