__init__.py 37.1 KB
Newer Older
1
import collections
2
3
import functools
import importlib.util
4
import inspect
5
import logging
6
import sys
7
from functools import partial
8
from glob import iglob
Baber's avatar
Baber committed
9
from pathlib import Path
Baber's avatar
nit  
Baber committed
10
11
12
13
14
15
16
17
18
19
20
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Generator,
    List,
    Mapping,
    Optional,
    Union,
)
21
22
23

import yaml
from yaml import YAMLError
&'s avatar
& committed
24

Lintang Sutawika's avatar
Lintang Sutawika committed
25
26
from lm_eval.api.group import ConfigurableGroup, GroupConfig
from lm_eval.evaluator_utils import get_subtask_list
Baber's avatar
nit  
Baber committed
27
28
29
30
31
from lm_eval.utils import pattern_match, setup_logging


if TYPE_CHECKING:
    from lm_eval.api.task import ConfigurableTask, Task
Lintang Sutawika's avatar
Lintang Sutawika committed
32
33
34


GROUP_ONLY_KEYS = list(GroupConfig().to_dict().keys())
35

lintangsutawika's avatar
lintangsutawika committed
36

Lintang Sutawika's avatar
Lintang Sutawika committed
37
eval_logger = logging.getLogger(__name__)
38
39
40
_Base = yaml.CLoader if getattr(yaml, "__with_libyaml__", False) else yaml.FullLoader


Baber's avatar
nit  
Baber committed
41
@functools.lru_cache(maxsize=None)  # ← reuse per (directory, simple) pair
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def _make_loader(yaml_dir: Path, simple: bool = False) -> type[yaml.Loader]:
    """
    Return a custom YAML Loader class bound to *yaml_dir*.

    yaml_dir
        Directory that holds the YAML file being parsed.
        We capture it so that !function look-ups can resolve relative
        Python files like  my_utils.some_fn  ➜  yaml_dir / "my_utils.py".
    simple
        If True we ignore !function completely (used by `mode="simple"`).
    """

    class Loader(_Base):
        """Dynamically-generated loader that knows its base directory."""

        # no extra state needed; the constructor stays the same

    # Register (or stub) the !function constructor **for this Loader only**
    if simple:
Baber's avatar
nit  
Baber committed
61
        yaml.add_constructor("!function", ignore_constructor, Loader=Loader)
62
63
64
65
66
67
68
69
70
71
72
73
74
75
    else:
        yaml.add_constructor(
            "!function",
            # capture yaml_dir once so the lambda is fast and pickle-able
            lambda ld, node, _dir=yaml_dir: _import_function(
                ld.construct_scalar(node),
                base_path=_dir,
            ),
            Loader=Loader,
        )

    return Loader


Baber's avatar
nit  
Baber committed
76
@functools.lru_cache(maxsize=None)  # ← cache module objects
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def _import_function(qualname: str, *, base_path: Path) -> Callable:
    mod_path, _, func_name = qualname.rpartition(".")
    if not mod_path:
        raise ValueError(f"{qualname!r} has no module part")
    file_path = base_path / f"{mod_path.replace('.', '/')}.py"
    module_name = f"_yaml_dynamic.{hash(file_path)}_{file_path.stem}"
    if module_name in sys.modules:
        mod = sys.modules[module_name]
    else:
        spec = importlib.util.spec_from_file_location(module_name, file_path)
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        sys.modules[module_name] = mod
    return getattr(mod, func_name)


def ignore_constructor(loader: yaml.Loader, node: yaml.Node) -> None:
    return None


Baber's avatar
nit  
Baber committed
97
@functools.lru_cache(maxsize=None)  #
98
99
100
101
102
103
def _parse_yaml_file(path: Path, mode: str) -> dict:
    loader_cls = _make_loader(path.parent, simple=(mode == "simple"))
    with path.open("rb") as fh:
        return yaml.load(fh, Loader=loader_cls)


Baber's avatar
nit  
Baber committed
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
@functools.lru_cache(maxsize=None)
def _get_cached_config(yaml_path: Path, mode: str) -> dict:
    """Load and cache resolved YAML configs with LRU eviction."""
    # Parse the YAML file
    yaml_config = _parse_yaml_file(yaml_path, mode)
    yaml_dir = yaml_path.parent

    # Handle includes
    include = yaml_config.pop("include", None)
    if not include:
        return yaml_config

    include_paths = include if isinstance(include, list) else [include]
    final_cfg: dict = {}

    for inc in reversed(include_paths):
        if inc is None:
            continue
        inc_path = Path(inc)
        if not inc_path.is_absolute():
            inc_path = (yaml_dir / inc_path).resolve()
        # Recursive call will use the cache
        included = _get_cached_config(inc_path, mode)
        final_cfg.update(included)

    final_cfg.update(yaml_config)  # local keys win
    return final_cfg


