Update app.py

app.py

@@ -68,7 +68,7 @@ def run_toy_evaluation():
 # 3. MMLU Evaluation call
 # ---------------------------------------------------------------------------
 @spaces.GPU(duration=120)  # Allow up to 2 minutes for full evaluation
-def run_mmlu_evaluation(all_subjects, num_subjects, num_shots, all_questions, num_questions):
+def run_mmlu_evaluation(all_subjects, num_subjects, num_shots, all_questions, num_questions, progress=gr.Progress()):
     """
     Runs the MMLU evaluation with the specified parameters.
     
@@ -78,13 +78,15 @@ def run_mmlu_evaluation(all_subjects, num_subjects, num_shots, all_questions, nu
         num_shots (int): Number of few-shot examples (0-5)
         all_questions (bool): Whether to evaluate all questions per subject
         num_questions (int): Number of examples per subject (1-20 or -1 for all)
+        progress (gr.Progress): Progress indicator
     """
 
     if not model_loaded:
         load_model()
 
     if not model_loaded:
-        return "⚠️ Model not loaded. Please load the model first.", None
+        return "⚠️ Model not loaded. Please load the model first.", None, gr.update(interactive=True), gr.update(visible=False), \
+              [gr.update(interactive=True) for _ in range(5)]
 
     # Convert num_subjects to -1 if all_subjects is True
     if all_subjects:
@@ -142,7 +144,9 @@ def run_mmlu_evaluation(all_subjects, num_subjects, num_shots, all_questions, nu
         f"* Evaluation completed in {elapsed_time:.2f} seconds\n"
     )
 
-    return report, results_df  # Return both text report and dataframe
+    # Return values that re-enable UI components after completion
+    return report, results_df, gr.update(interactive=True), gr.update(visible=False), \
+           [gr.update(interactive=True) for _ in range(5)]
     
 # ---------------------------------------------------------------------------
 # 4. Gradio Interface
@@ -200,20 +204,31 @@ with gr.Blocks() as demo:
             value=False,  # Default is unchecked
             info="When checked, evaluates all available questions for each subject"
         )
+        questions_info_text = gr.Markdown(visible=False, value="**All 14,042 questions across all subjects will be evaluated**")
+    
+    with gr.Row(elem_id="questions_selection_row"):
+        questions_container = gr.Column(scale=1, elem_id="questions_slider_container")
+    
+    # Move the slider into the container for easier visibility toggling
+    with questions_container:
         num_questions_slider = gr.Slider(
             minimum=1,
             maximum=20,
             value=10,  # Default is 10 questions
             step=1,
             label="Questions per Subject",
-            info="Choose a subset of questions (1-20), or click the checkbox for All Questions",
+            info="Choose a subset of questions (1-20)",
             interactive=True
         )
     
     with gr.Row():
-        eval_mmlu_button = gr.Button("Run MMLU Evaluation", variant="primary")
+        with gr.Column(scale=1):
+            eval_mmlu_button = gr.Button("Run MMLU Evaluation", variant="primary", interactive=True)
+            cancel_mmlu_button = gr.Button("Cancel MMLU Evaluation", variant="stop", visible=False)
         results_output = gr.Markdown(label="Evaluation Results")
-        results_table = gr.DataFrame(interactive=True, label="Detailed Results (Sortable)")
+        
+    with gr.Row():
+        results_table = gr.DataFrame(interactive=True, label="Detailed Results (Sortable)", visible=True)
     
     # Connect components
     load_button.click(fn=load_model, inputs=None, outputs=load_status)
@@ -238,21 +253,61 @@ with gr.Blocks() as demo:
         outputs=[num_subjects_slider]
     )
     
-    # Update num_questions_slider interactivity based on all_questions checkbox
-    def update_questions_slider(checked):
+    # Update interface based on all_questions checkbox
+    def update_questions_interface(checked):
         if checked:
-            return gr.update(interactive=False)
+            return gr.update(visible=False), gr.update(visible=True)
         else:
-            return gr.update(interactive=True)
+            return gr.update(visible=True), gr.update(visible=False)
     
     all_questions_checkbox.change(
-        fn=update_questions_slider,
+        fn=update_questions_interface,
         inputs=[all_questions_checkbox],
-        outputs=[num_questions_slider]
+        outputs=[questions_container, questions_info_text]
     )
     
-    # Connect MMLU evaluation button
+    # Function to disable UI components during evaluation
+    def disable_ui_for_evaluation():
+        return [
+            gr.update(interactive=False, info="MMLU Evaluation currently in progress"),  # all_subjects_checkbox
+            gr.update(interactive=False, info="MMLU Evaluation currently in progress"),  # num_subjects_slider
+            gr.update(interactive=False, info="MMLU Evaluation currently in progress"),  # num_shots_slider
+            gr.update(interactive=False, info="MMLU Evaluation currently in progress"),  # all_questions_checkbox
+            gr.update(interactive=False, info="MMLU Evaluation currently in progress"),  # num_questions_slider
+            gr.update(interactive=False),  # eval_mmlu_button
+            gr.update(visible=True)   # cancel_mmlu_button
+        ]
+    
+    # Function to handle cancel button click
+    def cancel_evaluation():
+        # This doesn't actually cancel the GPU job (which would require more backend support)
+        # But it does reset the UI state to be interactive again
+        return [
+            gr.update(interactive=True, info="When checked, evaluates all 57 MMLU subjects"),  # all_subjects_checkbox
+            gr.update(interactive=True, info="Number of subjects to evaluate (1-57). They will be loaded in alphabetical order."),  # num_subjects_slider
+            gr.update(interactive=True, info="Number of examples to use for few-shot learning (0-5). They will be loaded in alphabetical order."),  # num_shots_slider
+            gr.update(interactive=True, info="When checked, evaluates all available questions for each subject"),  # all_questions_checkbox
+            gr.update(interactive=True, info="Choose a subset of questions (1-20)"),  # num_questions_slider
+            gr.update(interactive=True),  # eval_mmlu_button
+            gr.update(visible=False),  # cancel_mmlu_button
+            "⚠️ Evaluation canceled by user", # results_output
+            None  # results_table
+        ]
+    
+    # Connect MMLU evaluation button - now disables UI and shows cancel button
     eval_mmlu_button.click(
+        fn=disable_ui_for_evaluation,
+        inputs=None,
+        outputs=[
+            all_subjects_checkbox,
+            num_subjects_slider,
+            num_shots_slider,
+            all_questions_checkbox,
+            num_questions_slider,
+            eval_mmlu_button,
+            cancel_mmlu_button
+        ]
+    ).then(
         fn=run_mmlu_evaluation,
         inputs=[
             all_subjects_checkbox,
@@ -261,7 +316,30 @@ with gr.Blocks() as demo:
             all_questions_checkbox,
             num_questions_slider
         ],
-        outputs=[results_output, results_table]
+        outputs=[
+            results_output,
+            results_table,
+            eval_mmlu_button,
+            cancel_mmlu_button,
+            [all_subjects_checkbox, num_subjects_slider, num_shots_slider, all_questions_checkbox, num_questions_slider]
+        ]
+    )
+    
+    # Connect cancel button
+    cancel_mmlu_button.click(
+        fn=cancel_evaluation,
+        inputs=None,
+        outputs=[
+            all_subjects_checkbox,
+            num_subjects_slider,
+            num_shots_slider,
+            all_questions_checkbox,
+            num_questions_slider,
+            eval_mmlu_button,
+            cancel_mmlu_button,
+            results_output,
+            results_table
+        ]
     )
 
-demo.launch()
+demo.launch()