Update mmlu_eval_original.py

mmlu_eval_original.py

@@ -153,9 +153,16 @@ def eval (subject, model, tokenizer, dev_df, test_df, num_questions_per_subject=
         
 def evaluate_mmlu(model, tokenizer, num_subjects=-1, num_questions=5, num_shots=5):
     """
-    Evaluates the model on MMLU across all subjects.
+    Evaluates the model on MMLU across specified number of subjects.
+    
+    Args:
+        model: The model to evaluate
+        tokenizer: The tokenizer to use
+        num_subjects (int): Number of subjects to evaluate. If -1, evaluates all subjects
+        num_questions (int): Number of questions per subject
+        num_shots (int): Number of few-shot examples to use
     """
-    model.eval()  # Ensure Dropout and BatchNorm behave appropriately for inference.
+    model.eval()  # Ensure Dropout and BatchNorm behave appropriately for inference
     
     dataset = load_dataset_from_hf(verbose=True)
 
@@ -167,12 +174,17 @@ def evaluate_mmlu(model, tokenizer, num_subjects=-1, num_questions=5, num_shots=
     test_df = test_df.sort_values(['subject', 'question'])
     dev_df = dev_df.sort_values(['subject', 'question'])
 
-    subjects = sorted(test_df['subject'].unique())
+    # Get all unique subjects
+    all_subjects = sorted(test_df['subject'].unique())
+    
+    # Select subjects based on num_subjects parameter
+    if num_subjects == -1 or num_subjects >= len(all_subjects):
+        subjects = all_subjects
+    else:
+        # Take the first num_subjects subjects
+        subjects = all_subjects[:num_subjects]
 
     results = {}
-    correct_examples = []
-    incorrect_examples = []
-    all_accuracies = []
     all_cors = []
     results_table = []
     
@@ -183,7 +195,16 @@ def evaluate_mmlu(model, tokenizer, num_subjects=-1, num_questions=5, num_shots=
         # Log subject and sample counts
         logger.info(f"Subject: {subject}, Test Samples: {len(test_samples)}, Dev Samples: {len(dev_samples)}")
 
-        cors, acc, probs = eval(subject, model, tokenizer, dev_samples, test_samples, num_questions_per_subject=num_questions, train_shots=num_shots)
+        cors, acc, probs = eval(
+            subject, 
+            model, 
+            tokenizer, 
+            dev_samples, 
+            test_samples, 
+            num_questions_per_subject=num_questions, 
+            train_shots=num_shots
+        )
+        
         results[subject] = acc
         all_cors.append(cors)
         
@@ -193,7 +214,7 @@ def evaluate_mmlu(model, tokenizer, num_subjects=-1, num_questions=5, num_shots=
             'Num_correct': int(np.sum(cors)), 
             'Accuracy': acc
         })
-        
+    
     weighted_acc = np.mean(np.concatenate(all_cors))
     
     min_acc_subject = min(results.items(), key=lambda x: x[1])[0]
@@ -203,5 +224,5 @@ def evaluate_mmlu(model, tokenizer, num_subjects=-1, num_questions=5, num_shots=
         "overall_accuracy": weighted_acc,
         "min_accuracy_subject": (min_acc_subject, results[min_acc_subject]),
         "max_accuracy_subject": (max_acc_subject, results[max_acc_subject]),
-        "full_accuracy_table": results_table
-    }
+        "full_accuracy_table": results_table,
+    }