mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
refactor: Inherit GCDDataset from ProceduralDataset base class
This commit is contained in:
parent
018bdcef6f
commit
c3170fd5f1
1 changed files with 3 additions and 16 deletions
|
|
@ -2,6 +2,7 @@
|
|||
from dataclasses import dataclass
|
||||
from random import Random
|
||||
from typing import List, Optional, Tuple
|
||||
from ..dataset import ProceduralDataset
|
||||
from math import gcd
|
||||
from functools import reduce
|
||||
|
||||
|
|
@ -24,27 +25,13 @@ class GCDConfig:
|
|||
assert self.max_value > self.min_value, "max_value must be > min_value"
|
||||
|
||||
|
||||
class GCDDataset:
|
||||
class GCDDataset(ProceduralDataset):
|
||||
"""Generates Greatest Common Divisor (GCD) tasks"""
|
||||
|
||||
def __init__(self, config: GCDConfig):
|
||||
self.config = config
|
||||
self.config.validate()
|
||||
self.seed = config.seed if config.seed is not None else Random().randint(0, 2**32)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.config.size
|
||||
|
||||
def __iter__(self):
|
||||
self._current_idx = 0
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if self._current_idx >= self.config.size:
|
||||
raise StopIteration
|
||||
item = self[self._current_idx]
|
||||
self._current_idx += 1
|
||||
return item
|
||||
super().__init__(seed=config.seed, size=config.size)
|
||||
|
||||
def _generate_numbers(self, rng: Random) -> Tuple[List[int], int]:
|
||||
"""Generate a list of random positive integers and their GCD.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue