refactor: Inherit GCDDataset from ProceduralDataset base class

This commit is contained in:
Andreas Koepf (aider) 2025-01-24 09:55:14 +01:00
parent 018bdcef6f
commit c3170fd5f1

View file

@ -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.