Update mmlu_eval_original.py

mmlu_eval_original.py

@@ -93,7 +93,7 @@ def gen_prompt(df, subject, k=-1):
 
 
 @torch.no_grad()
-def eval_batched(subject, model, tokenizer, dev_df, test_df, num_questions_per_subject=5, train_shots=5, batch_size=4):
+def eval_batched(subject, model, tokenizer, dev_df, test_df, num_questions_per_subject=5, train_shots=5, batch_size=8):
     """
     Improved eval function that uses batched processing on GPU 
     """
@@ -290,15 +290,26 @@ def evaluate_mmlu_batched(model, tokenizer, num_subjects=10, num_questions=10, n
         batch_size (int): Batch size for processing multiple examples at once
         auto_batch_size (bool): If True, automatically determine the optimal batch size
     """
-    
+    model.eval()  # Ensure Dropout and BatchNorm behave appropriately for inference
+
+    if tokenizer.pad_token is None:
+        logger.info("NO TOKENIZER PAD TOKEN")
+        tokenizer.pad_token = tokenizer.eos_token
+    if model.config.pad_token_id is None:
+        logger.info("NO PAD TOKEN ID")
+        model.config.pad_token_id = tokenizer.pad_token_id
+
+
+    dataset = load_dataset_from_hf(verbose=True)
+    test_df = pd.DataFrame(dataset['test'])
+    dev_df = pd.DataFrame(dataset['dev'])
+    test_df = test_df.sort_values(['subject', 'question'])
+    dev_df = dev_df.sort_values(['subject', 'question'])
+
+
     # If auto_batch_size is enabled, estimate the optimal batch size
     if auto_batch_size:
         # Get a sample prompt
-        dataset = load_dataset_from_hf(verbose=False)
-        test_df = pd.DataFrame(dataset['test'])
-        dev_df = pd.DataFrame(dataset['dev'])
-        test_df = test_df.sort_values(['subject', 'question'])
-        dev_df = dev_df.sort_values(['subject', 'question'])
         subject = test_df['subject'].iloc[0]
         test_sample = test_df[test_df['subject'] == subject].head(1)
         dev_sample = dev_df[dev_df['subject'] == subject].head(num_shots)
@@ -311,18 +322,7 @@ def evaluate_mmlu_batched(model, tokenizer, num_subjects=10, num_questions=10, n
         batch_size = get_max_batch_size(model, tokenizer, sample_prompt)
         logger.info(f"Auto-adjusted batch size: {batch_size}")
 
-    model.eval()  # Ensure Dropout and BatchNorm behave appropriately for inference
     
-    dataset = load_dataset_from_hf(verbose=True)
-
-    # Convert dataset partitions to pandas DataFrames
-    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'])
-
     # Get all unique subjects
     all_subjects = sorted(test_df['subject'].unique())