Update app.py

app.py

@@ -136,9 +136,9 @@ with gr.Blocks(css="""
     h1 {
         text-align: center;
     }
-    .section-spacing {
-        margin-top: 30px;
-        margin-bottom: 30px;
+    .section-divider {
+        border-top: 1px solid #ddd;
+        margin: 12px 0;
     }
     .config-box {
         border: 1px solid #ddd;
@@ -158,7 +158,7 @@ with gr.Blocks(css="""
     """)
     
     # Dataset Selection Section
-    gr.Markdown("## (A) Select Dataset for Evaluation", elem_classes=["section-spacing"])
+    gr.Markdown("## (A) Select Dataset for Evaluation")
     
     with gr.Row():
         dataset_dropdown = gr.Dropdown(
@@ -177,20 +177,15 @@ with gr.Blocks(css="""
             wrap=True,
             elem_id="preview_table"
         )
-        # Add vertical space after the preview
-        gr.Markdown(" ")
-        gr.Markdown(" ")
+        # Add a divider instead of lots of space
+        gr.Markdown("<div class='section-divider'></div>")
     
-    # Add more spacing between sections
-    gr.Markdown("&nbsp;", elem_classes=["section-spacing"])
-    gr.Markdown("&nbsp;", elem_classes=["section-spacing"])
+    # Preview data state to store the loaded preview data
+    preview_data_state = gr.State(None)
     
     # 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", elem_classes=["section-spacing"])
-        
-        # Add more spacing
-        gr.Markdown("&nbsp;")
+        gr.Markdown("## (B) Select Dataset Configuration Options")
         
         with gr.Row():
             # Left column for subject selection
@@ -208,27 +203,18 @@ with gr.Blocks(css="""
                     with gr.Column(visible=False) as num_subjects_container:
                         num_subjects_slider = gr.Slider(
                             minimum=1,
-                            maximum=14,
+                            maximum=14,  # Will be updated dynamically based on preview data
                             value=14,
                             step=1,
                             label="Number of Subjects",
-                            info="Number of subjects to evaluate (1-14). They will be loaded in alphabetical order."
+                            info="Number of subjects to evaluate. 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
+                        # Will be populated dynamically from the preview data
                         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)"
-                            ],
+                            choices=[],  # Will be populated from preview data
                             label="Select Specific Subjects",
                             info="Select which specific subjects to evaluate"
                         )
@@ -247,7 +233,7 @@ with gr.Blocks(css="""
                         info="Number of examples to use for few-shot learning (0-5)."
                     )
         
-        # Add spacing
+        # Add a small space
         gr.Markdown("&nbsp;")
         
         with gr.Row():
@@ -296,6 +282,39 @@ with gr.Blocks(css="""
     # Track preview visibility state
     preview_visibility = gr.State(False)
     
+    # Function to process dataset preview data
+    def get_subject_choices_from_preview(preview_data):
+        if not preview_data or 'subject_counts' not in preview_data:
+            return [], 0
+        
+        # Get subject counts from preview data
+        subject_counts = preview_data['subject_counts']
+        
+        # Sort subjects alphabetically
+        subjects = sorted(subject_counts.keys())
+        
+        # Format as "Subject (n=count)"
+        formatted_subjects = [f"{subject} (n={subject_counts[subject]})" for subject in subjects]
+        
+        return formatted_subjects, len(subjects)
+    
+    # Function to load preview data and update UI
+    def load_dataset_preview(dataset):
+        if dataset == "MMLU-Pro":
+            # Load the preview data
+            preview_data = mmlupro_dataset_preview(regenerate_preview=False)
+            
+            # Extract subject choices and count
+            subject_choices, subject_count = get_subject_choices_from_preview(preview_data)
+            
+            # Update the UI components
+            return (
+                preview_data,  # Store the preview data
+                gr.update(choices=subject_choices),  # Update checkbox choices
+                gr.update(maximum=subject_count, value=min(subject_count, 14))  # Update slider max
+            )
+        return None, gr.update(), gr.update()
+    
     # Function to show/hide configuration based on selected dataset
     def update_interface_based_on_dataset(dataset, current_visibility):
         if dataset == "MMLU-Pro":
@@ -317,24 +336,27 @@ with gr.Blocks(css="""
                 gr.update(value="Show Dataset Preview")  # Reset button text
             )
     
-    # Connect dataset dropdown to show/hide appropriate configuration
+    # Connect dataset dropdown to show/hide appropriate configuration and load preview data
     dataset_dropdown.change(
+        fn=load_dataset_preview,
+        inputs=[dataset_dropdown],
+        outputs=[preview_data_state, specific_subjects, num_subjects_slider],
+    ).then(
         fn=update_interface_based_on_dataset,
         inputs=[dataset_dropdown, preview_visibility],
         outputs=[mmlu_config_container, results_container, preview_toggle, dataset_preview_container, preview_visibility, preview_toggle]
     )
     
     # Function to toggle dataset preview visibility
-    def toggle_preview(dataset, preview_visibility):
+    def toggle_preview(dataset, preview_visibility, preview_data):
         # Toggle the visibility state
         is_visible = not preview_visibility
         
         # Update button text based on new state
         button_text = "Hide Dataset Preview" if is_visible else "Show Dataset Preview"
         
-        # Get preview data if becoming visible
+        # Format and show preview if becoming visible
         if is_visible and dataset == "MMLU-Pro":
-            preview_data = mmlupro_dataset_preview(regenerate_preview=False) # Change regenerate_preview=True if you want to regenerate the preview.
             formatted_preview = format_preview_for_display(preview_data)
             return is_visible, gr.update(visible=True), formatted_preview, gr.update(value=button_text)
         elif is_visible:
@@ -347,7 +369,7 @@ with gr.Blocks(css="""
     # Connect preview toggle to show/hide dataset information
     preview_toggle.click(
         fn=toggle_preview,
-        inputs=[dataset_dropdown, preview_visibility],
+        inputs=[dataset_dropdown, preview_visibility, preview_data_state],
         outputs=[preview_visibility, dataset_preview_container, preview_output, preview_toggle]
     )