Spaces:
Running
Running
| from .base_benchmark import BaseBenchmark | |
| from typing import Dict, Any, Optional, Tuple, List | |
| from datasets import load_dataset | |
| import re | |
| import random | |
| from .prompt_templates import get_mmlu_prompt | |
| from .evaluation_utils import extract_answer_mmlu | |
| class MMLUBenchmark(BaseBenchmark): | |
| """MMLU (Massive Multitask Language Understanding) benchmark""" | |
| SUBJECTS = [ | |
| 'abstract_algebra', 'anatomy', 'astronomy', 'business_ethics', | |
| 'clinical_knowledge', 'college_biology', 'college_chemistry', | |
| 'college_computer_science', 'college_mathematics', 'college_medicine', | |
| 'college_physics', 'computer_security', 'conceptual_physics', | |
| 'econometrics', 'electrical_engineering', 'elementary_mathematics', | |
| 'formal_logic', 'global_facts', 'high_school_biology', | |
| 'high_school_chemistry', 'high_school_computer_science', | |
| 'high_school_european_history', 'high_school_geography', | |
| 'high_school_government_and_politics', 'high_school_macroeconomics', | |
| 'high_school_mathematics', 'high_school_microeconomics', | |
| 'high_school_physics', 'high_school_psychology', 'high_school_statistics', | |
| 'high_school_us_history', 'high_school_world_history', 'human_aging', | |
| 'human_sexuality', 'international_law', 'jurisprudence', | |
| 'logical_fallacies', 'machine_learning', 'management', 'marketing', | |
| 'medical_genetics', 'miscellaneous', 'moral_disputes', 'moral_scenarios', | |
| 'nutrition', 'philosophy', 'prehistory', 'professional_accounting', | |
| 'professional_law', 'professional_medicine', 'professional_psychology', | |
| 'public_relations', 'security_studies', 'sociology', 'us_foreign_policy', | |
| 'virology', 'world_religions' | |
| ] | |
| def __init__(self): | |
| super().__init__(name="MMLU", dataset_name="cais/mmlu") | |
| async def load_dataset(self, sample_size: Optional[int] = None, **kwargs): | |
| """Load MMLU dataset""" | |
| subjects = kwargs.get('subjects', ['all']) | |
| if 'all' in subjects: | |
| subjects = self.SUBJECTS | |
| else: | |
| subjects = [s for s in subjects if s in self.SUBJECTS] | |
| self.dataset = [] | |
| self.few_shot_examples = {} # Store few-shot examples per subject | |
| for subject in subjects: | |
| try: | |
| # Load dev split for few-shot examples | |
| dev_ds = load_dataset(self.dataset_name, subject, split='dev') | |
| # Standard MMLU uses 5-shot | |
| self.few_shot_examples[subject] = [ | |
| { | |
| 'question': ex['question'], | |
| 'choices': ex['choices'], | |
| 'answer': ex['answer'] | |
| } | |
| for ex in list(dev_ds)[:5] | |
| ] | |
| # Load test split for evaluation | |
| test_ds = load_dataset(self.dataset_name, subject, split='test') | |
| for sample in test_ds: | |
| self.dataset.append({ | |
| 'subject': subject, | |
| 'question': sample['question'], | |
| 'choices': sample['choices'], | |
| 'answer': sample['answer'], # 0-3 index | |
| 'raw_sample': sample | |
| }) | |
| except Exception as e: | |
| print(f"Error loading {subject}: {e}") | |
| continue | |
| # Shuffle dataset | |
| random.shuffle(self.dataset) | |
| if sample_size and len(self.dataset) > sample_size: | |
| self.dataset = self.dataset[:sample_size] | |
| def format_prompt(self, sample: Dict[str, Any]) -> str: | |
| """Format MMLU question as prompt with few-shot examples""" | |
| subject = sample['subject'] | |
| few_shot_examples = self.few_shot_examples.get(subject, []) | |
| return get_mmlu_prompt( | |
| question=sample['question'], | |
| choices=sample['choices'], | |
| subject=subject.replace('_', ' ').title(), | |
| few_shot_examples=few_shot_examples | |
| ) | |
| async def evaluate_sample(self, api, sample: Dict[str, Any], **kwargs) -> Tuple[bool, Dict[str, Any]]: | |
| """Evaluate a single MMLU sample""" | |
| prompt = self.format_prompt(sample) | |
| try: | |
| response = await api.generate_with_retry(prompt, **kwargs) | |
| # Extract answer from response using standard extraction | |
| predicted_letter = extract_answer_mmlu(response) | |
| if predicted_letter: | |
| predicted_index = ord(predicted_letter) - ord('A') | |
| else: | |
| # If no clear answer, mark as incorrect | |
| predicted_index = -1 | |
| correct_index = sample['answer'] | |
| is_correct = predicted_index == correct_index | |
| result = { | |
| 'subject': sample['subject'], | |
| 'question': sample['question'], | |
| 'choices': sample['choices'], | |
| 'correct_answer': correct_index, | |
| 'predicted_answer': predicted_index, | |
| 'model_response': response, | |
| 'is_correct': is_correct | |
| } | |
| return is_correct, result | |
| except Exception as e: | |
| result = { | |
| 'subject': sample['subject'], | |
| 'question': sample['question'], | |
| 'error': str(e), | |
| 'is_correct': False | |
| } | |
| return False, result |