open-thought/reasoning-gym
1
0
Fork 0
mirror of https://github.com/open-thought/reasoning-gym.git synced 2026-04-28 17:29:39 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: Add category property to ProceduralDataset to extract category name (#248)

Browse source
This commit is contained in:
Andreas Köpf 2025-03-01 23:11:40 +01:00 • committed by GitHub
parent f549909c3d
commit e71d2a96b6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 9 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

9
reasoning_gym/dataset.py
View file

@ -19,6 +19,15 @@ class ProceduralDataset(ABC, Sized, Iterable[dict[str, Any]]):
self.size = size self.size = size
self.seed = seed if seed is not None else Random().randint(0, 2**32) self.seed = seed if seed is not None else Random().randint(0, 2**32)
@property
def category(self) -> str:
"""Extract category from the module name."""
module_name = self.__class__.__module__
parts = module_name.split(".")
if len(parts) >= 3:
return parts[1] # reasoning_gym.{category}.dataset_name
return "other"
def __len__(self) -> int: def __len__(self) -> int:
"""Return the virtual size of the dataset""" """Return the virtual size of the dataset"""
return self.size return self.size