Update mmlu_eval_original.py

mmlu_eval_original.py

@@ -4,12 +4,16 @@ from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import spaces
 import logging
+import numpy as np
+import pandas as pd
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 accuracy_metric = evaluate.load("accuracy")
+choices = ["A", "B", "C", "D"]
+MAX_CONTEXT_WINDOW = 4096 #Hard-coded for the moment, will be replaced later to be an input from the Model.
 
 def load_dataset_from_hf(verbose=False):
     mmlu_dataset = load_dataset("cais/mmlu", "all")
@@ -38,106 +42,140 @@ def load_dataset_from_hf(verbose=False):
         
             logger.info("=" * 50)  # Separator for readability
     return mmlu_dataset
-        
 
-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}
+def format_subject(subject):
+    l = subject.split("_")
+    s = ""
+    for entry in l:
+        s += " " + entry
+    return s
 
-Choices:
-{formatted_choices} [/INST]"""
+
+def format_example(df, idx, include_answer=True):
+    prompt = df.iloc[idx, 0]
+    k = df.shape[1] - 2
+    for j in range(k):
+        prompt += "\n{}. {}".format(choices[j], df.iloc[idx, j + 1])
+    prompt += "\nAnswer:"
+    if include_answer:
+        prompt += " {}\n\n".format(df.iloc[idx, k + 1])
+    return prompt
+
+
+def gen_prompt(df, subject, k=-1):
+    prompt = "The following are multiple choice questions (with answers) about {}.\n\n".format(
+        format_subject(subject)
+    )
+    if k == -1:
+        k = df.shape[0]
+    for i in range(k):
+        prompt += format_example(df, i)
     return prompt
 
-@spaces.GPU
-def generate_answer(model, tokenizer, question, choices):
-    """
-    Generates an answer using Mistral's instruction format for multiple choice questions.
-    """
-    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=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
-        )
-    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
 
 @torch.no_grad()
-def evaluate_mmlu(model, tokenizer, num_questions=5):
+def eval (subject, model, tokenizer, dev_df, test_df, num_questions_per_subject=5, train_shots=5):
+    cors = []
+    all_probs = []
+
+    if (train_shots < 0):
+        train_shots = 0 # Make positive.
+
+    for i in range(test_df.shape[0]):
+        prompt_end = format_example(test_df, i, include_answer=False)
+        train_prompt = gen_prompt(dev_df, subject, train_shots)
+        prompt = train_prompt + prompt_end
+
+        input_ids = tokenizer (prompt, return_tensors="pt").input_ids.to(model.device)
+
+
+        # Reduce number of shots in the prompt to fit in context window.
+        while (train_shots > 0 and input_ids.shape[-1] > MAX_CONTEXT_WINDOW):
+            train_shots -= 1
+            train_prompt = gen_prompt(dev_df, subject, train_shots)
+            prompt = train_prompt + prompt_end
+            input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(
+                model.device
+            )
+        logger.info (f"Prompt: {prompt}")
+
+        label = test_df.iloc[i, test_df.shape[1] - 1]
+
+        logits = model(input_ids=input_ids).logits[0, -1]
+
+
+        probs = (
+            torch.nn.functional.softmax(
+                torch.tensor(
+                    [
+                        logits[tokenizer("A").input_ids[-1]],
+                        logits[tokenizer("B").input_ids[-1]],
+                        logits[tokenizer("C").input_ids[-1]],
+                        logits[tokenizer("D").input_ids[-1]],
+                    ]
+                ).float(),
+                dim=0,
+            )
+            .detach()
+            .cpu()
+            .numpy()
+        )
+        pred = {0: "A", 1: "B", 2: "C", 3: "D"}[np.argmax(probs)]
+
+        cor = pred == label
+
+        cors.append(cor)
+        all_probs.append(probs)
+
+    acc = np.mean(cors)
+    cors = np.array(cors)
+
+    all_probs = np.array(all_probs)
+    print("Average accuracy {:.3f} - {}".format(acc, subject))
+
+    return cors, acc, all_probs
+
+        
+        
+
+def evaluate_mmlu(model, tokenizer, num_subjects=-1, num_questions=5, num_shots=5):
     """
-    Evaluates the model on MMLU across all tasks.
+    Evaluates the model on MMLU across all subjects.
     """
-    mmlu_dataset = load_dataset_from_hf(verbose=True)
+    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'])
+
+    subjects = sorted(test_df['subject'].unique())
+
     results = {}
     correct_examples = []
     incorrect_examples = []
+    all_accuracies = []
+    all_cors = []
     
-    # Filter out 'auxiliary_train' and other non-test splits
-    test_tasks = [k for k in mmlu_dataset.keys() if 'test' in k]
+    for subject in subjects:
+        test_samples = test_df[test_df['subject'] == subject].head(num_questions)
+        dev_samples = dev_df[dev_df['subject'] == subject].head(num_shots)
+
+        # Log subject and sample counts
+        logger.info(f"Subject: {subject}, Test Samples: {len(test_samples)}, Dev Samples: {len(dev_samples)}")
+
+        cors, acc, probs = eval(subject, model, tokenizer, dev_samples, test_samples, num_questions_per_subject=num_questions, train_shots=num_shots)
+        all_cors.append(cors)
     
-    for task_name in sorted(test_tasks):  # Sort tasks for deterministic order
-        dataset = mmlu_dataset[task_name]
-        # 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 ("TASK", task_name, "Sample", sample)
-            question = sample["question"]
-            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"])
-            print ("question:", question, "\n choices:", choices, "\n correct answer:", correct_answer)
-            
-            model_output = generate_answer(model, tokenizer, question, choices)
-            print ("model output:", model_output)
-            
-            predictions.append(model_output)
-            references.append(correct_answer)
-
-            # Store examples
-            if model_output == correct_answer and len(correct_examples) < 2:
-                correct_examples.append((task_name, question, model_output, correct_answer))
-            elif model_output != correct_answer and len(incorrect_examples) < 2:
-                incorrect_examples.append((task_name, question, model_output, correct_answer))
-
-        # Compute accuracy for the task
-        task_accuracy = accuracy_metric.compute(
-            predictions=predictions,
-            references=references
-        )["accuracy"]
-        
-        results[task_name] = task_accuracy
+    weighted_acc = np.mean(np.concatenate(all_cors))
 
-    # Compute overall statistics
-    overall_accuracy = sum(results.values()) / len(results)
-    min_task = min(results, key=results.get)
-    max_task = max(results, key=results.get)
 
     return {
-        "overall_accuracy": overall_accuracy,
-        "min_accuracy_task": (min_task, results[min_task]),
-        "max_accuracy_task": (max_task, results[max_task]),
+        "overall_accuracy": weighted_acc,
+        "min_accuracy_subject": (min_acc_subject, results[min_acc_subject]),
+        "max_accuracy_subject": (max_acc_subject, results[max_acc_subject]),
         "correct_examples": correct_examples,
         "incorrect_examples": incorrect_examples,
-        "all_results": results  # Added for detailed analysis
     }