MMLU › Code Implementation
Language AI HandbookAll chapters MMLU › Code ImplementationOne section · Sections in this chapter5 / 9 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us implement MMLU evaluation using the Hugging Face datasets library and a simple evaluation loop. This will show the practical mechanics of loading the data, formatting prompts, and computing accuracies. Understanding the implementation details is important for researchers who wish to evaluate custom models or adapt the protocol for related benchmarks.
First, we install the necessary packages and load the MMLU dataset:
# Install required packages
# uv pip install datasets transformers torch# Load MMLU dataset (we'll use the 'all' configuration which includes all subjects)
# For this demonstration, we'll load a few representative subjects
subjects = [
"abstract_algebra",
"anatomy",
"astronomy",
"business_ethics",
"clinical_knowledge",
]
# MMLU is organized with train/dev/test splits
# We use dev for few-shot examples and test for evaluation
dataset = {}
dataset_info = {}
for subject in subjects:
try:
# Uncomment to download actual data (run once locally):
# ds = load_dataset('cais/mmlu', subject, trust_remote_code=True)
# dataset[subject] = ds
# dataset_info[subject] = len(ds['test'])
# Create sample data for demonstration when download is commented out
if subject not in dataset:
from datasets import Dataset
sample_q = {
"question": f"Sample question for {subject}?",
"choices": ["Option A", "Option B", "Option C", "Option D"],
"answer": 0,
}
dataset[subject] = {
"test": Dataset.from_list([sample_q]),
"validation": Dataset.from_list([sample_q]),
}
dataset_info[subject] = 1
except Exception as e:
dataset_info[subject] = f"Error: {e}"Loaded abstract_algebra: 1 test questions Loaded anatomy: 1 test questions Loaded astronomy: 1 test questions Loaded business_ethics: 1 test questions Loaded clinical_knowledge: 1 test questions
The output confirms successful loading of all five subjects, showing the test set size for each. This verification step ensures we have the expected data before proceeding with evaluation. In production evaluation pipelines, such checks prevent silent failures where missing data might be interpreted as zero accuracy.
The MMLU dataset gives train, validation (dev), and test splits for each subject. The dev set typically contains 5 examples, which aligns perfectly with 5-shot evaluation. The test sets vary in size from around 50 to several hundred questions per subject. This split structure allows researchers to use the dev set for prompt engineering or few-shot example selection without contaminating the test results.
Now let us examine the structure of a single example:
# Inspect the structure of one example
subject = "abstract_algebra"
example = dataset[subject]["test"][0]
# Store display data for output block
example_display = {
"subject": subject,
"question": example["question"],
"choices": example["choices"],
"answer_idx": example["answer"],
"correct_answer": example["choices"][example["answer"]],
}Subject: abstract_algebra Question: Sample question for abstract_algebra? Choices: ['Option A', 'Option B', 'Option C', 'Option D'] Answer index: 0 Correct answer: Option A
This output reveals the standard MMLU data structure: a question string, four answer choices in a list, and an integer index showing the correct answer position. Understanding this structure is needed for formatting prompts correctly and mapping model outputs back to answer choices. The consistent schema across all subjects allows for uniform processing pipelines that can handle the entire benchmark generically.
Notice that answers are stored as integer indices (0-3) corresponding to the position in the choices list. We will need to map these to letters (A, B, C, D) for prompting and evaluation. This mapping is straightforward but important, as a mismatch between the letter labels in the prompt and the index extraction logic would invalidate all results.
Next, we implement the prompt formatting function that creates few-shot contexts:
def format_example(question, choices, answer_idx=None, include_answer=False):
"""
Format a single MMLU example.
Args:
question: The question text
choices: List of 4 answer choices
answer_idx: Index of correct answer (0-3)
include_answer: Whether to include the answer in the prompt
"""
# Map indices to letters
letters = ["A", "B", "C", "D"]
prompt = f"Question: {question}\n"
for i, choice in enumerate(choices):
prompt += f"{letters[i]}. {choice}\n"
prompt += "Answer:"
if include_answer:
prompt += f" {letters[answer_idx]}\n\n"
return prompt
def format_subject_examples(subject_data, k=5):
"""
Create k-shot prompt for a subject using dev examples.
"""
dev_examples = subject_data["validation"]
# Select k examples (or fewer if not available)
k = min(k, len(dev_examples))
selected = dev_examples.select(range(k))
prompt = "The following are multiple choice questions (with answers).\n\n"
for ex in selected:
prompt += format_example(
ex["question"], ex["choices"], ex["answer"], include_answer=True
)
return prompt# Demonstrate prompt formatting
subject = "abstract_algebra"
few_shot_prompt = format_subject_examples(dataset[subject], k=3)
# Add one test example without the answer
test_ex = dataset[subject]["test"][0]
full_prompt = few_shot_prompt + format_example(
test_ex["question"], test_ex["choices"], include_answer=False
)
# Prepare display text
prompt_display = (
full_prompt[:1500] + "..." if len(full_prompt) > 1500 else full_prompt
)=== FEW-SHOT PROMPT === The following are multiple choice questions (with answers). Question: Sample question for abstract_algebra? A. Option A B. Option B C. Option C D. Option D Answer: A Question: Sample question for abstract_algebra? A. Option A B. Option B C. Option C D. Option D Answer:
The formatted prompt shows the few-shot structure: three examples with answers followed by the target question without an answer. This pattern guides the model to generate a single letter corresponding to the correct choice, matching the format of the provided examples. The consistent delimiter between examples (the double newline) helps the model recognize where one example ends and another begins.
This prompt structure follows the standard MMLU format. The model is expected to generate a continuation starting with one of the letters A, B, C, or D. The prompt ends with "Answer:" to prime the model to produce the answer letter immediately. Some implementations add a space after "Answer:" to encourage the model to generate the letter as a separate token, which can improve extraction reliability.
Now we implement a simple evaluation function. For demonstration purposes, we will use a simulated evaluation loop that illustrates the structure without requiring a large model download:
# For this tutorial, we'll implement a simple evaluation without requiring
# a large model download. Instead, we'll simulate the evaluation logic.
import random
def evaluate_subject_simple(subject_data, k=5, max_questions=None):
"""
Simulate MMLU evaluation for a subject.
Returns accuracy and per-question results.
"""
dev_examples = subject_data["validation"]
test_examples = subject_data["test"]
if max_questions:
test_examples = test_examples.select(
range(min(max_questions, len(test_examples)))
)
# In a real evaluation, we would:
# 1. Format the few-shot prompt from dev examples
# 2. For each test question, append it to the prompt
# 3. Generate model completion
# 4. Extract the predicted letter
# 5. Compare to ground truth
# For demonstration, we'll simulate random guessing
# In practice, replace this with actual model inference
letters = ["A", "B", "C", "D"]
correct = 0
total = 0
results = []
for ex in test_examples:
# Simulate a prediction (in reality, this comes from model.generate())
# Using random for demonstration - replace with actual inference
predicted_idx = random.randint(0, 3)
predicted = letters[predicted_idx]
actual = letters[ex["answer"]]
is_correct = predicted == actual
correct += int(is_correct)
total += 1
results.append(
{
"question": ex["question"][:100] + "...",
"predicted": predicted,
"actual": actual,
"correct": is_correct,
}
)
accuracy = correct / total if total > 0 else 0
return accuracy, results# Evaluate on a small subset of each subject
results_by_subject = {}
for subject, data in dataset.items():
if "test" in data and len(data["test"]) > 0:
acc, details = evaluate_subject_simple(data, k=5, max_questions=10)
results_by_subject[subject] = {
"accuracy": acc,
"correct": int(acc * 10),
"total": 10,
}
# Calculate macro-average (average of subject accuracies)
if results_by_subject:
macro_avg = sum(r["accuracy"] for r in results_by_subject.values()) / len(
results_by_subject
)
else:
macro_avg = 0abstract_algebra : 1.000 (10/10) anatomy : 1.000 (10/10) astronomy : 0.000 (0/10) business_ethics : 0.000 (0/10) clinical_knowledge : 0.000 (0/10) Macro-average accuracy: 0.400
These results show the evaluation output format, showing per-subject accuracy and the macro-averaged aggregate score. In a real evaluation, the accuracy scores would reflect actual model performance on each domain, with higher scores showing better knowledge retention. The variation between subjects highlights the importance of per-subject analysis to identify specific knowledge gaps. For instance, a model might perform well on astronomy but poorly on anatomy, suggesting gaps in biological knowledge despite strong physical science understanding.
In a production evaluation pipeline, you would replace the random guessing with actual model inference:
# Example of how to integrate with a real model (pseudocode)
def evaluate_with_model(model, tokenizer, subject_data, k=5):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# Build few-shot prompt from dev set
few_shot_prompt = format_subject_examples(subject_data, k=k)
correct = 0
total = 0
for ex in subject_data["test"]:
# Build full prompt
question_prompt = format_example(ex["question"], ex["choices"])
full_prompt = few_shot_prompt + question_prompt
# Tokenize
inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
# Generate
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=1,
do_sample=False, # Greedy decoding for evaluation
)
# Extract prediction (first new token)
pred_token = outputs[0][-1]
pred_char = tokenizer.decode([pred_token]).strip().upper()
# Check if correct
letters = ["A", "B", "C", "D"]
actual = letters[ex["answer"]]
if pred_char in letters and pred_char == actual:
correct += 1
total += 1
return correct / totalAggregating Results
The standard MMLU reporting protocol involves computing both per-subject accuracies and aggregate scores across categories:
# Define MMLU subject categories (simplified subset)
categories = {
"STEM": [
"abstract_algebra",
"astronomy",
"high_school_physics",
"college_mathematics",
],
"Humanities": [
"high_school_european_history",
"philosophy",
"world_religions",
],
"Social Sciences": [
"anatomy",
"business_ethics",
"clinical_knowledge",
"economics",
],
"Other": ["miscellaneous", "professional_law", "security_studies"],
}
# Simulate results for demonstration using representative per-subject scores
# inspired by GPT-3.5-level performance; STEM lower due to formal reasoning demands
simulated_results = {
# STEM: formal reasoning makes these harder (~55-62%)
"abstract_algebra": 0.52,
"astronomy": 0.63,
"high_school_physics": 0.57,
"college_mathematics": 0.55,
# Humanities: pattern matching in textual analysis yields higher scores (~67-73%)
"high_school_european_history": 0.72,
"philosophy": 0.68,
"world_religions": 0.74,
# Social Sciences: mix of factual recall and applied reasoning (~61-68%)
"anatomy": 0.62,
"business_ethics": 0.67,
"clinical_knowledge": 0.65,
"economics": 0.64,
# Other: practical/professional domains (~60-66%)
"miscellaneous": 0.66,
"professional_law": 0.60,
"security_studies": 0.63,
}
# Calculate category averages
category_results = {}
for category, subjects in categories.items():
cat_results = [
simulated_results[s] for s in subjects if s in simulated_results
]
if cat_results:
category_results[category] = sum(cat_results) / len(cat_results)
# Overall macro-average
overall = (
sum(simulated_results.values()) / len(simulated_results)
if simulated_results
else 0
)Results by Category: ---------------------------------------- STEM : 0.567 Humanities : 0.713 Social Sciences : 0.645 Other : 0.630 ---------------------------------------- Overall (Macro Avg) : 0.634
The category breakdown reveals how performance varies across knowledge domains, with STEM subjects often requiring different reasoning patterns than humanities or social sciences. The macro-average ensures each subject contributes equally to the final score, preventing larger categories from dominating the metric and giving a balanced view of model capabilities across diverse fields. When reporting results, researchers often give both the aggregate score and the category breakdowns to give a complete picture of model capabilities.
Key Evaluation Parameters
Understanding the configurable aspects of MMLU evaluation helps you reproduce published results and adapt the protocol for your own needs. The key parameters are:
-
k (few-shot examples): Number of examples to include in the prompt context. Standard practice uses k=5, though k=0 (zero-shot) is also reported. Higher k gives more task guidance but increases context length and computation. The choice of k involves a trade-off between format clarification and context window limitations. Some researchers experiment with k=10 or higher for models with very large context windows, though returns typically diminish after 5 examples.
-
max_questions: Subset size for evaluation. Full evaluation uses all test questions (typically 100-500 per subject), but smaller values let rapid testing during development. When comparing models, it is important to use the same subset or the full set to ensure comparability. Random subsampling should use fixed seeds for reproducibility.
-
subject_selection: Which of the 57 subjects to evaluate. Full MMLU reporting requires all subjects, though subject-specific analysis can isolate domain performance. Partial evaluation is useful for ablation studies or when computational resources are limited, but partial scores should not be compared to full benchmark results.
-
random_state: Seed for reproducibility when sampling few-shot examples or subsetting questions. Essential for comparable results across runs. Different random seeds selecting different development set examples can produce score variations of 1-2%, so fixed seeds are necessary for rigorous comparison.
-
model_temperature: For generation, temperature=0 (greedy decoding) is standard to ensure deterministic outputs. Greedy decoding ensures that the model selects the most likely next token, which is appropriate for multiple-choice questions where creativity is not desired. Non-zero temperatures introduce stochasticity that makes results non-reproducible.
-
scoring_method: Whether to use generation-based scoring (extract the first character) or probability-based scoring (compare log-probabilities for A, B, C, D). The two methods can differ by several percentage points for the same model, so published papers should always specify which they used.