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

lint

Browse source
This commit is contained in:
Zafir Stojanovski 2025-02-11 14:44:46 +01:00
parent 3873c50ac6
commit d647498c43
3 changed files with 9 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

11
reasoning_gym/algorithmic/count_primes.py
View file

@ -13,6 +13,7 @@ from ..factory import ProceduralDataset, register_dataset
QUESTION_TEMPLATE = """Count how many prime numbers there are between {start} and {end} (inclusive) ?"""
@dataclass
class CountPrimesConfig:
"""Configuration for Count Primes dataset generation"""
@ -26,21 +27,22 @@ class CountPrimesConfig:
"""Validate configuration parameters"""
assert 1 <= self.max_n, "max_n must be at least 1"
class CountPrimesDataset(ProceduralDataset):
"""Generates Count Primes exercises with configurable difficulty"""
def __init__(self, config: CountPrimesConfig):
super().__init__(config=config, seed=config.seed, size=config.size)
self.primes = self._get_primes(config.max_n + 1)
def _get_primes(self, n: int) -> list[bool]:
if n <= 1:
return []
primes = [True] * n
primes[0] = primes[1] = False
for i in range(2, int(math.sqrt(n))+1):
for i in range(2, int(math.sqrt(n)) + 1):
if primes[i]:
for j in range(2*i, n, i):
for j in range(2 * i, n, i):
primes[j] = False
return primes
@ -49,7 +51,7 @@ class CountPrimesDataset(ProceduralDataset):
rng = Random(self.seed + idx)
start = rng.randint(1, self.config.max_n)
end = rng.randint(start, self.config.max_n)
primes = self.primes[start:end+1]
primes = self.primes[start : end + 1]
answer = sum(primes)
return {
"question": QUESTION_TEMPLATE.format(start=start, end=end),
@ -57,4 +59,5 @@ class CountPrimesDataset(ProceduralDataset):
"metadata": {"start": start, "end": end, "primes": primes, "solution": answer},
}
register_dataset("count_primes", CountPrimesDataset, CountPrimesConfig)