open-thought/reasoning-gym
1
0
Fork 0
mirror of https://github.com/open-thought/reasoning-gym.git synced 2026-04-28 17:29:39 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix primes representation in count_primes dataset metadata

Browse source
This commit is contained in:
AhmedSaif2 2025-02-26 14:58:21 +02:00
parent c0e5941fe5
commit e9e36f3a23
2 changed files with 17 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
reasoning_gym/algorithmic/count_primes.py
View file

@ -51,8 +51,8 @@ class CountPrimesDataset(ProceduralDataset):
rng = Random(self.seed + idx) rng = Random(self.seed + idx)
start = rng.randint(1, self.config.max_n) start = rng.randint(1, self.config.max_n)
end = rng.randint(start, self.config.max_n) end = rng.randint(start, self.config.max_n)
primes = self.primes[start : end + 1] primes = [i for i in range(start, end + 1) if self.primes[i]]
answer = sum(primes) answer = len(primes)
return { return {
"question": QUESTION_TEMPLATE.format(start=start, end=end), "question": QUESTION_TEMPLATE.format(start=start, end=end),
"answer": str(answer), "answer": str(answer),

15
tests/test_count_primes.py
View file

@ -86,3 +86,18 @@ def test_count_primes_answer():
assert primes[8] == False assert primes[8] == False
assert primes[9] == False assert primes[9] == False
assert primes[10] == False assert primes[10] == False
def test_count_primes_list():
"""Test that list of primes was correctly generated"""
config = CountPrimesConfig(max_n=100, size=100, seed=42)
dataset = CountPrimesDataset(config)
for item in dataset:
start = item["metadata"]["start"]
end = item["metadata"]["end"]
primes = item["metadata"]["primes"]
for p in primes:
assert p >= start
assert p <= end
assert dataset.primes[p] == True