Update app.py

app.py

@@ -24,13 +24,14 @@ model_name = "mistralai/Mistral-7B-v0.1"
 # 2. MMLU-Pro Evaluation
 # ---------------------------------------------------------------------------
 @spaces.GPU(duration=240)
-def run_mmlu_evaluation(all_subjects, num_subjects, num_shots, all_questions, num_questions, progress=gr.Progress()):
+def run_mmlu_evaluation(subject_selection_mode, num_subjects, selected_subjects, num_shots, all_questions, num_questions, progress=gr.Progress()):
     """
     Runs the MMLU evaluation with the specified parameters.
     
     Args:
-        all_subjects (bool): Whether to evaluate all subjects
+        subject_selection_mode (str): Mode of subject selection ("all", "number", or "specific")
         num_subjects (int): Number of subjects to evaluate (1-14)
+        selected_subjects (list): List of specific subjects to evaluate
         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-100 or all)
@@ -38,8 +39,11 @@ def run_mmlu_evaluation(all_subjects, num_subjects, num_shots, all_questions, nu
     """
     try:
         # Convert parameters if needed
-        if all_subjects:
+        if subject_selection_mode == "all":
             num_subjects = -1
+            selected_subjects = []
+        elif subject_selection_mode == "specific":
+            num_subjects = len(selected_subjects) if selected_subjects else -1
             
         if all_questions:
             num_questions = -1
@@ -50,7 +54,8 @@ def run_mmlu_evaluation(all_subjects, num_subjects, num_shots, all_questions, nu
             model_name,
             num_subjects=num_subjects,
             num_questions=num_questions,
-            num_shots=num_shots, 
+            num_shots=num_shots,
+            specific_subjects=selected_subjects if subject_selection_mode == "specific" else None
         )
         elapsed_time = time.time() - start_time
 
