Update mmlu_pro_eval_adapted.py

mmlu_pro_eval_adapted.py

@@ -149,7 +149,7 @@ def extract_final(text):
 def batch_inference(llm, sampling_params, inference_batch):
     start = time.time()
     outputs = llm.generate(inference_batch, sampling_params)
-    logging.info(str(len(inference_batch)) + " size batch costing time: " + str(time.time() - start))
+    logging.info("Batch of size: ", str(len(inference_batch)) + ". Time taken: " + str(time.time() - start))
     response_batch = []
     pred_batch = []
     for output in outputs:
@@ -157,9 +157,72 @@ def batch_inference(llm, sampling_params, inference_batch):
         response_batch.append(generated_text)
         pred = extract_answer(generated_text)
         pred_batch.append(pred)
-    logging.info("PRED BATCH: %s, RESPONSE BATCH: %s", pred_batch, response_batch)
     return pred_batch, response_batch
 
+def batch_inference_debug_mode(llm, sampling_params, inference_batch, tokenizer):
+    start = time.time()
+    outputs = llm.generate(inference_batch, sampling_params)
+    logging.info(str(len(inference_batch)) + " size batch costing time: " + str(time.time() - start))
+    response_batch = []
+    pred_batch = []
+    input_token_counts = []
+    output_token_counts = []
+
+    for i, output in enumerate(outputs):
+        generated_text = output.outputs[0].text
+        response_batch.append(generated_text)
+        pred = extract_answer(generated_text)
+        pred_batch.append(pred)
+        
+        # Proper token count using tokenizer
+        input_tokens = len(tokenizer.encode(inference_batch[i]))
+        output_tokens = len(tokenizer.encode(generated_text))
+        
+        input_token_counts.append(input_tokens)
+        output_token_counts.append(output_tokens)
+    
+    logging.info("PRED BATCH: %s", pred_batch)
+    logging.info("RESPONSE BATCH: %s", response_batch)
+    
+    # Convert to DataFrame for logging (handle cases with fewer than 40 requests)
+    num_samples = min(40, len(inference_batch))
+    summary_df = pd.DataFrame({
+        'Input': inference_batch[:num_samples],
+        'Response': response_batch[:num_samples]
+    })
+    logging.info("\nSummary of first %d requests and responses:\n%s", num_samples, summary_df.to_string())
+    
+    # Total and average input/output token statistics
+    total_input_tokens = sum(input_token_counts)
+    total_output_tokens = sum(output_token_counts)
+    avg_input_tokens = total_input_tokens / len(input_token_counts)
+    avg_output_tokens = total_output_tokens / len(output_token_counts)
+    
+    max_input_idx = np.argmax(input_token_counts)
+    max_output_idx = np.argmax(output_token_counts)
+    min_input_idx = np.argmin(input_token_counts)
+    min_output_idx = np.argmin(output_token_counts)
+    
+    logging.info("\nTotal input tokens: %d", total_input_tokens)
+    logging.info("Total output tokens: %d", total_output_tokens)
+    logging.info("Average input tokens: %.2f", avg_input_tokens)
+    logging.info("Average output tokens: %.2f", avg_output_tokens)
+    
+    logging.info("\nRequest with max input tokens: %d (Tokens: %d)\nInput: %s\nOutput: %s", 
+                 max_input_idx, input_token_counts[max_input_idx], inference_batch[max_input_idx], response_batch[max_input_idx])
+    
+    logging.info("\nRequest with max output tokens: %d (Tokens: %d)\nInput: %s\nOutput: %s", 
+                 max_output_idx, output_token_counts[max_output_idx], inference_batch[max_output_idx], response_batch[max_output_idx])
+    
+    logging.info("\nRequest with min input tokens: %d (Tokens: %d)\nInput: %s\nOutput: %s", 
+                 min_input_idx, input_token_counts[min_input_idx], inference_batch[min_input_idx], response_batch[min_input_idx])
+    
+    logging.info("\nRequest with min output tokens: %d (Tokens: %d)\nInput: %s\nOutput: %s", 
+                 min_output_idx, output_token_counts[min_output_idx], inference_batch[min_output_idx], response_batch[min_output_idx])
+
+    return pred_batch, response_batch
+
+
 
 def calculate_accuracy(res):
     """
@@ -190,7 +253,7 @@ def calculate_accuracy(res):
 
 
 @torch.no_grad()
-def eval_cot(subject, model, tokenizer, val_df, test_df, num_shots=5):
+def eval_cot(subject, model, tokenizer, val_df, test_df, num_shots=5, debug_mode=True):
     """
     Evaluate model using chain-of-thought prompting.
     
@@ -231,8 +294,9 @@ def eval_cot(subject, model, tokenizer, val_df, test_df, num_shots=5):
             
         inference_batches.append(prompt)
 
-    # Get model predictions
-    pred_batch, response_batch = batch_inference(llm, sampling_params, inference_batches)
+
+    batch_fn = batch_inference_debug_mode if debug_mode else batch_inference
+    pred_batch, response_batch = batch_fn(llm, sampling_params, inference_batches)
     
     # Add predictions to test DataFrame
     results_df = test_df.copy()
@@ -247,7 +311,7 @@ def eval_cot(subject, model, tokenizer, val_df, test_df, num_shots=5):
     return correctness, accuracy
 
 
-@spaces.GPU(duration=240)  # Extended to 3 minutes for larger evaluations
+@spaces.GPU(duration=240)  # Extended to 4 minutes for larger evaluations
 def evaluate_mmlu_pro(model_name, num_subjects=-1, num_questions=10, num_shots=5):
     """
     Main evaluation function for MMLU-Pro benchmark.