open-thought/reasoning-gym
1
0
Fork 0
mirror of https://github.com/open-thought/reasoning-gym.git synced 2026-04-19 12:58:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add ProductsDataset (multiplication tasks)

Browse source
This commit is contained in:
Andreas Koepf 2025-02-13 17:59:02 +01:00
parent ce30536627
commit 5410bb78a0
10 changed files with 56 additions and 56 deletions
Download patch file Download diff file Expand all files Collapse all files

10
reasoning_gym/arc/arc_1d.py
View file

@ -58,27 +58,27 @@ class Arc1DDataset(ProceduralDataset):
- metadata: dict with generation parameters
"""
# Create deterministic RNG from base seed and idx
item_rng = Random(self.seed + idx)
rng = Random(self.seed + idx)
# Select random task
task_name = item_rng.choice(self.task_names)
task_name = rng.choice(self.task_names)
task_func, task_kwargs = self.ARC_1D_TASKS[task_name]
# Generate training examples
train_examples = []
size = item_rng.randint(self.config.min_size, self.config.max_size)
size = rng.randint(self.config.min_size, self.config.max_size)
for _ in range(self.config.num_train):
example = None
while example is None:
example = task_func(item_rng, size, **task_kwargs)
example = task_func(rng, size, **task_kwargs)
train_examples.append(example)
# Generate test example
test_example = None
while test_example is None:
test_example = task_func(item_rng, size, **task_kwargs)
test_example = task_func(rng, size, **task_kwargs)
# Format question
question = "Find the common rule that maps an input grid to an output grid, given the examples below.\n\n"

8
reasoning_gym/arithmetic/__init__.py
View file

@ -4,7 +4,7 @@ Arithmetic tasks for training reasoning capabilities:
from .basic_arithmetic import BasicArithmeticDataset, BasicArithmeticDatasetConfig
from .calendar_arithmetic import CalendarArithmeticConfig, CalendarArithmeticDataset
from .chain_sum import ChainSum, ChainSumConfig
from .chain_sum import ChainSumConfig, ChainSumDataset
from .count_bits import CountBitsConfig, CountBitsDataset
from .dice import DiceConfig, DiceDataset
from .fraction_simplification import FractionSimplificationConfig, FractionSimplificationDataset
@ -14,13 +14,13 @@ from .lcm import LCMConfig, LCMDataset
from .leg_counting import LegCountingConfig, LegCountingDataset
from .power_function import PowerFunctionConfig, PowerFunctionDataset
from .prime_factorization import PrimeFactorizationConfig, PrimeFactorizationDataset
from .products import Products, ProductsConfig
from .products import ProductsConfig, ProductsDataset
from .time_intervals import TimeIntervalsConfig, TimeIntervalsDataset
__all__ = [
"BasicArithmeticDataset",
"BasicArithmeticDatasetConfig",
"ChainSum",
"ChainSumDataset",
"ChainSumConfig",
"CalendarArithmeticConfig",
"CalendarArithmeticDataset",
@ -36,7 +36,7 @@ __all__ = [
"PowerFunctionDataset",
"PrimeFactorizationConfig",
"PrimeFactorizationDataset",
"Products",
"ProductsDataset",
"ProductsConfig",
"GSMSymbolicDatasetConfig",
"GSMSymbolicDataset",

12
reasoning_gym/arithmetic/basic_arithmetic.py
View file

@ -78,17 +78,17 @@ class BasicArithmeticDataset(ProceduralDataset):
- metadata: dict with generation parameters
"""
# Create deterministic RNG from base seed and idx
item_rng = Random(self.seed + idx)
rng = Random(self.seed + idx)
num_terms = item_rng.randint(self.config.min_terms, self.config.max_terms)
num_digits = item_rng.randint(self.config.min_digits, self.config.max_digits)
num_terms = rng.randint(self.config.min_terms, self.config.max_terms)
num_digits = rng.randint(self.config.min_digits, self.config.max_digits)
if self.config.allow_parentheses:
expression, result = self._generate_complex_task(item_rng, num_terms, num_digits)
expression, result = self._generate_complex_task(rng, num_terms, num_digits)
else:
expression, result = self._generate_simple_task(item_rng, num_terms, num_digits)
expression, result = self._generate_simple_task(rng, num_terms, num_digits)
question = self._format_question(item_rng, expression)
question = self._format_question(rng, expression)
return {
"question": question,

6
reasoning_gym/arithmetic/calendar_arithmetic.py
View file

@ -122,9 +122,9 @@ class CalendarArithmeticDataset(ProceduralDataset):
self.tasks = [self.task_handlers[task] for task in self.config.tasks]
def __getitem__(self, idx: int) -> dict:
item_rng = random.Random(self.seed + idx)
task = item_rng.choice(self.tasks)
question, answer, metadata = task(item_rng)
rng = random.Random(self.seed + idx)
task = rng.choice(self.tasks)
question, answer, metadata = task(rng)
return {
"question": question,
"answer": str(answer),

12
reasoning_gym/arithmetic/chain_sum.py
View file

@ -32,7 +32,7 @@ class ChainSumConfig:
assert 10 ** (self.min_digits - 1) >= 1, "min_digits would result in invalid number range"
class ChainSum(ProceduralDataset):
class ChainSumDataset(ProceduralDataset):
"""Generates simple arithmetic tasks using only + and - operators"""
def __init__(self, config: ChainSumConfig):
@ -51,16 +51,16 @@ class ChainSum(ProceduralDataset):
- metadata: dict with generation parameters
"""
# Create deterministic RNG from base seed and idx
item_rng = random.Random(self.seed + idx)
rng = random.Random(self.seed + idx)
num_terms = item_rng.randint(self.config.min_terms, self.config.max_terms)
num_digits = item_rng.randint(self.config.min_digits, self.config.max_digits)
num_terms = rng.randint(self.config.min_terms, self.config.max_terms)
num_digits = rng.randint(self.config.min_digits, self.config.max_digits)
# Calculate value ranges based on number of digits
min_value = 0 if num_digits == 1 else 10 ** (num_digits - 1) # Special case for 1 digit
max_value = (10**num_digits) - 1 # e.g., 999 for 3 digits
expression, result = self._generate_task(item_rng, num_terms, min_value, max_value)
expression, result = self._generate_task(rng, num_terms, min_value, max_value)
return {
"question": f"{expression} =",
@ -143,4 +143,4 @@ class ChainSumCurriculum(BaseCurriculum):
# Register the dataset
register_dataset("chain_sum", ChainSum, ChainSumConfig)
register_dataset("chain_sum", ChainSumDataset, ChainSumConfig)

12
reasoning_gym/arithmetic/products.py
View file

@ -26,7 +26,7 @@ class ProductsConfig:
assert self.max_digits >= self.min_digits, "max_digits must be >= min_digits"
class Products(ProceduralDataset):
class ProductsDataset(ProceduralDataset):
"""Generates multiplication tasks with configurable number of terms"""
def __init__(self, config: ProductsConfig):
@ -45,16 +45,16 @@ class Products(ProceduralDataset):
- metadata: dict with generation parameters
"""
# Create deterministic RNG from base seed and idx
item_rng = random.Random(self.seed + idx)
rng = random.Random(self.seed + idx)
num_terms = item_rng.randint(self.config.min_terms, self.config.max_terms)
num_digits = item_rng.randint(self.config.min_digits, self.config.max_digits)
num_terms = rng.randint(self.config.min_terms, self.config.max_terms)
num_digits = rng.randint(self.config.min_digits, self.config.max_digits)
# Calculate value ranges based on number of digits
min_value = 0 if num_digits == 1 else 10 ** (num_digits - 1) # Special case for 1 digit
max_value = (10**num_digits) - 1 # e.g., 999 for 3 digits
expression, result = self._generate_task(item_rng, num_terms, min_value, max_value)
expression, result = self._generate_task(rng, num_terms, min_value, max_value)
return {
"question": f"{expression} =",
@ -127,4 +127,4 @@ class ProductsCurriculum(BaseCurriculum):
# Register the dataset
register_dataset("products", Products, ProductsConfig)
register_dataset("products", ProductsDataset, ProductsConfig)

8
reasoning_gym/arithmetic/time_intervals.py
View file

@ -82,14 +82,14 @@ class TimeIntervalsDataset(ProceduralDataset):
def __getitem__(self, idx: int) -> dict:
"""Generate a single time interval calculation task"""
item_rng = random.Random(self.seed + idx)
rng = random.Random(self.seed + idx)
# Randomly choose task type from config
task_type = item_rng.choice(self.config.task_types)
task_type = rng.choice(self.config.task_types)
start_time, end_time, format_str, expected_format = self._generate_times(item_rng, task_type)
start_time, end_time, format_str, expected_format = self._generate_times(rng, task_type)
template = item_rng.choice(self.TEMPLATES)
template = rng.choice(self.TEMPLATES)
question = template.format(start=start_time, end=end_time, format=expected_format)
# Calculate the actual difference

16
tests/test_chain_sum.py
View file

@ -1,6 +1,6 @@
import pytest
from reasoning_gym.arithmetic import ChainSum, ChainSumConfig
from reasoning_gym.arithmetic import ChainSumConfig, ChainSumDataset
from reasoning_gym.arithmetic.chain_sum import ChainSumCurriculum
@ -18,8 +18,8 @@ def test_chain_sum_config_validation():
def test_chain_sum_deterministic():
"""Test that dataset generates same items with same seed"""
config = ChainSumConfig(seed=42, size=10)
dataset1 = ChainSum(config)
dataset2 = ChainSum(config)
dataset1 = ChainSumDataset(config)
dataset2 = ChainSumDataset(config)
for i in range(len(dataset1)):
assert dataset1[i] == dataset2[i]
@ -28,7 +28,7 @@ def test_chain_sum_deterministic():
def test_chain_sum_items():
"""Test basic properties of generated items"""
config = ChainSumConfig(min_terms=2, max_terms=4, min_digits=1, max_digits=2, size=100, seed=42)
dataset = ChainSum(config)
dataset = ChainSumDataset(config)
for i in range(len(dataset)):
item = dataset[i]
@ -57,7 +57,7 @@ def test_chain_sum_number_ranges():
size=50,
seed=42,
)
dataset = ChainSum(config)
dataset = ChainSumDataset(config)
for i in range(len(dataset)):
item = dataset[i]
@ -71,7 +71,7 @@ def test_chain_sum_number_ranges():
# Test 1-digit numbers
config = ChainSumConfig(min_terms=2, max_terms=2, min_digits=1, max_digits=1, size=50, seed=42)
dataset = ChainSum(config)
dataset = ChainSumDataset(config)
for i in range(len(dataset)):
item = dataset[i]
expression = item["metadata"]["expression"]
@ -88,7 +88,7 @@ def test_chain_sum_negation():
config = ChainSumConfig(
min_terms=2, max_terms=2, min_digits=2, max_digits=2, size=100, seed=42, allow_negation=True
)
dataset = ChainSum(config)
dataset = ChainSumDataset(config)
# Track if we see both positive and negative numbers
has_positive = False
@ -112,7 +112,7 @@ def test_chain_sum_negation():
def test_chain_sum_iteration():
"""Test that iteration respects dataset size"""
config = ChainSumConfig(min_terms=2, max_terms=2, size=5, seed=42) # Small size for testing
dataset = ChainSum(config)
dataset = ChainSumDataset(config)
# Test manual iteration
items = []

6
tests/test_coaching.py
View file

@ -5,7 +5,7 @@ from pathlib import Path
import pytest
from reasoning_gym.arithmetic.chain_sum import ChainSum, ChainSumConfig
from reasoning_gym.arithmetic.chain_sum import ChainSumConfig, ChainSumDataset
from reasoning_gym.arithmetic.leg_counting import LegCountingConfig
from reasoning_gym.coaching import Coach, GroupedScores
from reasoning_gym.composite import CompositeConfig, CompositeDataset, DatasetSpec
@ -14,7 +14,7 @@ from reasoning_gym.composite import CompositeConfig, CompositeDataset, DatasetSp
def test_coach_with_chain_sum():
# Create a small ChainSum dataset
config = ChainSumConfig(min_terms=2, max_terms=3, min_digits=1, max_digits=2, size=10, seed=42)
dataset = ChainSum(config)
dataset = ChainSumDataset(config)
coach = Coach(dataset)
# Simulate an agent working on tasks
@ -208,7 +208,7 @@ def test_coach_score_logging(tmp_path):
# Create dataset and coach with logging
config = ChainSumConfig(min_terms=2, max_terms=3, min_digits=1, max_digits=2, size=10, seed=42)
dataset = ChainSum(config)
dataset = ChainSumDataset(config)
coach = Coach(dataset, score_log=log_file)
# Score a few answers

16
tests/test_products.py
View file

@ -1,6 +1,6 @@
import pytest
from reasoning_gym.arithmetic import Products, ProductsConfig
from reasoning_gym.arithmetic import ProductsConfig, ProductsDataset
from reasoning_gym.arithmetic.products import ProductsCurriculum
@ -18,8 +18,8 @@ def test_products_config_validation():
def test_products_deterministic():
"""Test that dataset generates same items with same seed"""
config = ProductsConfig(seed=42, size=10)
dataset1 = Products(config)
dataset2 = Products(config)
dataset1 = ProductsDataset(config)
dataset2 = ProductsDataset(config)
for i in range(len(dataset1)):
assert dataset1[i] == dataset2[i]
@ -28,7 +28,7 @@ def test_products_deterministic():
def test_products_items():
"""Test basic properties of generated items"""
config = ProductsConfig(min_terms=2, max_terms=4, min_digits=1, max_digits=2, size=100, seed=42)
dataset = Products(config)
dataset = ProductsDataset(config)
for i in range(len(dataset)):
item = dataset[i]
@ -57,7 +57,7 @@ def test_products_number_ranges():
size=50,
seed=42,
)
dataset = Products(config)
dataset = ProductsDataset(config)
for i in range(len(dataset)):
item = dataset[i]
@ -68,7 +68,7 @@ def test_products_number_ranges():
# Test 1-digit numbers
config = ProductsConfig(min_terms=2, max_terms=2, min_digits=1, max_digits=1, size=50, seed=42)
dataset = Products(config)
dataset = ProductsDataset(config)
for i in range(len(dataset)):
item = dataset[i]
expression = item["metadata"]["expression"]
@ -80,7 +80,7 @@ def test_products_number_ranges():
def test_products_iteration():
"""Test that iteration respects dataset size"""
config = ProductsConfig(min_terms=2, max_terms=2, size=5, seed=42) # Small size for testing
dataset = Products(config)
dataset = ProductsDataset(config)
# Test manual iteration
items = []
@ -101,7 +101,7 @@ def test_products_iteration():
def test_products_scoring():
"""Test that scoring works correctly"""
config = ProductsConfig(min_terms=2, max_terms=2, size=10, seed=42)
dataset = Products(config)
dataset = ProductsDataset(config)
# Test scoring with exact match
item = dataset[0]