unroll tags in groups; add tests

baberabb · baberabb · commit 65fc38736c38 · 2025-12-09T16:09:30.000+05:00
diff --git a/lm_eval/tasks/factory.py b/lm_eval/tasks/factory.py
@@ -7,20 +7,14 @@
 
 from lm_eval.api.group import ConfigurableGroup, GroupConfig
 from lm_eval.api.task import ConfigurableTask
-from lm_eval.tasks._config_loader import load_yaml as load_cfg
+from lm_eval.tasks._config_loader import load_yaml
 from lm_eval.tasks.index import Entry, Kind
 
 
-load_cfg_cached = load_cfg  # type: ignore[no-redef]
-
-
 class TaskFactory:
     """
     Turns a *Entry* (plus optional overrides) into a
-    *Task* (from task_v3) | *ConfigurableTask* | *GroupConfig* hierarchy.
-
-    For YAML tasks, uses the task_v3.Task builder pattern to automatically
-    select the appropriate Task subclass based on output_type.
+    *Task* | *ConfigurableTask* | *GroupConfig* hierarchy.
     """
 
     def __init__(self, *, meta: dict[str, Any] | None = None):
@@ -35,9 +29,9 @@ def build(
         registry: Mapping[str, Entry],
     ):
         """
-        • 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.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)
         """
         if entry.kind is Kind.TAG:
             return self._build_tag(entry, overrides, registry)
@@ -47,20 +41,22 @@ def build(
 
         return self._build_task(entry, overrides)
 
-    def _build_task(self, entry: Entry, overrides: dict[str, Any] | None) -> dict:
+    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)
+        # Use cfg["task"] as key (may be overridden, e.g., for namespacing)
+        task_name = cfg["task"]
 
         if "class" in cfg:  # PY_TASK route
             cls = cfg["class"]
             obj = cls(config=cfg) if _ctor_accepts_config(cls) else cls()
             if hasattr(obj, "config") and hasattr(obj.config, "task"):
-                obj.config.task = entry.name
-            return {entry.name: obj}
+                obj.config.task = task_name
+            return {task_name: obj}
 
         # Regular YAML task - use ConfigurableTask
         task_obj = ConfigurableTask(config=cfg)