133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
def load_yaml_config(
    yaml_path: Union[Path, str, None] = None,
    yaml_config: dict | None = None,
    yaml_dir: Path | None = None,
    mode: str = "full",
    *,
    _seen: set[tuple[Path, str]] | None = None,
    resolve_includes: bool = True,
) -> dict:
    """
    Parse a YAML config with optional include handling.

    Parameters
    ----------
    yaml_path
        Path to the main YAML file.  Needed unless *yaml_config* is
        supplied directly (e.g. by tests).
    yaml_config
        Pre-parsed dict to use instead of reading *yaml_path*.
    yaml_dir
        Base directory for resolving relative include paths.  Defaults
        to `yaml_path.parent`.
    mode
        "full"  – honour  !function  tags
        "simple" – ignore !function  (faster).
    _seen
        **Internal** recursion set: tuples of (absolute-path, mode).
        Prevents include cycles such as  A → B → A.
    """
    if yaml_config is None and yaml_path is None:
        raise ValueError("load_yaml_config needs either yaml_path or yaml_config")

    # ------------------------------------------------------------------ cycle guard
    if _seen is None:
        _seen = set()
    if yaml_path is not None:
        yaml_path = Path(yaml_path).expanduser().resolve()

Baber's avatar
nit  
Baber committed
171
172
173
        # ---------- fast-path: use LRU cached function ----------
        if yaml_config is None and resolve_includes:
            return _get_cached_config(yaml_path, mode)
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215

        key = (yaml_path.resolve(), mode)
        if key in _seen:
            raise ValueError(f"Include cycle detected at {yaml_path}")
        _seen.add(key)

    # ------------------------------------------------------------------ load / parse
    if yaml_config is None:  # ordinary path-based load
        yaml_config = _parse_yaml_file(yaml_path, mode)

    if yaml_dir is None and yaml_path is not None:
        yaml_dir = yaml_path.parent
    assert yaml_dir is not None, "yaml_dir must be set by caller or deduced from path"

    # ------------------------------------------------------------------ handle include
    include = yaml_config.pop("include", None)
    if not include and not resolve_includes:
        return yaml_config

    include_paths = include if isinstance(include, list) else [include]
    final_cfg: dict = {}

    for inc in reversed(include_paths):
        if inc is None:  # guard against explicit nulls
            continue
        inc_path = Path(inc)
        if not inc_path.is_absolute():
            inc_path = (yaml_dir / inc_path).resolve()
        included = load_yaml_config(
            yaml_path=inc_path,
            mode=mode,
            yaml_dir=inc_path.parent,
            _seen=_seen,  # <-- pass set downward
        )
        final_cfg.update(included)

    final_cfg.update(yaml_config)  # local keys win
    return final_cfg


def iter_yaml_files(root: Path) -> Generator[Path, Any, None]:
    # '**/*.yaml' is handled internally by os.scandir.
Baber's avatar
nit  
Baber committed
216
217
218
    for p in iglob("**/*.yaml", root_dir=root, recursive=True):
        # ignore check
        if p.startswith(("__pycache__", ".ipynb_checkpoints")):
219
            continue
Baber's avatar
nit  
Baber committed
220
        yield root / p
Lintang Sutawika's avatar
Lintang Sutawika committed
221

222

223
224
225
class TaskManager:
    """TaskManager indexes all tasks from the default `lm_eval/tasks/`
    and an optional directory if provided.
226

227
228
    """

229
230
    def __init__(
        self,
Lintang Sutawika's avatar
Lintang Sutawika committed
231
        verbosity: Optional[str] = None,
Baber's avatar
Baber committed
232
        include_path: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
233
        include_defaults: bool = True,
Baber Abbasi's avatar
Baber Abbasi committed
234
        metadata: Optional[dict] = None,
235
    ) -> None:
Lintang Sutawika's avatar
Lintang Sutawika committed
236
        if verbosity is not None:
Baber's avatar
nit  
Baber committed
237
            setup_logging(verbosity)
238
        self.include_path = include_path
Baber Abbasi's avatar
Baber Abbasi committed
239
        self.metadata = metadata
240
241
242
        self._task_index = self.initialize_tasks(
            include_path=include_path, include_defaults=include_defaults
        )
243
        self._all_tasks = sorted(list(self._task_index.keys()))
244

245
246
247
248
        self._all_groups = sorted(
            [x for x in self._all_tasks if self._task_index[x]["type"] == "group"]
        )
        self._all_subtasks = sorted(
249
250
251
252
253
            [
                x
                for x in self._all_tasks
                if self._task_index[x]["type"] in ["task", "python_task"]
            ]
254
255
256
257
258
        )
        self._all_tags = sorted(
            [x for x in self._all_tasks if self._task_index[x]["type"] == "tag"]
        )

259
        self.task_group_map = collections.defaultdict(list)
260

261
262
    def initialize_tasks(
        self,
Baber's avatar
Baber committed
263
        include_path: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
264
        include_defaults: bool = True,
Baber Abbasi's avatar
Baber Abbasi committed
265
266
    ) -> dict[str, dict]:
        """Creates a dictionary of tasks indexes.
267

268
269
270
271
272
        :param include_path: Union[str, List] = None
            An additional path to be searched for tasks recursively.
            Can provide more than one such path as a list.
        :param include_defaults: bool = True
            If set to false, default tasks (those in lm_eval/tasks/) are not indexed.
Baber Abbasi's avatar
Baber Abbasi committed
273
        return
274
275
            Dictionary of task names as key and task metadata
        """
