add hierarchical groups

Author: baberabb

lm_eval/tasks/factory.py

@@ -32,7 +32,25 @@ def build(
         * entry.kind == TASK / PY_TASK -> returns instantiated task object
         * entry.kind == GROUP -> returns (GroupConfig, mapping-of-subtasks)
         * entry.kind == TAG -> returns mapping-of-tasks (tag expansion)
+        * entry with ref_target -> resolves reference and builds target
+        * entry with tag_ref -> expands tag and builds tasks
         """
+        # Handle external references (ref: in children)
+        if entry.ref_target:
+            if entry.ref_target not in registry:
+                raise KeyError(
+                    f"Reference '{entry.ref_target}' not found for '{entry.name}'"
+                )
+            target_entry = registry[entry.ref_target]
+            return self.build(target_entry, overrides=overrides, registry=registry)
+
+        # Handle tag expansion (tag: in children)
+        if entry.tag_ref:
+            if entry.tag_ref not in registry:
+                raise KeyError(f"Tag '{entry.tag_ref}' not found for '{entry.name}'")
+            tag_entry = registry[entry.tag_ref]
+            return self._build_tag(tag_entry, overrides, registry)
+
         if entry.kind is Kind.TAG:
             return self._build_tag(entry, overrides, registry)
 
@@ -44,6 +62,11 @@ def build(
     def _build_task(self, entry: Entry, overrides: dict[str, Any] | None):
         """Build a task and return it wrapped in a dict {task_name: task_obj}."""
         cfg = self._load_full_config(entry, overrides)
+
+        # Remove structural keys that aren't part of task config
+        for key in ("children", "ref", "tag", "group"):
+            cfg.pop(key, None)
+
         # Use cfg["task"] as key (may be overridden, e.g., for namespacing)
         task_name = cfg["task"]
 
@@ -71,25 +94,91 @@ def _build_group(
         group_name = entry.name
 
         children: dict[str, Any] = {}
-        for item in group_obj.config["task"]:
+
+        # Handle new-style children: dict (hierarchical)
+        if "children" in raw_cfg:
+            children.update(
+                self._build_children(
+                    raw_cfg["children"], group_name, overrides, registry
+                )
+            )
+
+        # Handle old-style task: list (backward compatibility)
+        if "task" in grp_cfg and isinstance(grp_cfg["task"], list):
+            children.update(
+                self._build_task_list(grp_cfg["task"], group_name, overrides, registry)
+            )
+
+        return {group_obj: children}
+
+    def _build_children(
+        self,
+        children_cfg: dict[str, Any],
+        group_name: str,
+        overrides: dict[str, Any] | None,
+        registry: Mapping[str, Entry],
+    ) -> dict[str, Any]:
+        """Build children defined via children: dict."""
+        result: dict[str, Any] = {}
+
+        for child_name, child_cfg in children_cfg.items():
+            child_path = f"{group_name}::{child_name}"
+
+            # Look up pre-registered entry from index
+            if child_path in registry:
+                child_entry = registry[child_path]
+                child_overrides = overrides or {}
+
+                # Merge any inline overrides from child_cfg (excluding structural keys)
+                inline_overrides = {
+                    k: v
+                    for k, v in child_cfg.items()
+                    if k not in ("ref", "tag", "children")
+                }
+                if inline_overrides:
+                    child_overrides = {**child_overrides, **inline_overrides}
+
+                child = self.build(
+                    child_entry, overrides=child_overrides, registry=registry
+                )
+                result.update(child)
+            else:
+                # Fallback: inline task not pre-registered (shouldn't normally happen)
+                task_cfg = {**child_cfg, "task": child_path}
+                task_cfg["metadata"] = task_cfg.get("metadata", {}) | self._meta
+                result[child_path] = ConfigurableTask(config=task_cfg)
+
+        return result
+
+    def _build_task_list(
+        self,
+        task_list: list,
+        group_name: str,
+        overrides: dict[str, Any] | None,
+        registry: Mapping[str, Entry],
+    ) -> dict[str, Any]:
+        """Build children defined via task: list (backward compatibility)."""
+        result: dict[str, Any] = {}
+
+        for item in task_list:
             # Step 1: Normalize - extract base_name and item_overrides
             if isinstance(item, str):
                 base_name = item
                 item_overrides = overrides or {}
             elif isinstance(item, dict):
                 base_name = item["task"]
-                item_overrides = item
+                item_overrides = {**overrides, **item}
             else:
                 raise TypeError(
-                    f"Unsupported sub-entry {item!r} in group '{entry.name}'"
+                    f"Unsupported sub-entry {item!r} in group '{group_name}'"
                 )
 
             # Step 2: Handle inline task (not in registry)
             if base_name not in registry:
                 namespaced = f"{group_name}::{base_name}"
                 task_cfg = {**item_overrides, "task": namespaced}
                 task_cfg["metadata"] = task_cfg.get("metadata", {}) | self._meta
-                children[namespaced] = ConfigurableTask(config=task_cfg)
+                result[namespaced] = ConfigurableTask(config=task_cfg)
                 continue
 
             # Step 3: Build based on entry kind
@@ -118,9 +207,9 @@ def _build_group(
                     registry=registry,
                 )
 
-            children.update(child)
+            result.update(child)
 
-        return {group_obj: children}
+        return result
 
     def _build_tag(
         self,
@@ -137,7 +226,11 @@ def _build_tag(
     def _load_full_config(
         self, entry: Entry, overrides: dict[str, Any] | None
     ) -> dict[str, Any]:
-        if entry.yaml_path:
+        # For inline children (have parent), use the stored cfg directly
+        # instead of loading from YAML (which would load the parent's full config)
+        if entry.parent and entry.cfg:
+            cfg = deepcopy(entry.cfg)
+        elif entry.yaml_path:
             cfg = deepcopy(load_yaml(entry.yaml_path, resolve_func=True))
         else:
             cfg: dict[str, Any] = {

lm_eval/tasks/index.py

@@ -29,6 +29,12 @@ class Entry:
     cfg: dict[str, str] | None = None
     tags: set[str] = field(default_factory=set)
     task_list_path: Path | None = None
+    # Hierarchical task support
+    parent: str | None = (
+        None  # parent path for inline children (e.g., "mmlu" for "mmlu::stem")
+    )
+    ref_target: str | None = None  # for children with ref: points to external entry
+    tag_ref: str | None = None  # for children with tag: expands to tagged tasks
 
 
 log = logging.getLogger(__name__)
@@ -85,24 +91,33 @@ def process_cfg(
         cfg: dict[str, Any],
         path: Path,
         index: dict[str, Entry],
+        parent_path: str | None = None,
     ) -> None:
         kind = TaskIndex._kind_of(cfg)
         if kind is Kind.GROUP:
             grp_name = cfg["group"]
-            if grp_name in index:
+            # Build full path for hierarchical addressing
+            full_path = f"{parent_path}::{grp_name}" if parent_path else grp_name
+
+            if full_path in index:
                 log.debug(
-                    f"Duplicate group name '{grp_name}' found. "
-                    f"Already registered from: {index[grp_name].yaml_path}. "
+                    f"Duplicate group name '{full_path}' found. "
+                    f"Already registered from: {index[full_path].yaml_path}. "
                     f"Skipping duplicate from: {path}"
                 )
                 return
-            index[grp_name] = Entry(
-                name=grp_name,
+            index[full_path] = Entry(
+                name=full_path,
                 kind=Kind.GROUP,
                 yaml_path=path,
                 tags=TaskIndex._str_to_set(cfg.get("tag")),
                 cfg=cfg,
+                parent=parent_path,
             )
+
+            # Process inline children if present
+            if "children" in cfg:
+                TaskIndex._process_children(cfg["children"], full_path, path, index)
             return
 
         if kind is Kind.TASK or kind is Kind.PY_TASK:
@@ -204,3 +219,77 @@ def _str_to_set(tags: str | list[str] | None = None) -> set[str]:
             if isinstance(tags, str)
             else set()
         )
+
+    @staticmethod
+    def _process_children(
+        children: dict[str, Any],
+        parent_path: str,
+        yaml_path: Path,
+        index: dict[str, Entry],
+    ) -> None:
+        """Process inline children definitions within a group.
+
+        Children can be:
+        - Inline task: dict with task config fields (dataset_path, doc_to_text, etc.)
+        - Inline subgroup: dict with 'children' key
+        - External ref: dict with 'ref' key pointing to existing entry
+        - Tag expansion: dict with 'tag' key to expand tagged tasks
+        """
+        for child_name, child_cfg in children.items():
+            if not isinstance(child_cfg, dict):
+                log.warning(
+                    f"Invalid child config for '{child_name}' in '{parent_path}': "
+                    f"expected dict, got {type(child_cfg).__name__}"
+                )
+                continue
+
+            child_path = f"{parent_path}::{child_name}"
+
+            if child_path in index:
+                log.debug(f"Duplicate child '{child_path}' found, skipping.")
+                continue
+
+            if "ref" in child_cfg:
+                # External reference - register with ref_target for build-time resolution
+                index[child_path] = Entry(
+                    name=child_path,
+                    kind=Kind.GROUP,  # Assume group, will resolve at build time
+                    yaml_path=yaml_path,
+                    parent=parent_path,
+                    ref_target=child_cfg["ref"],
+                    cfg=child_cfg,
+                    tags=TaskIndex._str_to_set(child_cfg.get("tag")),
+                )
+
+            elif "tag" in child_cfg:
+                # Tag expansion - register with tag_ref for build-time expansion
+                index[child_path] = Entry(
+                    name=child_path,
+                    kind=Kind.TAG,
+                    yaml_path=yaml_path,
+                    parent=parent_path,
+                    tag_ref=child_cfg["tag"],
+                    cfg=child_cfg,
+                    tags=TaskIndex._str_to_set(child_cfg.get("tag")),
+                )
+
+            elif "children" in child_cfg:
+                # Nested inline group - recurse
+                nested_cfg = {**child_cfg, "group": child_name}
+                TaskIndex.process_cfg(
+                    nested_cfg, yaml_path, index, parent_path=parent_path
+                )
+
+            else:
+                # Inline task definition
+                task_cfg = {**child_cfg, "task": child_path}
+                index[child_path] = Entry(
+                    name=child_path,
+                    kind=Kind.TASK,
+                    yaml_path=yaml_path,
+                    parent=parent_path,
+                    cfg=task_cfg,
+                    tags=TaskIndex._str_to_set(child_cfg.get("tag")),
+                )
+                # Register tags for inline tasks
+                TaskIndex._register_tags(child_path, child_cfg.get("tag"), index)