From e71d2a96b6f293dd69e61d12454c90069b96ba52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andreas=20K=C3=B6pf?= <andreas.koepf@xamla.com>
Date: Sat, 1 Mar 2025 23:11:40 +0100
Subject: [PATCH] feat: Add `category` property to `ProceduralDataset` to
 extract category name (#248)

---
 reasoning_gym/dataset.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/reasoning_gym/dataset.py b/reasoning_gym/dataset.py
index 0733ac0b..17bb00e4 100644
--- a/reasoning_gym/dataset.py
+++ b/reasoning_gym/dataset.py
@@ -19,6 +19,15 @@ class ProceduralDataset(ABC, Sized, Iterable[dict[str, Any]]):
         self.size = size
         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:
         """Return the virtual size of the dataset"""
         return self.size