276
        if include_defaults:
Baber's avatar
Baber committed
277
            all_paths = [Path(__file__).parent]
278
279
        else:
            all_paths = []
280
        if include_path is not None:
Baber's avatar
Baber committed
281
            if isinstance(include_path, (str, Path)):
282
                include_path = [include_path]
Baber's avatar
Baber committed
283
284
            # Convert all paths to Path objects
            all_paths.extend(Path(p) for p in include_path)
285

286
287
288
289
        task_index = {}
        for task_dir in all_paths:
            tasks = self._get_task_and_group(task_dir)
            task_index = {**tasks, **task_index}
lintangsutawika's avatar
format  
lintangsutawika committed
290

291
292
293
        return task_index

    @property
Baber's avatar
Baber committed
294
    def all_tasks(self) -> List[str]:
295
296
        return self._all_tasks

297
    @property
Baber's avatar
Baber committed
298
    def all_groups(self) -> List[str]:
299
300
301
        return self._all_groups

    @property
Baber's avatar
Baber committed
302
    def all_subtasks(self) -> List[str]:
303
304
305
        return self._all_subtasks

    @property
Baber's avatar
Baber committed
306
    def all_tags(self) -> List[str]:
307
308
        return self._all_tags

309
    @property
Baber's avatar
Baber committed
310
    def task_index(self) -> Dict[str, Dict[str, Union[str, int, List[str]]]]:
311
312
        return self._task_index

313
    def list_all_tasks(
Baber's avatar
Baber committed
314
315
316
317
        self,
        list_groups: bool = True,
        list_tags: bool = True,
        list_subtasks: bool = True,
318
    ) -> str:
319
320
321
322
323
        """
        Return a Markdown table (as a string) listing groups, tags and/or subtasks
        known to this TaskManager.  Safe for configs whose yaml_path is -1 and for
        task configs whose `include:` is a list.
        """
324
325
        from pytablewriter import MarkdownTableWriter

326
        # ------------------------------------------------------------------ helpers
Baber's avatar
Baber committed
327
        def sanitize_path(path: str) -> str:
328
329
            # print a relative path for anything inside lm_eval/tasks/
            # path_str = str(path)
330
331
            if "lm_eval/tasks/" in path:
                return "lm_eval/tasks/" + path.split("lm_eval/tasks/")[-1]
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
            return path

        def first_output_type_from_includes(cfg: dict, base: Path) -> str:
            """Walk cfg['include'] (string or list) and return the first
            include that itself specifies an output_type."""
            inc_raw = cfg.get("include")
            if not inc_raw:
                return ""

            inc_list = inc_raw if isinstance(inc_raw, list) else [inc_raw]
            for inc in inc_list:
                inc_path = Path(inc)
                if not inc_path.is_absolute():  # treat as relative include
                    inc_path = base.parent / inc_path
                try:
                    inc_cfg = load_yaml_config(inc_path, mode="simple")
                except FileNotFoundError:
                    continue
                if "output_type" in inc_cfg:
                    return inc_cfg["output_type"]
            return ""

        # -------------------------------------------------------------- GROUP table
355
356
        group_table = MarkdownTableWriter()
        group_table.headers = ["Group", "Config Location"]
357
358
359
360
361
362
363
364
365
        group_table.value_matrix = [
            [
                g,
                "---"
                if self.task_index[g]["yaml_path"] == -1
                else sanitize_path(self.task_index[g]["yaml_path"]),
            ]
            for g in self.all_groups
        ]
366

367
        # ---------------------------------------------------------------- TAG table
368
369
370
371
        tag_table = MarkdownTableWriter()
        tag_table.headers = ["Tag"]
        tag_table.value_matrix = [[t] for t in self.all_tags]

372
        # ------------------------------------------------------------ SUBTASK table
373
374
        subtask_table = MarkdownTableWriter()
        subtask_table.headers = ["Task", "Config Location", "Output Type"]
375
376
        st_values: list[list[str]] = []

377
        for t in self.all_subtasks:
378
379
380
381
382
383
            raw_path = self.task_index[t]["yaml_path"]

            if raw_path == -1:
                # python-only task or generated at runtime
                display_path = "---"
                output_type = ""
384
            else:
385
386
387
388
389
390
391
392
393
394
395
396
                path_obj = Path(raw_path)
                display_path = sanitize_path(str(path_obj))

                # load minimal YAML to discover output_type
                cfg = load_yaml_config(path_obj, mode="simple")
                if "output_type" in cfg:
                    output_type = cfg["output_type"]
                else:
                    output_type = first_output_type_from_includes(cfg, path_obj)

            st_values.append([t, display_path, output_type])

397
398
        subtask_table.value_matrix = st_values

399
400
        # ------------------------------------------------------------- final string
        parts: list[str] = ["\n"]
401
        if list_groups:
402
403
            parts.append(group_table.dumps())
            parts.append("\n")
