diff --git a/reasoning_gym/arithmetic/gcd.py b/reasoning_gym/arithmetic/gcd.py index 0f0436ec..ba0a69cc 100644 --- a/reasoning_gym/arithmetic/gcd.py +++ b/reasoning_gym/arithmetic/gcd.py @@ -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.