open-thought/reasoning-gym
1
0
Fork 0
mirror of https://github.com/open-thought/reasoning-gym.git synced 2026-04-19 12:58:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: Add dataset factory with registration and creation functions

Browse source
This commit is contained in:
Andreas Koepf (aider) 2025-01-25 00:00:22 +01:00
parent a07df085cf
commit 2befe97151
2 changed files with 78 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

5
reasoning_gym/dataset.py
View file

@ -9,8 +9,9 @@ from typing import Any, Dict, Iterator, Optional
class ProceduralDataset(ABC, Sized, Iterable[Dict[str, Any]]):
"""Abstract base class for procedural dataset generators"""
def __init__(self, seed: Optional[int] = None, size: int = 500):
"""Initialize the dataset with optional seed and size"""
def __init__(self, config: Any, seed: Optional[int] = None, size: int = 500):
"""Initialize the dataset with config, optional seed and size"""
self.config = config
self.size = size
self.seed = seed if seed is not None else Random().randint(0, 2**32)