refactor: Inherit LCMDataset from ProceduralDataset base class

This commit is contained in:
Andreas Koepf (aider) 2025-01-24 09:56:19 +01:00
parent c3170fd5f1
commit 0dc19b831c

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 lcm
from functools import reduce
@ -24,27 +25,13 @@ class LCMConfig:
assert self.max_value > self.min_value, "max_value must be > min_value"
class LCMDataset:
class LCMDataset(ProceduralDataset):
"""Generates Least Common Multiple (LCM) tasks"""
def __init__(self, config: LCMConfig):
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 LCM.