404
        if list_tags:
405
406
            parts.append(tag_table.dumps())
            parts.append("\n")
407
        if list_subtasks:
408
409
410
411
            parts.append(subtask_table.dumps())
            parts.append("\n")

        return "".join(parts)
412

Baber Abbasi's avatar
Baber Abbasi committed
413
    def match_tasks(self, task_list: list[str]) -> list[str]:
Baber's avatar
nit  
Baber committed
414
        return pattern_match(task_list, self.all_tasks)
415

Baber Abbasi's avatar
Baber Abbasi committed
416
    def _name_is_registered(self, name: str) -> bool:
417
        return name in self.all_tasks
418

Baber Abbasi's avatar
Baber Abbasi committed
419
    def _name_is_task(self, name: str) -> bool:
420
421
422
        return (
            self._name_is_registered(name) and self.task_index[name]["type"] == "task"
        )
Lintang Sutawika's avatar
Lintang Sutawika committed
423

Baber Abbasi's avatar
Baber Abbasi committed
424
    def _name_is_tag(self, name: str) -> bool:
425
        return self._name_is_registered(name) and self.task_index[name]["type"] == "tag"
426

Baber Abbasi's avatar
Baber Abbasi committed
427
    def _name_is_group(self, name: str) -> bool:
428
429
430
        return (
            self._name_is_registered(name) and self.task_index[name]["type"] == "group"
        )
431

Baber Abbasi's avatar
Baber Abbasi committed
432
    def _name_is_python_task(self, name: str) -> bool:
433
434
435
436
        return (
            self._name_is_registered(name)
            and self.task_index[name]["type"] == "python_task"
        )
437

Baber Abbasi's avatar
Baber Abbasi committed
438
    def _config_is_task(self, config: dict) -> bool:
439
        return "task" in config and isinstance(config["task"], str)
440

Baber Abbasi's avatar
Baber Abbasi committed
441
    def _config_is_group(self, config: dict) -> bool:
442
        return "task" in config and isinstance(config["task"], list)
443

Baber Abbasi's avatar
Baber Abbasi committed
444
    def _config_is_python_task(self, config: dict) -> bool:
445
446
447
448
        return "class" in config

    def _config_is_task_list(self, config: dict) -> bool:
        return "task_list" in config and isinstance(config["task_list"], list)
449

Baber's avatar
Baber committed
450
    def _get_yaml_path(self, name: str) -> Union[str, int]:
451
452
        if name not in self.task_index:
            raise ValueError
453
454
        return self.task_index[name]["yaml_path"]

Baber's avatar
Baber committed
455
    def _get_config(self, name: str) -> Dict:
456
457
        if name not in self.task_index:
            raise ValueError
458
459
460
461
        yaml_path = self._get_yaml_path(name)
        if yaml_path == -1:
            return {}
        else:
462
            return load_yaml_config(Path(yaml_path), mode="full")
463

Baber's avatar
Baber committed
464
    def _get_tasklist(self, name: str) -> Union[List[str], int]:
465
466
        if self._name_is_task(name):
            raise ValueError
467
468
        return self.task_index[name]["task"]

469
470
471
472
473
    def _register_task(
        self,
        task_name: str,
        task_type: str,
        yaml_path: str,
Baber's avatar
Baber committed
474
475
476
477
        tasks_and_groups: Dict[str, Dict],
        config: Optional[Dict] = None,
        populate_tags_fn: Optional[callable] = None,
    ) -> None:
478
479
480
481
482
483
484
        """Helper method to register a task in the tasks_and_groups dict"""
        tasks_and_groups[task_name] = {
            "type": task_type,
            "yaml_path": yaml_path,
        }
        # Only populate tags for configs that support it (not groups)
        if config and task_type != "group" and populate_tags_fn:
485
            populate_tags_fn(config, task_name, tasks_and_groups)
486
487

    def _merge_task_configs(
Baber's avatar
Baber committed
488
489
        self, base_config: Dict, task_specific_config: Dict, task_name: str
    ) -> Dict:
490
491
492
493
494
495
496
        """Merge base config with task-specific overrides for task_list configs"""
        if task_specific_config:
            task_specific_config = task_specific_config.copy()
            task_specific_config.pop("task", None)
            return {**base_config, **task_specific_config, "task": task_name}
        return {**base_config, "task": task_name}

Baber's avatar
Baber committed
497
498
499
    def _process_tag_subtasks(
        self, tag_name: str, update_config: Optional[Dict] = None
    ) -> Dict:
500
501
502
503
504
505
506
507
        """Process subtasks for a tag and return loaded tasks"""
        subtask_list = self._get_tasklist(tag_name)
        fn = partial(
            self._load_individual_task_or_group,
            update_config=update_config,
        )
        return dict(collections.ChainMap(*map(fn, reversed(subtask_list))))

Baber's avatar
Baber committed
508
    def _process_alias(self, config: Dict, group: Optional[str] = None) -> Dict:
