open-thought/reasoning-gym
1
0
Fork 0
mirror of https://github.com/open-thought/reasoning-gym.git synced 2026-04-22 16:49:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

pass config to ProceduralDataset base

Browse source
This commit is contained in:
Andreas Koepf 2025-01-25 00:23:05 +01:00
parent b3a2b719c7
commit 0d2d8ba6a0
20 changed files with 45 additions and 80 deletions
Download patch file Download diff file Expand all files Collapse all files

7
reasoning_gym/arithmetic/chain_sum.py
View file

@ -18,7 +18,7 @@ class ChainSumConfig:
seed: Optional[int] = None
size: int = 500
def validate(self):
def validate(self) -> None:
"""Validate configuration parameters"""
assert self.min_terms > 0, "min_terms must be positive"
assert self.max_terms >= self.min_terms, "max_terms must be >= min_terms"
@ -34,9 +34,7 @@ class ChainSum(ProceduralDataset):
"""Generates simple arithmetic tasks using only + and - operators"""
def __init__(self, config: ChainSumConfig):
self.config = config
self.config.validate()
super().__init__(seed=config.seed, size=config.size)
super().__init__(config=config, seed=config.seed, size=config.size)
def __getitem__(self, idx: int) -> dict:
"""Generate a single chain sum task
@ -145,5 +143,6 @@ def chain_sum_dataset(
)
return ChainSum(config)
# Register the dataset
register_dataset("chain_sum", ChainSum, ChainSumConfig)