Update dataset_previews.py

dataset_previews.py

@@ -4,7 +4,7 @@ import pandas as pd
 import numpy as np
 from typing import Dict, Any, List, Tuple
 import collections
-from mmlu_pro_eval_adapted import load_mmlu_pro, preprocess  # Import preprocess also
+from mmlu_pro_eval_adapted import load_mmlu_pro
 
 def calculate_dataset_statistics():
     """
@@ -15,8 +15,12 @@ def calculate_dataset_statistics():
     """
     try:
         # Load MMLU-Pro data using the function from mmlu_pro_eval_adapted
-        mmlu_data = load_mmlu_pro(num_subjects=-1, num_questions=-1)
+        test_df, val_df = load_mmlu_pro()
         
+        test_df = test_df.sort_values(['category', 'question_id'])
+        
+        all_subjects = sorted(test_df['category'].unique())
+
         # Calculate total questions and questions per subject
         total_questions = 0
         subject_counts = {}
@@ -24,14 +28,15 @@ def calculate_dataset_statistics():
         # Count options per question
         options_counts = []
         
-        for subject_name, subject_data in mmlu_data.items():
-            num_questions = len(subject_data["test_examples"])
-            subject_counts[subject_name] = num_questions
+        for subject in all_subjects:
+            test_samples = test_df[test_df['category'] == subject]
+            num_questions = len(test_samples)
+            subject_counts[subject] = num_questions
             total_questions += num_questions
             
             # Count options for each question
-            for test_example in subject_data["test_examples"]:
-                options_counts.append(len(test_example["options"]))
+            for sample in test_samples:
+                options_counts.append(len(sample["options"]))
         
         max_options = max(options_counts)
         avg_options = sum(options_counts) / len(options_counts)