mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
refactor: Inherit NumberSortingDataset from ProceduralDataset
This commit is contained in:
parent
65c52912f8
commit
336fdad55c
1 changed files with 4 additions and 17 deletions
|
|
@ -1,10 +1,11 @@
|
|||
"""Number sorting task generator"""
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from random import Random
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from ..dataset import ProceduralDataset
|
||||
|
||||
|
||||
@dataclass
|
||||
class NumberSortingConfig:
|
||||
|
|
@ -28,27 +29,13 @@ class NumberSortingConfig:
|
|||
assert self.min_value < self.max_value, "max_value must be > min_value"
|
||||
|
||||
|
||||
class NumberSortingDataset:
|
||||
class NumberSortingDataset(ProceduralDataset):
|
||||
"""Generates number sorting tasks"""
|
||||
|
||||
def __init__(self, config: NumberSortingConfig):
|
||||
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 _format_number(self, num: float, decimals: int) -> str:
|
||||
"""Format number with specified decimal places"""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue