Update dataset_previews.py

dataset_previews.py

@@ -17,33 +17,21 @@ def calculate_dataset_statistics():
         # Load MMLU-Pro data using the function from mmlu_pro_eval_adapted
         test_df, val_df = load_mmlu_pro()
         
+        # Ensure consistent ordering
         test_df = test_df.sort_values(['category', 'question_id'])
-        
-        all_subjects = sorted(test_df['category'].unique())
 
-        # Calculate total questions and questions per subject
+        # Calculate total questions
         total_questions = len(test_df)
-        subject_counts = {}
-        
-        # Count options per question
-        options_counts = []
-        
-        for subject in all_subjects:
-            print("Subject", subject)
-            test_samples = test_df[test_df['category'] == subject]
-            num_questions = len(test_samples)
-            subject_counts[subject] = num_questions
-            print("First sample", test_samples.head(1), "\t Num Questions:", num_questions)
-            
-            # Count options for each question
-            for sample in test_samples:
-                print ("SAMPLE", sample)
-                options_counts.append(len(sample['options']))
         
+        # Calculate subject-wise question counts efficiently
+        subject_counts = test_df['category'].value_counts().to_dict()
+
+        # Count options per question efficiently using `.apply()`
+        options_counts = test_df['options'].apply(len).tolist()
         max_options = max(options_counts)
         avg_options = sum(options_counts) / len(options_counts)
-        
-        # Count questions with each number of options
+
+        # Count frequency of each option count
         options_distribution = collections.Counter(options_counts)
         
         return {