Update app.py

app.py

@@ -4,7 +4,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
 from huggingface_hub import login
 from toy_dataset_eval import evaluate_toy_dataset
-from mmlu_eval_original import evaluate_mmlu_batched
+from mmlu_pro_eval_adapted import evaluate_mmlu_pro
 import spaces
 import pandas as pd
 import time  # Added for timing functionality
@@ -23,49 +23,9 @@ model_name = "mistralai/Mistral-7B-v0.1"
 tokenizer = None
 model = None
 model_loaded = False
-
-@spaces.GPU
-def load_model():
-    """Loads the Mistral model and tokenizer and updates the load status."""
-    global tokenizer, model, model_loaded
-    start_time = time.time()  # Start timing
-    try:
-        if tokenizer is None:
-            tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
-        if model is None:
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                token=hf_token,
-                torch_dtype=torch.float16
-            )
-            model.to('cuda')
-        model_loaded = True
-        elapsed_time = time.time() - start_time  # Calculate elapsed time
-        return f"✅ Model Loaded in {elapsed_time:.2f} seconds!"
-    except Exception as e:
-        model_loaded = False
-        return f"❌ Model Load Failed: {str(e)}"
-# ---------------------------------------------------------------------------
-# 2. Toy Evaluation
-# ---------------------------------------------------------------------------
-@spaces.GPU(duration=120)
-def run_toy_evaluation():
-    """Runs the toy dataset evaluation."""
-    if not model_loaded:
-        load_model()
-
-    if not model_loaded:
-        return "⚠️ Model not loaded. Please load the model first."
-    
-    start_time = time.time()  # Start timing
-    results = evaluate_toy_dataset(model, tokenizer)
-    elapsed_time = time.time() - start_time  # Calculate elapsed time
-    
-    return f"{results}\n\nEvaluation completed in {elapsed_time:.2f} seconds.", \
-           f"<div>Time taken: {elapsed_time:.2f} seconds</div>"  # Return timing info
     
 # ---------------------------------------------------------------------------
-# 3. MMLU Evaluation call
+# 1. MMLU-Pro 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, progress=gr.Progress()):
@@ -80,17 +40,7 @@ def run_mmlu_evaluation(all_subjects, num_subjects, num_shots, all_questions, nu
         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, 
-                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))
-
+    
     # Convert num_subjects to -1 if all_subjects is True
     if all_subjects:
         num_subjects = -1
@@ -101,14 +51,11 @@ def run_mmlu_evaluation(all_subjects, num_subjects, num_shots, all_questions, nu
 
     # Run evaluation with timing
     start_time = time.time()  # Start timing
-    results = evaluate_mmlu_batched(
-        model, 
-        tokenizer,
+    results = evaluate_mmlu_pro(
+        model_name,
         num_subjects=num_subjects,
         num_questions=num_questions,
         num_shots=num_shots, 
-        batch_size=32,
-        auto_batch_size=True
     )
     elapsed_time = time.time() - start_time  # Calculate elapsed time
 
@@ -159,23 +106,11 @@ def run_mmlu_evaluation(all_subjects, num_subjects, num_shots, all_questions, nu
 # 4. Gradio Interface
 # ---------------------------------------------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# Mistral-7B on MMLU - Evaluation Demo")
+    gr.Markdown("# Mistral-7B on MMLU-Pro Evaluation Demo")
     gr.Markdown("""
-    This demo evaluates Mistral-7B on the MMLU Dataset.
+    This demo evaluates Mistral-7B-v0.1 on the MMLU-Pro Dataset (available here: https://huggingface.co/datasets/TIGER-Lab/MMLU-Pro).
     """)
 
-    # Load Model Section
-    with gr.Row():
-        load_button = gr.Button("Load Model", variant="primary")
-        load_status = gr.Textbox(label="Model Status", interactive=False)
-
-    # Toy Dataset Evaluation Section
-    gr.Markdown("### Toy Dataset Evaluation")
-    with gr.Row():
-        eval_toy_button = gr.Button("Run Toy Evaluation", variant="primary")
-        toy_output = gr.Textbox(label="Results")
-        toy_plot = gr.HTML(label="Visualization and Details")
-    
     # MMLU Evaluation Section
     gr.Markdown("### MMLU Evaluation")
     
@@ -183,15 +118,15 @@ with gr.Blocks() as demo:
         all_subjects_checkbox = gr.Checkbox(
             label="Evaluate All Subjects",
             value=False,  # Default is unchecked
-            info="When checked, evaluates all 57 MMLU subjects"
+            info="When checked, evaluates all 14 MMLU-Pro subjects"
         )
         num_subjects_slider = gr.Slider(
             minimum=1,
-            maximum=57,
-            value=10,  # Default is 10 subjects
+            maximum=14,
+            value=14,  # Default is all subjects
             step=1,
             label="Number of Subjects",
-            info="Number of subjects to evaluate (1-57). They will be loaded in alphabetical order.",
+            info="Number of subjects to evaluate (1-14). They will be loaded in alphabetical order.",
             interactive=True
         )
     
@@ -211,7 +146,7 @@ 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**")
+        questions_info_text = gr.Markdown(visible=False, value="**All 12,032 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")
@@ -220,37 +155,27 @@ with gr.Blocks() as demo:
     with questions_container:
         num_questions_slider = gr.Slider(
             minimum=1,
-            maximum=20,
-            value=10,  # Default is 10 questions
+            maximum=40,
+            value=20,  # Default is 10 questions
             step=1,
             label="Questions per Subject",
-            info="Choose a subset of questions (1-20)",
+            info="Choose a subset of questions (1-40) per subject. They will be loaded in order of question_id for reproducibility. ",
             interactive=True
         )
     
     with gr.Row():
         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)
+            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)
         results_output = gr.Markdown(label="Evaluation Results")
         
     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)
-    
-    # Connect toy evaluation
-    eval_toy_button.click(
-        fn=run_toy_evaluation,
-        inputs=None,
-        outputs=[toy_output, toy_plot]
-    )
-    
+        
     # Update num_subjects_slider interactivity based on all_subjects checkbox
     def update_subjects_slider(checked):
         if checked:
-            return gr.update(value=57, interactive=False)
+            return gr.update(value=14, interactive=False)
         else:
             return gr.update(interactive=True)
     
@@ -290,11 +215,11 @@ with gr.Blocks() as demo:
         # 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="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-20)"),  # num_questions_slider
+            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
             gr.update(interactive=True),  # eval_mmlu_button
             gr.update(visible=False),  # cancel_mmlu_button
             "⚠️ Evaluation canceled by user", # results_output