509
510
511
512
513
514
515
516
        # If the group is not the same as the original
        # group which the group alias was intended for,
        # Set the group_alias to None instead.
        if ("group_alias" in config) and ("group" in config) and group is not None:
            if config["group"] != group:
                config["group_alias"] = None
        return config

Baber's avatar
Baber committed
517
    def _class_has_config_in_constructor(self, cls) -> bool:
518
519
520
521
522
523
524
        constructor = getattr(cls, "__init__", None)
        return (
            "config" in inspect.signature(constructor).parameters
            if constructor
            else False
        )

525
    def _load_individual_task_or_group(
526
        self,
Baber's avatar
Baber committed
527
        name_or_config: Optional[Union[str, Dict]] = None,
528
        parent_name: Optional[str] = None,
Baber's avatar
Baber committed
529
        update_config: Optional[Dict] = None,
530
    ) -> Mapping:
Baber's avatar
nit  
Baber committed
531
532
        from lm_eval.api.task import ConfigurableTask, Task

Baber's avatar
Baber committed
533
534
        def _load_task(
            config: Dict, task: str, yaml_path: Optional[str] = None
Baber's avatar
nit  
Baber committed
535
        ) -> Dict[str, Union["ConfigurableTask", "Task"]]:
536
            if "include" in config:
537
538
539
                # Store the task name to preserve it after include processing
                original_task_name = config.get("task", task)

540
                config = {
541
542
                    **load_yaml_config(
                        yaml_path=Path(yaml_path),
543
                        yaml_config={"include": config.pop("include")},
544
                        mode="full" if yaml_path else "simple",
545
546
                    ),
                    **config,
547
                    "task": original_task_name,
548
                }
549
550
551
552

                # Ensure the task name from the group config is preserved
                # This prevents tasks with the same include from being treated as duplicates

553
            if self._config_is_python_task(config):
554
555
556
557
558
559
                if self._class_has_config_in_constructor(config["class"]):
                    task_object = config["class"](config=config)
                else:
                    task_object = config["class"]()
                if isinstance(task_object, ConfigurableTask):
                    # very scuffed: set task name here. TODO: fixme?
560
                    task_object.config.task = task
561
            else:
Baber Abbasi's avatar
Baber Abbasi committed
562
563
564
565
                if self.metadata is not None:
                    config["metadata"] = config.get("metadata", {}) | self.metadata
                else:
                    config["metadata"] = config.get("metadata", {})
566
                task_object = ConfigurableTask(config=config)
Lintang Sutawika's avatar
Lintang Sutawika committed
567

568
569
            return {task: task_object}

Baber Abbasi's avatar
Baber Abbasi committed
570
        def _get_group_and_subtask_from_config(
Baber's avatar
Baber committed
571
572
            config: Dict,
        ) -> tuple[ConfigurableGroup, List[str]]:
Baber Abbasi's avatar
Baber Abbasi committed
573
574
            if self.metadata is not None:
                config["metadata"] = config.get("metadata", {}) | self.metadata
Lintang Sutawika's avatar
Lintang Sutawika committed
575
576
577
578
579
580
581
582
583
            group_name = ConfigurableGroup(config=config)
            subtask_list = []
            for task in group_name.config["task"]:
                if isinstance(task, str) and self._name_is_tag(task):
                    subtask_list.extend(self._get_tasklist(task))
                else:
                    subtask_list.append(task)
            return group_name, subtask_list

Baber Abbasi's avatar
Baber Abbasi committed
584
        def _process_group_config(
Baber's avatar
Baber committed
585
586
            config: Dict, update_config: Optional[Dict] = None
        ) -> tuple[Dict, Optional[Dict]]:
Lintang Sutawika's avatar
Lintang Sutawika committed
587
588
589
590
591
592
593
594
595
596
597
            if update_config is not None:
                config = {**config, **update_config}
            _update_config = {
                k: v for k, v in config.items() if k not in GROUP_ONLY_KEYS
            }
            if not bool(_update_config):
                _update_config = None

            group_config = {k: v for k, v in config.items() if k in GROUP_ONLY_KEYS}
            return group_config, _update_config

598
599
600
601
        if isinstance(name_or_config, str):
            if update_config is not None:
                # Process name_or_config as a dict instead
                name_or_config = {"task": name_or_config, **update_config}
Lintang Sutawika's avatar
Lintang Sutawika committed
602
603
604
            elif self._name_is_task(name_or_config) or self._name_is_python_task(
                name_or_config
            ):
605
606
                # Get the yaml_path for this task
                yaml_path = self._get_yaml_path(name_or_config)
607
                task_config = self._get_config(name_or_config)
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636

                # Handle task_list configs
                if "task_list" in task_config:
                    # Find the specific task entry
                    task_specific_config = None
                    for task_entry in task_config["task_list"]:
                        if (
                            isinstance(task_entry, dict)
                            and task_entry.get("task") == name_or_config
                        ):
                            task_specific_config = task_entry
                            break

                    if task_specific_config:
                        # Create base config without task_list
                        base_config = {
                            k: v for k, v in task_config.items() if k != "task_list"
                        }
                        # Merge using helper method
                        task_config = self._merge_task_configs(
                            base_config, task_specific_config, name_or_config
                        )
                    else:
                        # Task not found in task_list, shouldn't happen if indexing worked correctly
                        eval_logger.warning(
                            f"Task {name_or_config} not found in task_list"
                        )
                        task_config = {"task": name_or_config}

