Update app.py

app.py

@@ -1,109 +1,111 @@
 import gradio as gr
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
 from huggingface_hub import login
-from toy_dataset_eval import evaluate_toy_dataset
 from mmlu_pro_eval_adapted import evaluate_mmlu_pro
 import spaces
 import pandas as pd
-import time  # Added for timing functionality
+import time
+import traceback
 
 # Read token and login
 hf_token = os.getenv("HF_READ_WRITE_TOKEN")
 if hf_token:
     login(hf_token)
 else:
-    print("⚠️ No HF_TOKEN_READ_WRITE found in environment")
+    print("⚠️ No HF_READ_WRITE_TOKEN found in environment")
 
 # ---------------------------------------------------------------------------
-# 1. Model and tokenizer setup and Loading
+# 1. Model configuration
 # ---------------------------------------------------------------------------
 model_name = "mistralai/Mistral-7B-v0.1"
-tokenizer = None
-model = None
-model_loaded = False
-    
+
 # ---------------------------------------------------------------------------
-# 1. MMLU-Pro Evaluation call
+# 2. MMLU-Pro Evaluation
 # ---------------------------------------------------------------------------
-@spaces.GPU(duration=120)  # Allow up to 2 minutes for full evaluation
+@spaces.GPU(duration=180)  # Extended to 3 minutes for larger evaluations
 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.
     
     Args:
         all_subjects (bool): Whether to evaluate all subjects
-        num_subjects (int): Number of subjects to evaluate (1-57)
+        num_subjects (int): Number of subjects to evaluate (1-14)
         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)
