Spaces:
Running
Running
| from .base_benchmark import BaseBenchmark | |
| from typing import Dict, Any, Optional, Tuple | |
| from datasets import load_dataset | |
| import re | |
| from .prompt_templates import get_gsm8k_cot_prompt | |
| class GSM8KBenchmark(BaseBenchmark): | |
| """GSM8K (Grade School Math 8K) benchmark""" | |
| def __init__(self): | |
| super().__init__(name="GSM8K", dataset_name="gsm8k") | |
| async def load_dataset(self, sample_size: Optional[int] = None, **kwargs): | |
| """Load GSM8K dataset""" | |
| dataset = load_dataset(self.dataset_name, 'main', split='test') | |
| self.dataset = [] | |
| for sample in dataset: | |
| self.dataset.append({ | |
| 'question': sample['question'], | |
| 'answer': sample['answer'], | |
| 'raw_sample': sample | |
| }) | |
| # Shuffle dataset | |
| import random | |
| random.shuffle(self.dataset) | |
| if sample_size and len(self.dataset) > sample_size: | |
| self.dataset = self.dataset[:sample_size] | |
| def extract_answer_from_solution(self, solution: str) -> Optional[str]: | |
| """Extract numerical answer from GSM8K solution string""" | |
| # GSM8K answers are in format: "... #### number" | |
| match = re.search(r'#### ([\-0-9\.\,]+)', solution) | |
| if match: | |
| answer_str = match.group(1).replace(',', '') | |
| return answer_str | |
| return None | |
| def extract_number_from_response(self, response: str) -> Optional[str]: | |
| """Extract the final numerical answer from model response""" | |
| # Official lm-eval uses these patterns in order: | |
| # 1. Look for "The answer is X" pattern (CoT standard) | |
| match = re.search(r'The answer is ([\-0-9\.\,]+)\.?', response, re.IGNORECASE) | |
| if match: | |
| return match.group(1).replace(',', '') | |
| # 2. Look for #### format (if model knows GSM8K format) | |
| match = re.search(r'#### ([\-0-9\.\,]+)', response) | |
| if match: | |
| return match.group(1).replace(',', '') | |
| # 3. Flexible extraction: find all numbers and take the last one | |
| # This matches lm-eval's flexible-extract with group_select: -1 | |
| numbers = re.findall(r'(-?[$0-9.,]{2,})|(-?[0-9]+)', response) | |
| if numbers: | |
| # Flatten tuples and get last non-empty match | |
| flat_numbers = [n for group in numbers for n in group if n] | |
| if flat_numbers: | |
| last_number = flat_numbers[-1] | |
| # Clean the number | |
| cleaned = last_number.replace('$', '').replace(',', '') | |
| try: | |
| # Validate it's a proper number | |
| float(cleaned) | |
| return cleaned | |
| except: | |
| pass | |
| return None | |
| def format_prompt(self, sample: Dict[str, Any]) -> str: | |
| """Format GSM8K question as prompt with CoT examples""" | |
| # Use the standard CoT prompt from lm-eval | |
| return get_gsm8k_cot_prompt(sample['question']) | |
| async def evaluate_sample(self, api, sample: Dict[str, Any], **kwargs) -> Tuple[bool, Dict[str, Any]]: | |
| """Evaluate a single GSM8K sample""" | |
| prompt = self.format_prompt(sample) | |
| try: | |
| response = await api.generate_with_retry(prompt, **kwargs) | |
| # Extract correct answer | |
| correct_answer = self.extract_answer_from_solution(sample['answer']) | |
| # Extract model's answer | |
| model_answer = self.extract_number_from_response(response) | |
| # Check if answers match (exact string match after normalization) | |
| is_correct = False | |
| if correct_answer is not None and model_answer is not None: | |
| # GSM8K uses exact match on normalized strings | |
| is_correct = correct_answer == model_answer | |
| result = { | |
| 'question': sample['question'], | |
| 'correct_answer': correct_answer, | |
| 'model_answer': model_answer, | |
| 'model_response': response, | |
| 'is_correct': is_correct, | |
| 'solution': sample['answer'] | |
| } | |
| return is_correct, result | |
| except Exception as e: | |
| result = { | |
| 'question': sample['question'], | |
| 'error': str(e), | |
| 'is_correct': False | |
| } | |
| return False, result |