637
                return _load_task(task_config, task=name_or_config, yaml_path=yaml_path)
638
            else:
639
640
641
                subtask_list = self._get_tasklist(name_or_config)
                if subtask_list == -1:
                    group_config = self._get_config(name_or_config)
Lintang Sutawika's avatar
Lintang Sutawika committed
642
643
644
645
646
647
                    group_config, update_config = _process_group_config(group_config)
                    group_name, subtask_list = _get_group_and_subtask_from_config(
                        group_config
                    )
                else:
                    if self._name_is_tag(name_or_config):
648
649
650
                        return self._process_tag_subtasks(
                            name_or_config,
                            name_or_config
Lintang Sutawika's avatar
Lintang Sutawika committed
651
652
653
654
655
656
657
                            if isinstance(name_or_config, dict)
                            else None,
                        )
                    else:
                        group_name = ConfigurableGroup(
                            config={"group": name_or_config, "task": subtask_list}
                        )
658

659
660
        if isinstance(name_or_config, dict):
            if self._config_is_task(name_or_config):
Lintang Sutawika's avatar
Lintang Sutawika committed
661
662
663
                name = name_or_config.pop("task")
                if update_config is not None:
                    name_or_config = {**name_or_config, **update_config}
664
665
                # If the name is registered as a group
                if self._name_is_group(name):
Lintang Sutawika's avatar
Lintang Sutawika committed
666
667
668
669
670
671
672
673
674
                    group_config = self._get_config(name)

                    group_config, update_config = _process_group_config(
                        group_config, name_or_config
                    )
                    group_name, subtask_list = _get_group_and_subtask_from_config(
                        group_config
                    )
                elif self._name_is_tag(name):
675
                    return self._process_tag_subtasks(name, name_or_config)
676
                else:
677
                    yaml_path = None
678
                    if self._name_is_registered(name):
679
                        yaml_path = self._get_yaml_path(name)
680
681
682
683
                        base_task_config = self._get_config(name)

                        # Check if this is a duplicate.
                        if parent_name is not None:
684
685
686
687
688
689
690
691
                            num_duplicate = len(
                                list(
                                    filter(
                                        lambda x: x.startswith(name),
                                        self.task_group_map[parent_name],
                                    )
                                )
                            )
692
693
694
695
                            if num_duplicate > 0:
                                name = f"{name}-{num_duplicate}"
                            self.task_group_map[parent_name].append(name)

696
697
698
699
                        task_config = {
                            **base_task_config,
                            **name_or_config,
                        }
700
701
                    else:
                        task_config = name_or_config
702
                    return _load_task(task_config, task=name, yaml_path=yaml_path)
703
            else:
Lintang Sutawika's avatar
Lintang Sutawika committed
704
705
706
707
                group_config, update_config = _process_group_config(name_or_config)
                group_name, subtask_list = _get_group_and_subtask_from_config(
                    group_config
                )
708

709
710
711
712
713
        fn = partial(
            self._load_individual_task_or_group,
            parent_name=group_name,
            update_config=update_config,
        )
Lintang Sutawika's avatar
Lintang Sutawika committed
714
715
        return {
            group_name: dict(collections.ChainMap(*map(fn, reversed(subtask_list))))
716
        }
717

Baber's avatar
Baber committed
718
719
720
    def load_task_or_group(
        self, task_list: Optional[Union[str, List[str]]] = None
    ) -> Dict:
721
        """Loads a dictionary of task objects from a list
722

723
724
        :param task_list: Union[str, list] = None
            Single string or list of string of task names to be loaded
725

726
727
728
729
730
        :return
            Dictionary of task objects
        """
        if isinstance(task_list, str):
            task_list = [task_list]
731

732
        all_loaded_tasks = dict(
Baber Abbasi's avatar
Baber Abbasi committed
733
734
735
736
737
738
            collections.ChainMap(
                *map(
                    lambda task: self._load_individual_task_or_group(task),
                    task_list,
                )
            )
739
740
741
        )
        return all_loaded_tasks

Baber's avatar
Baber committed
742
    def load_config(self, config: Dict) -> Mapping:
743
744
        return self._load_individual_task_or_group(config)

Baber's avatar
Baber committed
745
    def _get_task_and_group(self, task_dir: Union[str, Path]) -> Dict[str, Dict]:
