Update mmlu_eval_original.py

mmlu_eval_original.py

@@ -3,9 +3,42 @@ import evaluate
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import spaces
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 accuracy_metric = evaluate.load("accuracy")
-mmlu_dataset = load_dataset("cais/mmlu", "all")
+
+def load_dataset_from_hf(verbose=False):
+    mmlu_dataset = load_dataset("cais/mmlu", "all")
+    if verbose:
+        for split in mmlu_dataset.keys():
+            dataset = mmlu_dataset[split]  # Access the dataset split
+            
+            # Log number of rows and columns
+            num_rows = len(dataset)
+            num_cols = len(dataset.column_names)
+            
+            logger.info(f"Dataset Split: {split}")
+            logger.info(f"Number of Rows: {num_rows}")
+            logger.info(f"Number of Columns: {num_cols}")
+            
+            # Log column names and their types
+            column_types = {col: str(dataset.features[col].dtype) for col in dataset.column_names}
+            logger.info(f"Column Names: {dataset.column_names}")
+            logger.info(f"Column Types: {column_types}")
+        
+            # Log a sample of 5 rows
+            sample_rows = dataset.select(range(min(5, num_rows)))  # Ensure we don't exceed available rows
+            logger.info("Sample Rows:")
+            for row in sample_rows:
+                logger.info(row)
+        
+            logger.info("=" * 50)  # Separator for readability
+    return mmlu_dataset
+        
 
 def format_mmlu_prompt(question, choices):
     """
@@ -46,10 +79,11 @@ def generate_answer(model, tokenizer, question, choices):
     return response[:1]  # Fallback: take first character
 
 @torch.no_grad()
-def eval(model, tokenizer, num_questions_per_task=5, dev_df, test_df):
+def evaluate_mmlu(model, tokenizer, num_questions=5):
     """
     Evaluates the model on MMLU across all tasks.
     """
+    mmlu_dataset = load_dataset_from_hf(verbose=True)
     results = {}
     correct_examples = []
     incorrect_examples = []