+        num_questions (int): Number of examples per subject (1-40 or all)
         progress (gr.Progress): Progress indicator
     """
-    
-    # Convert num_subjects to -1 if all_subjects is True
-    if all_subjects:
-        num_subjects = -1
-        
-    # Convert num_questions to -1 if all_questions is True
-    if all_questions:
-        num_questions = -1
+    try:
+        # Convert parameters if needed
+        if all_subjects:
+            num_subjects = -1
+            
+        if all_questions:
+            num_questions = -1
 
-    # Run evaluation with timing
-    start_time = time.time()  # Start timing
-    results = evaluate_mmlu_pro(
-        model_name,
-        num_subjects=num_subjects,
-        num_questions=num_questions,
-        num_shots=num_shots, 
-    )
-    elapsed_time = time.time() - start_time  # Calculate elapsed time
+        # Run evaluation with timing
+        start_time = time.time()
+        results = evaluate_mmlu_pro(
+            model_name,
+            num_subjects=num_subjects,
+            num_questions=num_questions,
+            num_shots=num_shots, 
+        )
+        elapsed_time = time.time() - start_time
 
-    # Format results
-    overall_acc = results["overall_accuracy"]
-    min_subject, min_acc = results["min_accuracy_subject"]
-    max_subject, max_acc = results["max_accuracy_subject"]
-    
-    # Create DataFrame from results table
-    results_df = pd.DataFrame(results["full_accuracy_table"])
-    
-    # Calculate totals for the overall row
-    total_samples = results_df['Num_samples'].sum()
-    total_correct = results_df['Num_correct'].sum()
-    
-    # Create overall row
-    overall_row = pd.DataFrame({
-        'Subject': ['**Overall**'],
-        'Num_samples': [total_samples],
-        'Num_correct': [total_correct],
-        'Accuracy': [overall_acc]
-    })
-    
-    # Concatenate overall row with results
-    results_df = pd.concat([overall_row, results_df], ignore_index=True)
-    
-    # Verify that the overall accuracy is consistent with the total correct/total samples
-    assert abs(overall_acc - (total_correct / total_samples)) < 1e-6, \
-        "Overall accuracy calculation mismatch detected"
-    
-    # Format the report
-    report = (
-        f"### Overall Results\n"
-        f"* Overall Accuracy: {overall_acc:.3f}\n"
-        f"* Best Performance: {max_subject} ({max_acc:.3f})\n"
-        f"* Worst Performance: {min_subject} ({min_acc:.3f})\n"
-        f"* Evaluation completed in {elapsed_time:.2f} seconds\n"
-    )
+        # Format results
+        overall_acc = results["overall_accuracy"]
+        min_subject, min_acc = results["min_accuracy_subject"]
+        max_subject, max_acc = results["max_accuracy_subject"]
+        
+        # Create DataFrame from results table
+        results_df = pd.DataFrame(results["full_accuracy_table"])
+        
+        # Calculate totals for the overall row
+        total_samples = results_df['Num_samples'].sum()
+        total_correct = results_df['Num_correct'].sum()
+        
+        # Create overall row
+        overall_row = pd.DataFrame({
+            'Subject': ['**Overall**'],
+            'Num_samples': [total_samples],
+            'Num_correct': [total_correct],
+            'Accuracy': [overall_acc]
+        })
+        
+        # Concatenate overall row with results
+        results_df = pd.concat([overall_row, results_df], ignore_index=True)
+        
+        # Format the report
+        report = (
+            f"### Overall Results\n"
+            f"* Overall Accuracy: {overall_acc:.3f}\n"
+            f"* Best Performance: {max_subject} ({max_acc:.3f})\n"
+            f"* Worst Performance: {min_subject} ({min_acc:.3f})\n"
+            f"* Evaluation completed in {elapsed_time:.2f} seconds\n"
+        )
 
-    # 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), gr.update(interactive=True), 
-            gr.update(interactive=True), gr.update(interactive=True), 
-            gr.update(interactive=True))
+        # 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), gr.update(interactive=True), 
+                gr.update(interactive=True), gr.update(interactive=True), 
+                gr.update(interactive=True))
+                
+    except Exception as e:
+        # Handle errors gracefully
+        error_trace = traceback.format_exc()
+        error_message = f"### Error during evaluation\n```\n{error_trace}\n```"
+        
+        # Re-enable UI components on error
+        return (error_message, None,
+                gr.update(interactive=True), gr.update(visible=False),
+                gr.update(interactive=True), gr.update(interactive=True), 
+                gr.update(interactive=True), gr.update(interactive=True), 
+                gr.update(interactive=True))
     
 # ---------------------------------------------------------------------------
-# 4. Gradio Interface
+# 3. Gradio Interface
 # ---------------------------------------------------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("# Mistral-7B on MMLU-Pro Evaluation Demo")
@@ -117,13 +119,13 @@ with gr.Blocks() as demo:
     with gr.Row():
         all_subjects_checkbox = gr.Checkbox(
             label="Evaluate All Subjects",
-            value=False,  # Default is unchecked
+            value=False,
             info="When checked, evaluates all 14 MMLU-Pro subjects"
         )
         num_subjects_slider = gr.Slider(
             minimum=1,
             maximum=14,
-            value=14,  # Default is all subjects
+            value=14,
             step=1,
             label="Number of Subjects",
             info="Number of subjects to evaluate (1-14). They will be loaded in alphabetical order.",
@@ -134,16 +136,16 @@ with gr.Blocks() as demo:
         num_shots_slider = gr.Slider(
             minimum=0,
             maximum=5,
-            value=5,  # Default is 5 few-shot examples
+            value=5,
             step=1,
             label="Number of Few-shot Examples",
-            info="Number of examples to use for few-shot learning (0-5). They will be loaded in alphabetical order."
+            info="Number of examples to use for few-shot learning (0-5)."
         )
     
     with gr.Row():
         all_questions_checkbox = gr.Checkbox(
             label="Evaluate All Questions",
-            value=False,  # Default is unchecked
+            value=False,
             info="When checked, evaluates all available questions for each subject"
         )
         questions_info_text = gr.Markdown(visible=False, value="**All 12,032 questions across all subjects will be evaluated**")
@@ -151,33 +153,32 @@ with gr.Blocks() as demo:
     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=40,
-            value=20,  # Default is 10 questions
+            value=20,
             step=1,
             label="Questions per Subject",
-            info="Choose a subset of questions (1-40) per subject. They will be loaded in order of question_id for reproducibility. ",
+            info="Choose a subset of questions (1-40) per subject. They will be loaded in order of question_id.",
             interactive=True
         )
     
     with gr.Row():
         with gr.Column(scale=1):
             eval_mmlu_button = gr.Button("Run MMLU-Pro Evaluation", variant="primary", interactive=True)
-            cancel_mmlu_button = gr.Button("Cancel MMLU-Pro Evaluation", variant="stop", visible=False)
+            cancel_mmlu_button = gr.Button("Cancel Evaluation", variant="stop", visible=False)
         results_output = gr.Markdown(label="Evaluation Results")
         
     with gr.Row():
         results_table = gr.DataFrame(interactive=True, label="Detailed Results (Sortable)", visible=True)
-        
+    
+    # Track evaluation state - used to prevent multiple evaluations
+    evaluation_state = gr.State({"running": False})
+    
     # Update num_subjects_slider interactivity based on all_subjects checkbox
     def update_subjects_slider(checked):
-        if checked:
-            return gr.update(value=14, interactive=False)
-        else:
-            return gr.update(interactive=True)
+        return gr.update(interactive=not checked)
     
     all_subjects_checkbox.change(
         fn=update_subjects_slider,
@@ -199,45 +200,75 @@ with gr.Blocks() as demo:
     )
     
     # Function to disable UI components during evaluation
-    def disable_ui_for_evaluation():
+    def start_evaluation(state):
+        if state["running"]:
+            return [
+                state,
+                gr.update(interactive=False),
+                gr.update(interactive=False),
+                gr.update(interactive=False),
+                gr.update(interactive=False),
+                gr.update(interactive=False),
+                gr.update(interactive=False),
+                gr.update(visible=False),
+                "Evaluation already in progress. Please wait.",
+                None
+            ]
+        
+        # Update state to running
+        state["running"] = True
+        
         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
+            state,
+            gr.update(interactive=False),  # all_subjects_checkbox
+            gr.update(interactive=False),  # num_subjects_slider
+            gr.update(interactive=False),  # num_shots_slider
+            gr.update(interactive=False),  # all_questions_checkbox
+            gr.update(interactive=False),  # num_questions_slider
             gr.update(interactive=False),  # eval_mmlu_button
-            gr.update(visible=True)   # cancel_mmlu_button
+            gr.update(visible=True),       # cancel_mmlu_button
+            "Starting evaluation...",      # results_output
+            None                           # results_table
         ]
     
+    # Function to reset UI after evaluation
+    def finish_evaluation(state):
+        state["running"] = False
+        return state
+    
     # 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
+    def cancel_evaluation(state):
+        # Note: This doesn't actually stop the evaluation process
+        # It only updates the UI state to appear canceled
+        state["running"] = False
         return [
-            gr.update(interactive=True, info="When checked, evaluates all 14 MMLU-Pro subjects"),  # all_subjects_checkbox
-            gr.update(interactive=True, info="Number of subjects to evaluate (1-14). 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-40) per subject. They will be loaded in order of question_id for reproducibility."),  # num_questions_slider
+            state,
+            gr.update(interactive=True),  # all_subjects_checkbox
+            gr.update(interactive=True),  # num_subjects_slider
+            gr.update(interactive=True),  # num_shots_slider
+            gr.update(interactive=True),  # all_questions_checkbox
+            gr.update(interactive=True),  # 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
+            gr.update(visible=False),     # cancel_mmlu_button
+            "⚠️ Evaluation canceled by user (note: backend process may continue running)", # results_output
+            None                          # results_table
         ]
     
-    # Connect MMLU evaluation button - now disables UI and shows cancel button
+    # Connect MMLU evaluation button with state tracking
     eval_mmlu_button.click(
-        fn=disable_ui_for_evaluation,
-        inputs=None,
+        fn=start_evaluation,
+        inputs=[evaluation_state],
         outputs=[
+            evaluation_state,
             all_subjects_checkbox,
             num_subjects_slider,
             num_shots_slider,
             all_questions_checkbox,
             num_questions_slider,
             eval_mmlu_button,
-            cancel_mmlu_button
+            cancel_mmlu_button,
+            results_output,
+            results_table
         ]
     ).then(
         fn=run_mmlu_evaluation,
@@ -259,13 +290,18 @@ with gr.Blocks() as demo:
             all_questions_checkbox,
             num_questions_slider
         ]
+    ).then(
+        fn=finish_evaluation,
+        inputs=[evaluation_state],
+        outputs=[evaluation_state]
     )
     
     # Connect cancel button
     cancel_mmlu_button.click(
         fn=cancel_evaluation,
-        inputs=None,
+        inputs=[evaluation_state],
         outputs=[
+            evaluation_state,
             all_subjects_checkbox,
             num_subjects_slider,
             num_shots_slider,