746
        """Creates a dictionary of tasks index with the following metadata,
Lintang Sutawika's avatar
Lintang Sutawika committed
747
        - `type`, that can be either `task`, `python_task`, `group` or `tags`.
748
749
            `task` refer to regular task configs, `python_task` are special
            yaml files that only consists of `task` and `class` parameters.
Lintang Sutawika's avatar
Lintang Sutawika committed
750
751
            `group` are group configs. `tags` are labels that can be assigned
            to tasks to assist in sorting and calling tasks of certain themes.
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
        - `yaml_path`, path to the yaml file. If the entry is a `group` that
            was configured through a task config, the yaml_path will be -1
            and all subtasks will be listed in `task` (see below)
        - `task`, reserved for entries with `type` as `group`. This will list
            all subtasks. When a group config is created (as opposed to task
            config having `group` parameter set), this will be set to -1 to
            avoid recursive indexing. The whole list of subtasks will be loaded
            at evaluation.

        :param task_dir: str
            A directory to check for tasks

        :return
            Dictionary of task names as key and task metadata
        """
767

Baber's avatar
Baber committed
768
769
770
        def _populate_tags_and_groups(
            config: Dict, task: str, tasks_and_groups: Dict[str, Dict]
        ) -> None:
771
            # TODO: remove group in next release
772
773
774
775
776
777
778
779
780
781
782
783
784
            if "tag" in config:
                attr_list = config["tag"]
                if isinstance(attr_list, str):
                    attr_list = [attr_list]

                for tag in attr_list:
                    if tag not in tasks_and_groups:
                        tasks_and_groups[tag] = {
                            "type": "tag",
                            "task": [task],
                            "yaml_path": -1,
                        }
                    elif tasks_and_groups[tag]["type"] != "tag":
Lintang Sutawika's avatar
Lintang Sutawika committed
785
                        eval_logger.info(
786
787
                            f"The tag '{tag}' is already registered as a group, this tag will not be registered. "
                            "This may affect tasks you want to call."
788
                        )
789
790
791
                        break
                    else:
                        tasks_and_groups[tag]["task"].append(task)
792

Lintang Sutawika's avatar
Lintang Sutawika committed
793
        # TODO: remove group in next release
794
795
796
797
        # ignore_dirs = [
        #     "__pycache__",
        #     ".ipynb_checkpoints",
        # ]
798
        tasks_and_groups = collections.defaultdict()
Baber's avatar
Baber committed
799
800
        task_dir_path = Path(task_dir)

801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
        for yaml_path in iter_yaml_files(task_dir_path):
            try:
                config = load_yaml_config(
                    yaml_path, mode="simple", resolve_includes=False
                )
            except (FileNotFoundError, YAMLError, OSError) as err:
                eval_logger.debug(f"File {yaml_path} could not be loaded ({err})")
                continue
            if self._config_is_python_task(config):
                # This is a python class config
                task = config["task"]
                self._register_task(
                    task,
                    "python_task",
                    str(yaml_path),
                    tasks_and_groups,
                    config,
                    _populate_tags_and_groups,
                )
            elif self._config_is_group(config):
                # This is a group config
                tasks_and_groups[config["group"]] = {
                    "type": "group",
                    "task": -1,  # This signals that
                    # we don't need to know
                    # the task list for indexing
                    # as it can be loaded
                    # when called.
                    "yaml_path": str(yaml_path),
                }
lintangsutawika's avatar
lintangsutawika committed
831

832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
                # # Registered the level 1 tasks from a group config
                # for config in config["task"]:
                #     if isinstance(config, dict) and self._config_is_task(config):
                #         task = config["task"]
                #         tasks_and_groups[task] = {
                #             "type": "task",
                #             "yaml_path": yaml_path,
                #             }

            elif self._config_is_task(config):
                # This is a task config
                task = config["task"]
                self._register_task(
                    task,
                    "task",
                    str(yaml_path),
                    tasks_and_groups,
                    config,
                    _populate_tags_and_groups,
                )
            elif self._config_is_task_list(config):
                # This is a task_list config
                for task_entry in config["task_list"]:
                    if isinstance(task_entry, dict) and "task" in task_entry:
                        task_name = task_entry["task"]
857
                        self._register_task(
858
                            task_name,
859
                            "task",
Baber's avatar
Baber committed
860
                            str(yaml_path),
861
862
863
                            tasks_and_groups,
                            config,
                            _populate_tags_and_groups,
864
                        )
865
866
            else:
                eval_logger.debug(f"File {yaml_path} could not be loaded")
867
868

        return tasks_and_groups
lintangsutawika's avatar
lintangsutawika committed
869

870

871
872
873
874
875
876
877
def get_task_name_from_config(task_config: Dict[str, str]) -> str:
    if "task" in task_config:
        return task_config["task"]
    if "dataset_name" in task_config:
        return "{dataset_path}_{dataset_name}".format(**task_config)
    else:
        return "{dataset_path}".format(**task_config)
lintangsutawika's avatar
lintangsutawika committed
878

879

Baber's avatar
nit  
Baber committed
880
def get_task_name_from_object(task_object: Union["ConfigurableTask", "Task"]) -> str:
881
882
    if hasattr(task_object, "config"):
        return task_object._config["task"]
lintangsutawika's avatar
lintangsutawika committed
883
884
885
886
887
888
889
890
891

    # TODO: scrap this
    # this gives a mechanism for non-registered tasks to have a custom name anyways when reporting
    return (
        task_object.EVAL_HARNESS_NAME
        if hasattr(task_object, "EVAL_HARNESS_NAME")
        else type(task_object).__name__
    )

