Update mmlu_eval_original.py

mmlu_eval_original.py

@@ -76,6 +76,9 @@ def gen_prompt(df, subject, k=-1):
 
 @torch.no_grad()
 def eval (subject, model, tokenizer, dev_df, test_df, num_questions_per_subject=5, train_shots=5):
+    assert all(dev_df['subject'] == subject), f"Not all items in dev_df match subject {subject}"
+    assert all(test_df['subject'] == subject), f"Not all items in test_df match subject {subject}"
+    
     cors = []
     all_probs = []
 
@@ -125,6 +128,12 @@ def eval (subject, model, tokenizer, dev_df, test_df, num_questions_per_subject=
 
         cor = pred == label
 
+        logger.info(f"Label: {label}")
+        logger.info(f"Logits: {logits}")
+        logger.info(f"Probabilities: {probs}")
+        logger.info(f"Prediction: {pred}")
+        logger.info(f"Correct: {cor}")
+
         cors.append(cor)
         all_probs.append(probs)
 
@@ -151,6 +160,10 @@ def evaluate_mmlu(model, tokenizer, num_subjects=-1, num_questions=5, num_shots=
     test_df = pd.DataFrame(dataset['test'])
     dev_df = pd.DataFrame(dataset['dev'])
 
+    # Sort datasets by subject and other relevant columns
+    test_df = test_df.sort_values(['subject', 'question'])
+    dev_df = dev_df.sort_values(['subject', 'question'])
+
     subjects = sorted(test_df['subject'].unique())
 
     results = {}
@@ -158,6 +171,7 @@ def evaluate_mmlu(model, tokenizer, num_subjects=-1, num_questions=5, num_shots=
     incorrect_examples = []
     all_accuracies = []
     all_cors = []
+    results_table = []
     
     for subject in subjects:
         test_samples = test_df[test_df['subject'] == subject].head(num_questions)
@@ -167,15 +181,27 @@ def evaluate_mmlu(model, tokenizer, num_subjects=-1, num_questions=5, num_shots=
         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)
+        results[subject] = acc
         all_cors.append(cors)
-    
-    weighted_acc = np.mean(np.concatenate(all_cors))
+        
+        results_table.append({
+            'Subject': subject,
+            'Num_samples': len(test_samples),
+            '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]
+    max_acc_subject = max(results.items(), key=lambda x: x[1])[0]
+    
     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]),
-        "correct_examples": correct_examples,
-        "incorrect_examples": incorrect_examples,
+        "full_accuracy_table": results_table
     }