Update app.py

app.py

@@ -60,71 +60,141 @@ def run_toy_evaluation():
 # 3. MMLU Evaluation call
 # ---------------------------------------------------------------------------
 @spaces.GPU(duration=120)  # Allow up to 2 minutes for full evaluation
-def run_mmlu_evaluation(num_questions):
-    if not model_loaded:
-        load_model()
-
-    if not model_loaded:
-        return "⚠️ Model not loaded. Please load the model first."
+def run_mmlu_evaluation(all_subjects, num_subjects, num_shots, num_examples):
     """
-    Runs the MMLU evaluation with the specified number of questions per task.
-    Also displays two correct and two incorrect examples.
+    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_shots (int): Number of few-shot examples (0-5)
+        num_examples (int): Number of examples per subject (1-10 or -1 for all)
     """
-    results = evaluate_mmlu(model, tokenizer, num_questions)
-
-    overall_accuracy = results["overall_accuracy"]
-    min_task, min_acc = results["min_accuracy_task"]
-    max_task, max_acc = results["max_accuracy_task"]
-    correct_examples = results["correct_examples"]
-    incorrect_examples = results["incorrect_examples"]
-
-    # Format examples for readability
-    def format_example(example):
-        task, question, model_output, correct_answer = example
-        return f"**Task:** {task}\n**Question:** {question}\n**Model Output:** {model_output}\n**Correct Answer:** {correct_answer}\n"
+    if not model_loaded:
+        return "⚠️ Model not loaded. Please load the model first."
 
-    correct_text = "\n\n".join(format_example(ex) for ex in correct_examples)
-    incorrect_text = "\n\n".join(format_example(ex) for ex in incorrect_examples)
+    # Convert num_subjects to -1 if all_subjects is True
+    if all_subjects:
+        num_subjects = -1
+
+    # Run evaluation
+    results = evaluate_mmlu(
+        model, 
+        tokenizer,
+        num_subjects=num_subjects,
+        num_questions=num_examples,
+        num_shots=num_shots
+    )
 
+    # 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"])
+    
+    # Format the report
     report = (
-        f"### Overall Accuracy: {overall_accuracy:.2f}\n"
-        f"**Min Accuracy:** {min_acc:.2f} on `{min_task}`\n"
-        f"**Max Accuracy:** {max_acc:.2f} on `{max_task}`\n\n"
-        f"---\n\n"
-        f"### ✅ Correct Examples\n{correct_text if correct_examples else 'No correct examples available.'}\n\n"
-        f"### ❌ Incorrect Examples\n{incorrect_text if incorrect_examples else 'No incorrect examples available.'}"
+        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\n"
+        f"### Detailed Results Table\n"
+        f"{results_df.to_markdown()}\n"
     )
 
     return report
-
+    
 # ---------------------------------------------------------------------------
 # 4. Gradio Interface
 # ---------------------------------------------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# Mistral-7B Math Evaluation Demo")
+    gr.Markdown("# Mistral-7B on MMLU - Evaluation Demo")
     gr.Markdown("""
-    This demo evaluates Mistral-7B on Various Datasets.
+    This demo evaluates Mistral-7B on the MMLU Dataset.
     """)
 
-    # Load Model Button
-    load_button = gr.Button("Load Model", variant="primary")
-    load_status = gr.Textbox(label="Model Status", interactive=False)
-    load_button.click(fn=load_model, inputs=None, outputs=load_status)
+    # 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
+    # Toy Dataset Evaluation Section
     gr.Markdown("### Toy Dataset Evaluation")
-    eval_button = gr.Button("Run Evaluation", variant="primary")
-    output_text = gr.Textbox(label="Results")
-    output_plot = gr.HTML(label="Visualization and Details")
-
-    eval_button.click(fn=run_toy_evaluation, inputs=None, outputs=[output_text, output_plot])
-
-    # MMLU 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")
-    num_questions_input = gr.Number(label="Questions per Task (Total of 57 tasks)", value=5, precision=0)
-    eval_mmlu_button = gr.Button("Run MMLU Evaluation", variant="primary")
-    mmlu_output = gr.Textbox(label="MMLU Evaluation Results")
-
-    eval_mmlu_button.click(fn=run_mmlu_evaluation, inputs=[num_questions_input], outputs=[mmlu_output])
+    
+    with gr.Row():
+        all_subjects_checkbox = gr.Checkbox(
+            label="Evaluate All Subjects",
+            value=True,
+            info="When checked, evaluates all 57 MMLU subjects"
+        )
+        num_subjects_slider = gr.Slider(
+            minimum=1,
+            maximum=57,
+            value=57,
+            step=1,
+            label="Number of Subjects",
+            info="Number of subjects to evaluate (1-57)",
+            interactive=True
+        )
+    
+    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)"
+        )
+        num_examples_slider = gr.Slider(
+            minimum=1,
+            maximum=10,
+            value=5,
+            step=1,
+            label="Examples per Subject",
+            info="Number of test examples per subject (1-10)"
+        )
+    
+    with gr.Row():
+        eval_mmlu_button = gr.Button("Run MMLU Evaluation", variant="primary")
+        results_output = gr.Markdown(label="Evaluation Results")
+    
+    # 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
+    all_subjects_checkbox.change(
+        fn=lambda x: gr.update(interactive=not x),
+        inputs=[all_subjects_checkbox],
+        outputs=[num_subjects_slider]
+    )
+    
+    # Connect MMLU evaluation button
+    eval_mmlu_button.click(
+        fn=run_mmlu_evaluation,
+        inputs=[
+            all_subjects_checkbox,
+            num_subjects_slider,
+            num_shots_slider,
+            num_examples_slider
+        ],
+        outputs=results_output
+    )
 
-demo.launch()
+demo.launch()