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

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

6
reasoning_gym/arithmetic/basic_arithmetic.py
View file

@ -21,7 +21,7 @@ class BasicArithmeticDatasetConfig:
format_style: Literal["simple", "natural"] = "simple"
whitespace: Literal["no_space", "single", "random"] = "single" # Whitespace style between terms
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"
@ -63,9 +63,7 @@ class BasicArithmeticDataset(ProceduralDataset):
"""Dataset that generates basic arithmetic tasks with configurable complexity"""
def __init__(self, config: BasicArithmeticDatasetConfig):
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[str, Any]:
"""Generate a single arithmetic task