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

refactor: Move min_value and max_value calculations to __getitem__

Browse source
This commit is contained in:
Andreas Koepf (aider) 2025-01-23 12:05:55 +01:00
parent 1ff01627cc
commit e1734b0b11
1 changed files with 8 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

15
reasoning_gym/arithmetic/chain_sum.py
View file

@ -56,7 +56,11 @@ class ChainSum:
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)
expression, result = self._generate_task(item_rng, num_terms, num_digits)
# Calculate value ranges based on number of digits
min_value = 10 ** (num_digits - 1) # e.g., 100 for 3 digits
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)
return {
"question": f"{expression} =",
@ -68,21 +72,18 @@ class ChainSum:
}
}
def _generate_task(self, rng: random.Random, num_terms: int, num_digits: int) -> tuple[str, int]:
def _generate_task(self, rng: random.Random, num_terms: int, min_value: int, max_value: int) -> tuple[str, int]:
"""Generate a chain sum task
Args:
rng: Random number generator
num_terms: Number of terms in the expression
num_digits: Number of digits for each number
min_value: Minimum value for generated numbers
max_value: Maximum value for generated numbers
Returns:
Tuple of (expression string, result integer)
"""
# Generate numbers with at least min_digits
min_value = 10 ** (num_digits - 1) # e.g., 100 for 3 digits
max_value = (10 ** num_digits) - 1 # e.g., 999 for 3 digits
if self.config.allow_negation:
# Allow both positive and negative numbers in the range
constants = [rng.randint(-max_value, max_value) for _ in range(num_terms)]