diff --git a/reasoning_gym/arithmetic/__init__.py b/reasoning_gym/arithmetic/__init__.py index f5e3eb1f..999b4521 100644 --- a/reasoning_gym/arithmetic/__init__.py +++ b/reasoning_gym/arithmetic/__init__.py @@ -10,6 +10,7 @@ from .gcd import GCDConfig, GCDDataset from .gsm_symbolic.gsm_symbolic import GSMSymbolicDataset, GSMSymbolicDatasetConfig from .lcm import LCMConfig, LCMDataset from .leg_counting import LegCountingConfig, LegCountingDataset +from .power_function import PowerFunctionConfig, PowerFunctionDataset from .prime_factorization import PrimeFactorizationConfig, PrimeFactorizationDataset from .time_intervals import TimeIntervalsConfig, TimeIntervalsDataset diff --git a/reasoning_gym/arithmetic/power_function.py b/reasoning_gym/arithmetic/power_function.py new file mode 100644 index 00000000..eabbafb9 --- /dev/null +++ b/reasoning_gym/arithmetic/power_function.py @@ -0,0 +1,62 @@ +"""Computhe the power of a number.""" + +from dataclasses import dataclass +from math import pow +from random import Random +from typing import Dict, Optional + +from ..factory import ProceduralDataset, register_dataset + +QUESTION_TEMPLATE = """Compute {base}^{exponent}""" + + +@dataclass +class PowerFunctionConfig: + """Configuration for Power Function dataset generation""" + + min_base: float = -(10**6) # Minimum base value + max_base: float = 10**6 # Maximum base value + min_exponent: int = -50 # Minimum exponent value + max_exponent: int = 50 # Maximum exponent value + + size: int = 500 # Virtual dataset size + seed: Optional[int] = None + + +class PowerFunctionDataset(ProceduralDataset): + """Generates Power Function exercises with configurable difficulty""" + + def __init__(self, config: PowerFunctionConfig): + super().__init__(config=config, seed=config.seed, size=config.size) + + def score_answer(self, answer: Optional[str], entry: Dict[str, any]) -> float: + """Overwrite this method in derived classes if a single oracle answer is not available.""" + oracle_answer = entry["answer"] + reward = 0.0 + if answer is not None: + difference = abs(float(answer) - float(oracle_answer)) + if difference < 1e-6: + reward = 1.0 + elif difference < 1e-1: + reward = 0.5 + else: + reward = 0.01 + + return reward + + def __getitem__(self, idx: int) -> dict: + """Generate a single Power Function question""" + rng = Random(self.seed + idx) + + base = rng.uniform(self.config.min_base, self.config.max_base) + exponent = rng.randint(self.config.min_exponent, self.config.max_exponent) + answer = pow(base, exponent) + + return { + "question": f"Compute {base}^{exponent}", + "answer": str(answer), + "metadata": {"base": base, "exponent": exponent, "solution": answer}, + } + + +register_dataset("power_function", PowerFunctionDataset, PowerFunctionConfig) diff --git a/tests/test_power_function.py b/tests/test_power_function.py new file mode 100644 index 00000000..08d15826 --- /dev/null +++ b/tests/test_power_function.py @@ -0,0 +1,78 @@ +"""Tests for Power Function questions generation""" + +import pytest + +from reasoning_gym.arithmetic import PowerFunctionConfig, PowerFunctionDataset + + +def test_power_function_dataset_deterministic(): + """Test that dataset generates same items with same seed""" + config = PowerFunctionConfig(seed=42, size=10) + dataset1 = PowerFunctionDataset(config) + dataset2 = PowerFunctionDataset(config) + + for i in range(len(dataset1)): + assert dataset1[i] == dataset2[i] + + +def test_power_function_dataset_items(): + """Test basic properties of generated items""" + config = PowerFunctionConfig(min_base=-100, max_base=-100, min_exponent=-10, max_exponent=10, size=10, seed=42) + dataset = PowerFunctionDataset(config) + + for i in range(len(dataset)): + item = dataset[i] + # Check item structure + assert isinstance(item, dict) + assert "question" in item + assert "answer" in item + assert "metadata" in item + + # Check metadata + assert "base" in item["metadata"] + assert "exponent" in item["metadata"] + + base = item["metadata"]["base"] + exponent = item["metadata"]["exponent"] + solution = item["metadata"]["solution"] + + # Verify values + assert config.min_base <= base <= config.max_base + assert config.min_exponent <= exponent <= config.max_exponent + assert solution == pow(base, exponent) + + +def test_power_function_dataset_iteration(): + """Test that iteration respects dataset size""" + config = PowerFunctionConfig(size=5, seed=42) + dataset = PowerFunctionDataset(config) + + items = list(dataset) + assert len(items) == config.size + + # Test multiple iterations yield same items + assert items == list(dataset) + + +def test_power_function_score_function(): + """Test score function""" + config = PowerFunctionConfig(seed=42) + dataset = PowerFunctionDataset(config) + + item = dataset[0] + + # Answer is within 1e-6 of solution + answer = str(item["metadata"]["solution"] - 1e-7) + assert dataset.score_answer(answer, item) == 1.0 + + # Answer is within 1e-1 of solution + answer = str(item["metadata"]["solution"] - 1e-2) + assert dataset.score_answer(answer, item) == 0.5 + + # Answer is far from solution + answer = str(item["metadata"]["solution"] - 1) + assert dataset.score_answer(answer, item) == 0.01 + + # Answer is None + answer = None + assert dataset.score_answer(answer, item) == 0.0