Update mmlu_pro_eval_adapted.py

mmlu_pro_eval_adapted.py

@@ -310,17 +310,18 @@ def eval_cot(subject, model, tokenizer, val_df, test_df, num_shots=5, debug_mode
 
     return correctness, accuracy
 
-
-@spaces.GPU(duration=240)  # Extended to 4 minutes for larger evaluations
-def evaluate_mmlu_pro(model_name, num_subjects=-1, num_questions=10, num_shots=5):
+def evaluate_mmlu_pro(model_name, num_subjects=-1, num_questions=10, num_shots=5, specific_subjects=None, flash_attention=False, regex_pattern=None):
     """
     Main evaluation function for MMLU-Pro benchmark.
     
     Args:
         model_name: Name/path of model to evaluate
         num_subjects: Number of subjects to test (-1 for all)
-        num_questions: Number of questions per subject
+        num_questions: Number of questions per subject (-1 for all)
         num_shots: Number of examples to include in prompts
+        specific_subjects: List of specific subjects to evaluate (overrides num_subjects)
+        flash_attention: Whether to use flash attention (currently ignored)
+        regex_pattern: Regex pattern for answer extraction (currently ignored)
     """
     print(f"Is CUDA available: {torch.cuda.is_available()}")
     print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
@@ -332,12 +333,13 @@ def evaluate_mmlu_pro(model_name, num_subjects=-1, num_questions=10, num_shots=5
     # Sort DataFrames
     test_df = test_df.sort_values(['category', 'question_id'])
     val_df = val_df.sort_values(['category', 'question_id'])
-
     # Get unique subjects
     all_subjects = sorted(test_df['category'].unique())
     
-    # Select subjects based on num_subjects parameter
-    if num_subjects == -1 or num_subjects >= len(all_subjects):
+    # Select subjects based on parameters
+    if specific_subjects is not None:
+        selected_subjects = [subject for subject in specific_subjects if subject in all_subjects]
+    elif num_subjects == -1 or num_subjects >= len(all_subjects):
         selected_subjects = all_subjects
     else:
         selected_subjects = all_subjects[:num_subjects]
@@ -348,13 +350,17 @@ def evaluate_mmlu_pro(model_name, num_subjects=-1, num_questions=10, num_shots=5
     results = {}
     all_correctness = []
     results_table = []
-
     # Process each subject
     for subject in tqdm(selected_subjects, desc="Processing Selected Categories"):
         # Filter data for current subject
-        test_samples = test_df[test_df['category'] == subject].head(num_questions)
+        if num_questions == -1:
+            # Use all questions for this subject
+            test_samples = test_df[test_df['category'] == subject]
+        else:
+            # Use specified number of questions
+            test_samples = test_df[test_df['category'] == subject].head(num_questions)
+            
         val_samples = val_df[val_df['category'] == subject].head(num_shots)
-
         # Run evaluation
         correctness, acc = eval_cot(
             subject, 
@@ -377,13 +383,13 @@ def evaluate_mmlu_pro(model_name, num_subjects=-1, num_questions=10, num_shots=5
     
     # Calculate overall metrics
     weighted_acc = np.mean(all_correctness)
-    min_acc_subject = min(results.items(), key=lambda x: x[1])[0]
-    max_acc_subject = max(results.items(), key=lambda x: x[1])[0]
+    min_acc_subject = min(results.items(), key=lambda x: x[1])
+    max_acc_subject = max(results.items(), key=lambda x: x[1])
     
     # Return results summary
     return {
         "overall_accuracy": weighted_acc,
-        "min_accuracy_subject": (min_acc_subject, results[min_acc_subject]),
-        "max_accuracy_subject": (max_acc_subject, results[max_acc_subject]),
+        "min_accuracy_subject": min_acc_subject,
+        "max_accuracy_subject": max_acc_subject,
         "full_accuracy_table": results_table,
     }