@@ -131,6 +136,17 @@ with gr.Blocks(css="""
     h1 {
         text-align: center;
     }
+    .section-spacing {
+        margin-top: 30px;
+        margin-bottom: 30px;
+    }
+    .config-box {
+        border: 1px solid #ddd;
+        border-radius: 8px;
+        padding: 15px;
+        margin: 10px;
+        background-color: #f9f9f9;
+    }
 """) as demo:
     gr.Markdown("# Head-to-Head Model Evaluation Comparator")
     gr.Markdown("""
@@ -142,7 +158,7 @@ with gr.Blocks(css="""
     """)
     
     # Dataset Selection Section
-    gr.Markdown("## (A) Select Dataset for Evaluation")
+    gr.Markdown("## (A) Select Dataset for Evaluation", elem_classes=["section-spacing"])
     
     with gr.Row():
         dataset_dropdown = gr.Dropdown(
@@ -165,35 +181,74 @@ with gr.Blocks(css="""
         gr.Markdown(" ")
         gr.Markdown(" ")
     
+    # Add more spacing between sections
+    gr.Markdown(" ", elem_classes=["section-spacing"])
+    gr.Markdown(" ", elem_classes=["section-spacing"])
+    
     # MMLU Config Container - Initially hidden until dataset is selected
     with gr.Column(visible=False) as mmlu_config_container:
-        gr.Markdown("## (B) Select Dataset Configuration Options")
+        gr.Markdown("## (B) Select Dataset Configuration Options", elem_classes=["section-spacing"])
         
-        with gr.Row():
-            all_subjects_checkbox = gr.Checkbox(
-                label="Evaluate All Subjects",
-                value=False,
-                info="When checked, evaluates all 14 MMLU-Pro subjects"
-            )
-            num_subjects_slider = gr.Slider(
-                minimum=1,
-                maximum=14,
-                value=14,
-                step=1,
-                label="Number of Subjects",
-                info="Number of subjects to evaluate (1-14). They will be loaded in alphabetical order.",
-                interactive=True
-            )
+        # Add more spacing
+        gr.Markdown(" ")
         
         with gr.Row():
-            num_shots_slider = gr.Slider(
-                minimum=0,
-                maximum=5,
-                value=5,
-                step=1,
-                label="Number of Few-shot Examples",
-                info="Number of examples to use for few-shot learning (0-5)."
-            )
+            # Left column for subject selection
+            with gr.Column(scale=1):
+                with gr.Box(elem_classes=["config-box"]):
+                    gr.Markdown("### Choose Subjects")
+                    
+                    subject_selection_mode = gr.Radio(
+                        choices=["Evaluate All Subjects", "Choose Number of Subjects", "Specify which Subjects to Evaluate"],
+                        value="Evaluate All Subjects",
+                        label="Subject Selection Mode"
+                    )
+                    
+                    # Subject number slider - initially hidden, shown when "Choose Number of Subjects" is selected
+                    with gr.Column(visible=False) as num_subjects_container:
+                        num_subjects_slider = gr.Slider(
+                            minimum=1,
+                            maximum=14,
+                            value=14,
+                            step=1,
+                            label="Number of Subjects",
+                            info="Number of subjects to evaluate (1-14). They will be loaded in alphabetical order."
+                        )
+                    
+                    # Subject checkboxes - initially hidden, shown when "Specify which Subjects to Evaluate" is selected
+                    with gr.Column(visible=False) as specific_subjects_container:
+                        # We'll populate this with checkboxes for each subject
+                        # The actual subjects will come from the dataset preview
+                        specific_subjects = gr.CheckboxGroup(
+                            choices=[
+                                "Biology (n=717)",
+                                "Chemistry (n=500)",
+                                "Physics (n=650)",
+                                "Mathematics (n=800)",
+                                "Computer Science (n=450)",
+                                "History (n=300)",
+                                "Literature (n=250)"
+                            ],
+                            label="Select Specific Subjects",
+                            info="Select which specific subjects to evaluate"
+                        )
+            
+            # Right column for few-shot examples
+            with gr.Column(scale=1):
+                with gr.Box(elem_classes=["config-box"]):
+                    gr.Markdown("### Few-shot Configuration")
+                    
+                    num_shots_slider = gr.Slider(
+                        minimum=0,
+                        maximum=5,
+                        value=5,
+                        step=1,
+                        label="Number of Few-shot Examples",
+                        info="Number of examples to use for few-shot learning (0-5)."
+                    )
+        
+        # Add spacing
+        gr.Markdown(" ")
         
         with gr.Row():
             all_questions_checkbox = gr.Checkbox(
@@ -296,14 +351,20 @@ with gr.Blocks(css="""
         outputs=[preview_visibility, dataset_preview_container, preview_output, preview_toggle]
     )
     
-    # Update num_subjects_slider interactivity based on all_subjects checkbox
-    def update_subjects_slider(checked):
-        return gr.update(interactive=not checked)
+    # Function to update UI based on subject selection mode
+    def update_subject_selection_ui(mode):
+        if mode == "Evaluate All Subjects":
+            return gr.update(visible=False), gr.update(visible=False)
+        elif mode == "Choose Number of Subjects":
+            return gr.update(visible=True), gr.update(visible=False)
+        else:  # "Specify which Subjects to Evaluate"
+            return gr.update(visible=False), gr.update(visible=True)
     
-    all_subjects_checkbox.change(
-        fn=update_subjects_slider,
-        inputs=[all_subjects_checkbox],
-        outputs=[num_subjects_slider]
+    # Connect subject selection mode to UI updates
+    subject_selection_mode.change(
+        fn=update_subject_selection_ui,
+        inputs=[subject_selection_mode],
+        outputs=[num_subjects_container, specific_subjects_container]
     )
     
     # Update interface based on all_questions checkbox
@@ -319,6 +380,20 @@ with gr.Blocks(css="""
         outputs=[questions_container, questions_info_text]
     )
     
+    # Function to convert subject selection mode to parameters
+    def get_subject_mode_param(mode):
+        if mode == "Evaluate All Subjects":
+            return "all"
+        elif mode == "Choose Number of Subjects":
+            return "number"
+        else:  # "Specify which Subjects to Evaluate"
+            return "specific"
+    
+    # Function to extract subject names from checkboxes
+    def get_subject_names(selected_subjects):
+        # Extract just the subject name without the count
+        return [subject.split(" (")[0] for subject in selected_subjects]
+    
     # Function to disable UI components during evaluation
     def start_evaluation(state):
         if state["running"]:
@@ -341,8 +416,9 @@ with gr.Blocks(css="""
         
         return [
             state,
-            gr.update(interactive=False),  # all_subjects_checkbox
+            gr.update(interactive=False),  # subject_selection_mode
             gr.update(interactive=False),  # num_subjects_slider
+            gr.update(interactive=False),  # specific_subjects
             gr.update(interactive=False),  # num_shots_slider
             gr.update(interactive=False),  # all_questions_checkbox
             gr.update(interactive=False),  # num_questions_slider
@@ -365,8 +441,9 @@ with gr.Blocks(css="""
         state["running"] = False
         return [
             state,
-            gr.update(interactive=True),  # all_subjects_checkbox
+            gr.update(interactive=True),  # subject_selection_mode
             gr.update(interactive=True),  # num_subjects_slider
+            gr.update(interactive=True),  # specific_subjects
             gr.update(interactive=True),  # num_shots_slider
             gr.update(interactive=True),  # all_questions_checkbox
             gr.update(interactive=True),  # num_questions_slider
@@ -383,8 +460,9 @@ with gr.Blocks(css="""
         inputs=[evaluation_state],
         outputs=[
             evaluation_state,
-            all_subjects_checkbox,
+            subject_selection_mode,
             num_subjects_slider,
+            specific_subjects,
             num_shots_slider,
             all_questions_checkbox,
             num_questions_slider,
@@ -395,10 +473,19 @@ with gr.Blocks(css="""
             results_table_container
         ]
     ).then(
-        fn=run_mmlu_evaluation,
+        fn=lambda mode, num, subjects, shots, all_q, num_q: 
+            run_mmlu_evaluation(
+                get_subject_mode_param(mode),
+                num,
+                get_subject_names(subjects),
+                shots,
+                all_q,
+                num_q
+            ),
         inputs=[
-            all_subjects_checkbox,
+            subject_selection_mode,
             num_subjects_slider,
+            specific_subjects,
             num_shots_slider,
             all_questions_checkbox,
             num_questions_slider
@@ -408,7 +495,7 @@ with gr.Blocks(css="""
             results_table,
             eval_mmlu_button, 
             cancel_mmlu_button,
-            all_subjects_checkbox,
+            subject_selection_mode,
             num_subjects_slider,
             num_shots_slider,
             all_questions_checkbox,
@@ -427,8 +514,9 @@ with gr.Blocks(css="""
         inputs=[evaluation_state],
         outputs=[
             evaluation_state,
-            all_subjects_checkbox,
+            subject_selection_mode,
             num_subjects_slider,
+            specific_subjects,
             num_shots_slider,
             all_questions_checkbox,
             num_questions_slider,