Update app.py

app.py

@@ -60,7 +60,7 @@ def add_new_eval(
         gr.Warning("Your submission has not been processed. Please check your representation files!")
         return -1
 
-    # Even if save is False, store the submission (e.g. temporarily) so that the leaderboard includes it.
+    # Even if save is False, we store the submission (e.g., temporarily) so that the leaderboard includes it.
     if save:
         save_results(representation_name, benchmark_types, results)
     else:
@@ -100,27 +100,31 @@ def generate_plots_based_on_submission(benchmark_types, similarity_tasks, functi
     for btype in benchmark_types:
         # For each benchmark type, choose plotting parameters based on additional selections.
         if btype == "similarity":
+            # Use the user-selected similarity tasks (if provided) to determine the metrics.
             x_metric = similarity_tasks[0] if similarity_tasks and len(similarity_tasks) > 0 else None
             y_metric = similarity_tasks[1] if similarity_tasks and len(similarity_tasks) > 1 else None
         elif btype == "function":
             x_metric = function_prediction_aspect if function_prediction_aspect else None
             y_metric = function_prediction_dataset if function_prediction_dataset else None
         elif btype == "family":
+            # For family, assume that family_prediction_dataset is a list of datasets.
             x_metric = family_prediction_dataset[0] if family_prediction_dataset and len(family_prediction_dataset) > 0 else None
             y_metric = family_prediction_dataset[1] if family_prediction_dataset and len(family_prediction_dataset) > 1 else None
         elif btype == "affinity":
-            x_metric, y_metric = None, None  # Use default plotting for affinity
+            # For affinity, you may use default plotting parameters.
+            x_metric, y_metric = None, None
         else:
             x_metric, y_metric = None, None
 
         # Generate the plot using your benchmark_plot function.
+        # Here, aspect, dataset, and single_metric are passed as None, but you could extend this logic.
         plot_img = benchmark_plot(btype, method_names, x_metric, y_metric, None, None, None)
         plot_file = os.path.join(tmp_dir, f"{btype}.png")
         if isinstance(plot_img, plt.Figure):
             plot_img.savefig(plot_file)
             plt.close(plot_img)
         else:
-            # Assume plot_img is a file path already.
+            # If benchmark_plot already returns a file path, use it directly.
             plot_file = plot_img
         plot_files.append(plot_file)
 
@@ -143,11 +147,11 @@ def submission_callback(
     function_prediction_dataset,
     family_prediction_dataset,
     save_checkbox,
-    return_leaderboard,  # Checkbox: if checked, return leaderboard CSV
-    return_plots       # Checkbox: if checked, return plot results ZIP
+    return_option,  # New radio selection: "Leaderboard CSV" or "Plot Results"
 ):
     """
-    Runs the evaluation and returns files based on selected output options.
+    Runs the evaluation and then returns either a downloadable CSV of the leaderboard
+    (which includes the new submission) or a ZIP file of plots generated based on the submission's selections.
     """
     eval_status = add_new_eval(
         human_file,
@@ -163,26 +167,22 @@ def submission_callback(
     )
 
     if eval_status == -1:
-        return "Submission failed. Please check your files and selections.", None, None
-
-    csv_file = None
-    plots_file = None
-    msg = "Submission processed. "
+        return "Submission failed. Please check your files and selections.", None
 
-    if return_leaderboard:
-        csv_file = download_leaderboard_csv()
-        msg += "Leaderboard CSV is ready. "
-    if return_plots:
-        plots_file = generate_plots_based_on_submission(
+    if return_option == "Leaderboard CSV":
+        csv_path = download_leaderboard_csv()
+        return "Your leaderboard CSV (including your submission) is ready for download.", csv_path
+    elif return_option == "Plot Results":
+        zip_path = generate_plots_based_on_submission(
             benchmark_types,
             similarity_tasks,
             function_prediction_aspect,
             function_prediction_dataset,
             family_prediction_dataset,
         )
-        msg += "Plot results ZIP is ready."
-
-    return msg, csv_file, plots_file
+        return "Your plots are ready for download.", zip_path
+    else:
+        return "Submission processed, but no output option was selected.", None
 
 
 # --------------------------
@@ -195,7 +195,7 @@ with block:
 
     with gr.Tabs(elem_classes="tab-buttons") as tabs:
         with gr.TabItem("🏅 PROBE Leaderboard", elem_id="probe-benchmark-tab-table", id=1):
-            # Leaderboard Tab (unchanged)
+            # Leaderboard tab (unchanged from before)
             leaderboard = get_baseline_df(None, None)
             method_names = leaderboard['Method'].unique().tolist()
             metric_names = leaderboard.columns.tolist()
@@ -266,7 +266,7 @@ with block:
                     Select options to update the visualization.
                     """
                 )
-            # Plotting section remains available as before.
+            # (Plotting section remains available as before; not the focus of the submission callback)
             benchmark_type_selector_plot = gr.Dropdown(
                 choices=list(benchmark_specific_metrics.keys()),
                 label="Select Benchmark Type for Plotting",
@@ -346,15 +346,6 @@ with block:
                         label="Save results for leaderboard and visualization",
                         value=True
                     )
-                    # New independent checkboxes for output return options:
-                    return_leaderboard = gr.Checkbox(
-                        label="Return Leaderboard CSV",
-                        value=False
-                    )
-                    return_plots = gr.Checkbox(
-                        label="Return Plot Results",
-                        value=False
-                    )
             with gr.Row():
                 human_file = gr.components.File(
                     label="The representation file (csv) for Human dataset",
@@ -366,11 +357,16 @@ with block:
                     file_count="single",
                     type='filepath'
                 )
+            # New radio button for output selection.
+            return_option = gr.Radio(
+                choices=["Leaderboard CSV", "Plot Results"],
+                label="Return Output",
+                value="Leaderboard CSV",
+                interactive=True,
+            )
             submit_button = gr.Button("Submit Eval")
             submission_result_msg = gr.Markdown()
-            # Two file outputs: one for CSV, one for Plot ZIP.
-            submission_csv_file = gr.File(label="Leaderboard CSV", visible=True)
-            submission_plots_file = gr.File(label="Plot Results ZIP", visible=True)
+            submission_result_file = gr.File()
             submit_button.click(
                 submission_callback,
                 inputs=[
@@ -384,10 +380,9 @@ with block:
                     function_dataset,
                     family_prediction_dataset,
                     save_checkbox,
-                    return_leaderboard,
-                    return_plots,
+                    return_option,
                 ],
-                outputs=[submission_result_msg, submission_csv_file, submission_plots_file]
+                outputs=[submission_result_msg, submission_result_file]
             )
 
     with gr.Row():