Update mmlu_eval.py

mmlu_eval.py

@@ -1,63 +1,78 @@
 import torch
-import random
 import evaluate
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import spaces
 
-# Load Accuracy Metric
 accuracy_metric = evaluate.load("accuracy")
-
-# Load MMLU dataset
 mmlu_dataset = load_dataset("cais/mmlu", "all")
 
+def format_mmlu_prompt(question, choices):
+    """
+    Formats the prompt according to Mistral's official instruction format.
+    Source: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3
+    """
+    formatted_choices = "\n".join([f"{chr(65 + i)}. {choice}" for i, choice in enumerate(choices)])
+    prompt = f"""<s>[INST] You are taking a multiple choice test. Select the correct answer by responding with only the letter (A, B, C, or D) of the correct choice.
+
+Question: {question}
+
+Choices:
+{formatted_choices} [/INST]"""
+    return prompt
+
 @spaces.GPU
-def generate_answer(model, tokenizer, question):
+def generate_answer(model, tokenizer, question, choices):
     """
-    Generates an answer using Mistral's instruction format.
+    Generates an answer using Mistral's instruction format for multiple choice questions.
     """
-    prompt = f"<s>[INST] {question}. Provide only the correct answer. [/INST]"
+    prompt = format_mmlu_prompt(question, choices)
     
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=50,
-            temperature=0.0,
+            max_new_tokens=5,  # We only need a single letter
+            do_sample=False,   # Use deterministic greedy decoding
+            num_beams=1,       # Use simple greedy search
             pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id
         )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+    # Extract just the letter answer
+    for char in response:
+        if char in 'ABCD':
+            return char
+    return response[:1]  # Fallback: take first character
 
 def evaluate_mmlu(model, tokenizer, num_questions_per_task=5):
     """
-    Evaluates the model on MMLU across all 57 tasks.
-
-    Returns:
-        - Overall accuracy
-        - Min accuracy task
-        - Max accuracy task
-        - Two correct examples
-        - Two incorrect examples
+    Evaluates the model on MMLU across all tasks.
     """
     results = {}
     correct_examples = []
     incorrect_examples = []
-
-    for task_name in mmlu_dataset.keys():
-        print ("TASK NAME: ", task_name)
+    
+    # Filter out 'auxiliary_train' and other non-test splits
+    test_tasks = [k for k in mmlu_dataset.keys() if 'test' in k]
+    
+    for task_name in sorted(test_tasks):  # Sort tasks for deterministic order
         dataset = mmlu_dataset[task_name]
-        sampled_questions = random.sample(list(dataset), min(num_questions_per_task, len(dataset)))
+        # Instead of random sampling, take the first n questions
+        total_questions = min(num_questions_per_task, len(dataset))
+        sampled_questions = [dataset[i] for i in range(total_questions)]
 
         predictions = []
         references = []
 
         for sample in sampled_questions:
-            print ("SAMPLE", sample)
             question = sample["question"]
-            correct_answer = str(sample["answer"]).strip().lower()
-            model_output = generate_answer(model, tokenizer, question).strip().lower()
-
+            choices = [sample["choices"][i] for i in range(4)]
+            # Convert numeric answer to letter (0->A, 1->B, etc.)
+            correct_answer = chr(65 + sample["answer"])
+            
+            model_output = generate_answer(model, tokenizer, question, choices)
+            
             predictions.append(model_output)
             references.append(correct_answer)
 
@@ -68,10 +83,11 @@ def evaluate_mmlu(model, tokenizer, num_questions_per_task=5):
                 incorrect_examples.append((task_name, question, model_output, correct_answer))
 
         # Compute accuracy for the task
-        norm_preds = [str(p).lower().strip() for p in predictions]
-        norm_refs = [str(r).lower().strip() for r in references]
-        task_accuracy = accuracy_metric.compute(predictions=norm_preds, references=norm_refs)["accuracy"]
-
+        task_accuracy = accuracy_metric.compute(
+            predictions=predictions,
+            references=references
+        )["accuracy"]
+        
         results[task_name] = task_accuracy
 
     # Compute overall statistics
@@ -85,4 +101,5 @@ def evaluate_mmlu(model, tokenizer, num_questions_per_task=5):
         "max_accuracy_task": (max_task, results[max_task]),
         "correct_examples": correct_examples,
         "incorrect_examples": incorrect_examples,
-    }
+        "all_results": results  # Added for detailed analysis
+    }