-        return {entry.name: task_obj}
+        return {task_name: task_obj}
 
     def _build_group(
         self,
@@ -71,46 +67,58 @@ def _build_group(
         raw_cfg = self._load_full_config(entry, None)
         grp_cfg = {k: v for k, v in raw_cfg.items() if k in GroupConfig.__annotations__}
         grp_cfg["metadata"] = grp_cfg.get("metadata", {}) | self._meta
-        # Use ConfigurableGroup (hashable) instead of GroupConfig (dict, unhashable)
         group_obj = ConfigurableGroup(config=grp_cfg)
         group_name = entry.name
 
         children: dict[str, Any] = {}
         for item in group_obj.config["task"]:
+            # Step 1: Normalize - extract base_name and item_overrides
             if isinstance(item, str):
-                # Case 1: String reference - look up in registry
                 base_name = item
-                child = self.build(
-                    registry[item],
-                    overrides=overrides,  # group-level overrides propagate
-                    registry=registry,
-                )
+                item_overrides = overrides or {}
             elif isinstance(item, dict):
                 base_name = item["task"]
-                if base_name in registry:
-                    # Case 2: Modify existing indexed task
-                    child = self.build(
-                        registry[base_name],
-                        overrides=item,  # per-item override
-                        registry=registry,
-                    )
-                else:
-                    # Case 3: Create new task inline (not indexed)
-                    task_cfg = {**item}
-                    task_cfg["metadata"] = task_cfg.get("metadata", {}) | self._meta
-                    task_obj = ConfigurableTask(config=task_cfg)
-                    child = {base_name: task_obj}
+                item_overrides = item
             else:
                 raise TypeError(
                     f"Unsupported sub-entry {item!r} in group '{entry.name}'"
                 )
 
-            # Namespace ALL child tasks with group_name::task_name
-            namespaced_child = {}
-            for task_name, task_obj in child.items():
-                namespaced_name = f"{group_name}::{task_name}"
-                namespaced_child[namespaced_name] = task_obj
-            children.update(namespaced_child)
+            # 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)
+                continue
+
+            # Step 3: Build based on entry kind
+            child_entry = registry[base_name]
+
+            if child_entry.kind is Kind.GROUP:
+                child = self.build(
+                    child_entry, overrides=item_overrides, registry=registry
+                )
+            elif child_entry.kind is Kind.TAG:
+                child = {}
+                for task_name in child_entry.tags:
+                    namespaced = f"{group_name}::{task_name}"
+                    child.update(
+                        self.build(
+                            registry[task_name],
+                            overrides={"task": namespaced, **item_overrides},
+                            registry=registry,
+                        )
+                    )
+            else:  # TASK or PY_TASK
+                namespaced = f"{group_name}::{base_name}"
+                child = self.build(
+                    child_entry,
+                    overrides={"task": namespaced, **item_overrides},
+                    registry=registry,
+                )
+
+            children.update(child)
 
         return {group_obj: children}
 
@@ -119,7 +127,7 @@ def _build_tag(
         entry: Entry,
         overrides: dict[str, Any] | None,
         registry: Mapping[str, Entry],
-    ) -> dict:
+    ):
         """Build all tasks in a tag and return merged dict."""
         result = {}
         for name in entry.tags:
@@ -130,7 +138,7 @@ def _load_full_config(
         self, entry: Entry, overrides: dict[str, Any] | None
     ) -> dict[str, Any]:
         if entry.yaml_path:
-            cfg = deepcopy(load_cfg_cached(entry.yaml_path, resolve_func=True))
+            cfg = deepcopy(load_yaml(entry.yaml_path, resolve_func=True))
         else:
             cfg: dict[str, Any] = {
                 "metadata": {"config": "unknown"}
diff --git a/lm_eval/tasks/index.py b/lm_eval/tasks/index.py
@@ -90,7 +90,7 @@ def process_cfg(
         if kind is Kind.GROUP:
             grp_name = cfg["group"]
             if grp_name in index:
-                log.warning(
+                log.debug(
                     f"Duplicate group name '{grp_name}' found. "
                     f"Already registered from: {index[grp_name].yaml_path}. "
                     f"Skipping duplicate from: {path}"
@@ -105,26 +105,7 @@ def process_cfg(
             )
             return
 
-        if kind is Kind.PY_TASK:
-            name = cfg["task"]
-            if name in index:
-                log.warning(
-                    f"Duplicate task name '{name}' found. "
-                    f"Already registered from: {index[name].yaml_path}. "
-                    f"Skipping duplicate from: {path}"
-                )
-                return
-            index[name] = Entry(
-                name=name,
-                kind=Kind.PY_TASK,
-                yaml_path=path,
-                tags=TaskIndex._str_to_set(cfg.get("tag")),
-                cfg=cfg,
-            )
-            TaskIndex._register_tags(name, cfg.get("tag"), index)
-            return
-
-        if kind is Kind.TASK:
+        if kind is Kind.TASK or kind is Kind.PY_TASK:
             name = cfg["task"]
             if name in index:
                 log.warning(
diff --git a/tests/test_configs/tag_parent_group.yaml b/tests/test_configs/tag_parent_group.yaml
@@ -0,0 +1,5 @@
+# Parent group containing a subgroup - simulates mmlu pattern
+# Structure: tag_parent_group -> tag_subgroup -> test_tag_tasks (TAG) -> tag_task_1, tag_task_2, tag_task_3
+group: tag_parent_group
+task:
+  - tag_subgroup
diff --git a/tests/test_configs/tag_subgroup.yaml b/tests/test_configs/tag_subgroup.yaml
@@ -0,0 +1,5 @@
+# Subgroup that references a TAG - simulates mmlu_humanities pattern
+# This group contains a tag reference (test_tag_tasks) which expands to multiple tasks
+group: tag_subgroup
+task:
+  - test_tag_tasks
diff --git a/tests/test_configs/tag_task_1.yaml b/tests/test_configs/tag_task_1.yaml
@@ -0,0 +1,16 @@
+# Task 1 with tag - simulates mmlu_formal_logic pattern
+task: tag_task_1
+tag: test_tag_tasks
+dataset_path: json
+dataset_kwargs:
+  data_files:
+    test: tests/test_configs/test_data.json
+output_type: multiple_choice
+doc_to_text: "{{question}}"
+doc_to_target: "{{choices[answer]}}"
+test_split: test
+num_fewshot: 0
+metric_list:
+  - metric: acc
+    aggregation: mean
+    higher_is_better: true
diff --git a/tests/test_configs/tag_task_2.yaml b/tests/test_configs/tag_task_2.yaml
@@ -0,0 +1,16 @@
+# Task 2 with tag - simulates mmlu_high_school_history pattern
+task: tag_task_2
+tag: test_tag_tasks
+dataset_path: json
+dataset_kwargs:
+  data_files:
+    test: tests/test_configs/test_data.json
+output_type: multiple_choice
+doc_to_text: "{{question}}"
+doc_to_target: "{{choices[answer]}}"
+test_split: test
+num_fewshot: 0
+metric_list:
+  - metric: acc
+    aggregation: mean
+    higher_is_better: true
diff --git a/tests/test_configs/tag_task_3.yaml b/tests/test_configs/tag_task_3.yaml
@@ -0,0 +1,16 @@
+# Task 3 with tag - simulates mmlu_philosophy pattern
+task: tag_task_3
+tag: test_tag_tasks
+dataset_path: json
+dataset_kwargs:
+  data_files:
+    test: tests/test_configs/test_data.json
+output_type: multiple_choice
+doc_to_text: "{{question}}"
+doc_to_target: "{{choices[answer]}}"
+test_split: test
+num_fewshot: 0
+metric_list:
+  - metric: acc
+    aggregation: mean
+    higher_is_better: true
diff --git a/tests/test_task_manager.py b/tests/test_task_manager.py
@@ -745,3 +745,64 @@ def test_include_defaults_true_with_new_tasks(self, shared_task_manager):
             assert len(task_manager.all_tasks) > len(shared_task_manager.all_tasks), (
                 "Should have more tasks when including custom path"
             )
+
+    def test_tag_expansion_in_group(self, test_configs_task_manager):
+        """Test that TAGs inside groups are expanded and each task is namespaced individually.
+
+        This tests the MMLU-like structure: GROUP -> TAG -> multiple tasks
+        Without proper TAG handling, all tasks in the tag get the same namespaced name
+        and collide, leaving only one task.
+        """
+        # Load the subgroup that contains a TAG reference
+        result = test_configs_task_manager.load_task_or_group(["tag_subgroup"])
+
+        # Get the children dict from the group
+        group_key = list(result.keys())[0]
+        children = result[group_key]
+
+        # All 3 tasks from the tag should be expanded and namespaced
+        assert "tag_subgroup::tag_task_1" in children, (
+            "tag_task_1 should be namespaced under tag_subgroup"
+        )
+        assert "tag_subgroup::tag_task_2" in children, (
+            "tag_task_2 should be namespaced under tag_subgroup"
+        )
+        assert "tag_subgroup::tag_task_3" in children, (
+            "tag_task_3 should be namespaced under tag_subgroup"
+        )
+
+        # Verify we have exactly 3 tasks (not 1 due to collision)
+        assert len(children) == 3, (
+            f"Should have 3 tasks from TAG expansion, got {len(children)}"
+        )
+
+    def test_nested_group_with_tag(self, test_configs_task_manager):
+        """Test nested groups with TAG: parent_group -> subgroup -> TAG -> tasks.
+
+        This simulates the full MMLU structure where:
+        - mmlu (GROUP) contains mmlu_humanities (GROUP)
+        - mmlu_humanities contains mmlu_humanities_tasks (TAG)
+        - The TAG expands to individual tasks
+        """
+        # Load the parent group
+        result = test_configs_task_manager.load_task_or_group(["tag_parent_group"])
+
+        # Navigate the nested structure
+        parent_key = list(result.keys())[0]
+        parent_children = result[parent_key]
+
+        # Should contain the subgroup
+        assert len(parent_children) == 1, "Parent should have 1 child (the subgroup)"
+
+        # Get the subgroup
+        subgroup_key = list(parent_children.keys())[0]
+        subgroup_children = parent_children[subgroup_key]
+
+        # The subgroup should have all 3 tasks expanded from the TAG
+        # Tasks are namespaced under their immediate parent group (tag_subgroup)
+        assert "tag_subgroup::tag_task_1" in subgroup_children
+        assert "tag_subgroup::tag_task_2" in subgroup_children
+        assert "tag_subgroup::tag_task_3" in subgroup_children
+        assert len(subgroup_children) == 3, (
+            f"Subgroup should have 3 tasks, got {len(subgroup_children)}"
+        )
