Update mmlu_eval.py

mmlu_eval.py

@@ -35,30 +35,39 @@ def evaluate_mmlu(model, tokenizer, num_questions_per_task=5):
         - Overall accuracy
         - Min accuracy task
         - Max accuracy task
+        - Two correct examples
+        - Two incorrect examples
     """
     results = {}
-    
+    correct_examples = []
+    incorrect_examples = []
+
     for task_name in mmlu_dataset.keys():
         dataset = mmlu_dataset[task_name]
         sampled_questions = random.sample(list(dataset), min(num_questions_per_task, len(dataset)))
-        
+
         predictions = []
         references = []
-        
+
         for sample in sampled_questions:
             question = sample["question"]
-            correct_answer = sample["answer"]  # Assuming dataset provides direct answers
-
-            model_output = generate_answer(model, tokenizer, question)
+            correct_answer = str(sample["answer"]).strip().lower()
+            model_output = generate_answer(model, tokenizer, question).strip().lower()
 
             predictions.append(model_output)
             references.append(correct_answer)
 
+            # Store examples
+            if model_output == correct_answer and len(correct_examples) < 2:
+                correct_examples.append((task_name, question, model_output, correct_answer))
+            elif model_output != correct_answer and len(incorrect_examples) < 2:
+                incorrect_examples.append((task_name, question, model_output, correct_answer))
+
         # Compute accuracy for the task
         norm_preds = [str(p).lower().strip() for p in predictions]
         norm_refs = [str(r).lower().strip() for r in references]
         task_accuracy = accuracy_metric.compute(predictions=norm_preds, references=norm_refs)["accuracy"]
-        
+
         results[task_name] = task_accuracy
 
     # Compute overall statistics
@@ -70,4 +79,6 @@ def evaluate_mmlu(model, tokenizer, num_questions_per_task=5):
         "overall_accuracy": overall_accuracy,
         "min_accuracy_task": (min_task, results[min_task]),
         "max_accuracy_task": (max_task, results[max_task]),
+        "correct_examples": correct_examples,
+        "incorrect_examples": incorrect_examples,
     }