892

Baber's avatar
Baber committed
893
def _check_duplicates(task_dict: Dict[str, List[str]]) -> None:
Lintang Sutawika's avatar
Lintang Sutawika committed
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
    """helper function solely used in validating get_task_dict output.
    Takes the output of lm_eval.evaluator_utils.get_subtask_list and
    returns a list of all leaf subtasks contained within, and errors if any such leaf subtasks are
    "oversubscribed" to several disjoint groups.
    """
    subtask_names = []
    for key, value in task_dict.items():
        subtask_names.extend(value)

    duplicate_tasks = {
        task_name for task_name in subtask_names if subtask_names.count(task_name) > 1
    }

    # locate the potentially problematic groups that seem to 'compete' for constituent subtasks
    competing_groups = [
        group
        for group in task_dict.keys()
        if len(set(task_dict[group]).intersection(duplicate_tasks)) > 0
    ]

    if len(duplicate_tasks) > 0:
        raise ValueError(
            f"Found 1 or more tasks while trying to call get_task_dict() that were members of more than 1 called group: {list(duplicate_tasks)}. Offending groups: {competing_groups}. Please call groups which overlap their constituent tasks in separate evaluation runs."
        )


920
def get_task_dict(
Baber's avatar
nit  
Baber committed
921
    task_name_list: Union[str, List[Union[str, Dict, "Task"]]],
922
    task_manager: Optional[TaskManager] = None,
Baber's avatar
nit  
Baber committed
923
) -> Dict[str, Union["ConfigurableTask", "Task"]]:
924
    """Creates a dictionary of task objects from either a name of task, config, or prepared Task object.
lintangsutawika's avatar
lintangsutawika committed
925

926
927
928
929
930
931
932
    :param task_name_list: List[Union[str, Dict, Task]]
        Name of model or LM object, see lm_eval.models.get_model
    :param task_manager: TaskManager = None
        A TaskManager object that stores indexed tasks. If not set,
        task_manager will load one. This should be set by the user
        if there are additional paths that want to be included
        via `include_path`
933

934
935
936
    :return
        Dictionary of task objects
    """
Baber's avatar
nit  
Baber committed
937
    from lm_eval.api.task import ConfigurableTask, Task
Lintang Sutawika's avatar
Lintang Sutawika committed
938

939
    task_name_from_string_dict = {}
940
941
942
    task_name_from_config_dict = {}
    task_name_from_object_dict = {}

943
    if isinstance(task_name_list, str):
lintangsutawika's avatar
lintangsutawika committed
944
        task_name_list = [task_name_list]
945
946
947
948
949
950
951
952
953
    elif isinstance(task_name_list, list):
        if not all([isinstance(task, (str, dict, Task)) for task in task_name_list]):
            raise TypeError(
                "Expected all list items to be of types 'str', 'dict', or 'Task', but at least one entry did not match."
            )
    else:
        raise TypeError(
            f"Expected a 'str' or 'list' but received {type(task_name_list)}."
        )
lintangsutawika's avatar
lintangsutawika committed
954

955
    string_task_name_list = [task for task in task_name_list if isinstance(task, str)]
956
957
958
    others_task_name_list = [
        task for task in task_name_list if not isinstance(task, str)
    ]
959
960
961
    if len(string_task_name_list) > 0:
        if task_manager is None:
            task_manager = TaskManager()
lintangsutawika's avatar
lintangsutawika committed
962

963
964
965
        task_name_from_string_dict = task_manager.load_task_or_group(
            string_task_name_list
        )
966

967
968
    for task_element in others_task_name_list:
        if isinstance(task_element, dict):
969
970
            task_name_from_config_dict = {
                **task_name_from_config_dict,
971
                **task_manager.load_config(config=task_element),
972
973
974
975
976
            }

        elif isinstance(task_element, Task):
            task_name_from_object_dict = {
                **task_name_from_object_dict,
lintangsutawika's avatar
lintangsutawika committed
977
                get_task_name_from_object(task_element): task_element,
978
            }
lintangsutawika's avatar
lintangsutawika committed
979

980
    if not set(task_name_from_string_dict.keys()).isdisjoint(
lintangsutawika's avatar
lintangsutawika committed
981
        set(task_name_from_object_dict.keys())
982
983
    ):
        raise ValueError
984

Lintang Sutawika's avatar
Lintang Sutawika committed
985
    final_task_dict = {
986
        **task_name_from_string_dict,
lintangsutawika's avatar
lintangsutawika committed
987
        **task_name_from_config_dict,
988
        **task_name_from_object_dict,
lintangsutawika's avatar
lintangsutawika committed
989
    }
Lintang Sutawika's avatar
Lintang Sutawika committed
990
991
992
993
994
995
996
997

    # behavior can get odd if one tries to invoke several groups that "compete" for the same task.
    # (notably, because one could request several num_fewshot values at once in GroupConfig overrides for the subtask
    # and we'd be unsure which to use and report.)
    # we explicitly check and error in this case.
    _check_duplicates(get_subtask_list(final_task_dict))

    return final_task_dict