__init__.py 36.6 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

import yaml
Baber's avatar
nit  
Baber committed
23
from memory_profiler import profile
24
from yaml import YAMLError
&'s avatar
& committed
25

Lintang Sutawika's avatar
Lintang Sutawika committed
26
27
from lm_eval.api.group import ConfigurableGroup, GroupConfig
from lm_eval.evaluator_utils import get_subtask_list
Baber's avatar
nit  
Baber committed
28
29
30
31
32
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
33
34
35


GROUP_ONLY_KEYS = list(GroupConfig().to_dict().keys())
36
37
_CONFIG_CACHE: dict[tuple[Path, str], dict] = {}

lintangsutawika's avatar
lintangsutawika committed
38

Lintang Sutawika's avatar
Lintang Sutawika committed
39
eval_logger = logging.getLogger(__name__)
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
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
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
171
172
173
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
_Base = yaml.CLoader if getattr(yaml, "__with_libyaml__", False) else yaml.FullLoader


@functools.lru_cache(maxsize=128)  # ← reuse per (directory, simple) pair
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:
        yaml.add_constructor("!function", lambda *_: None, Loader=Loader)
    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


@functools.lru_cache(maxsize=1000)  # ← cache module objects
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


@functools.lru_cache(maxsize=1000)  #
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)


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()

        # ---------- fast-path: return memoised, already-resolved cfg ----------
        cache_key = (yaml_path, mode)
        if yaml_config is None and resolve_includes and cache_key in _CONFIG_CACHE:
            return _CONFIG_CACHE[cache_key]

        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
    # -------- memoise after *all* includes are merged ----------
    if yaml_config is None and resolve_includes:
        _CONFIG_CACHE[cache_key] = final_cfg

    return final_cfg


def iter_yaml_files(root: Path) -> Generator[Path, Any, None]:
    # '**/*.yaml' is handled internally by os.scandir.
    for path in iglob("**/*.yaml", root_dir=root, recursive=True):
        # quick ignore check
        if "/__pycache__/" in path or "/.ipynb_checkpoints/" in path:
            continue
        yield root / path
Lintang Sutawika's avatar
Lintang Sutawika committed
199

200

201
202
203
class TaskManager:
    """TaskManager indexes all tasks from the default `lm_eval/tasks/`
    and an optional directory if provided.
204

205
206
    """

Baber's avatar
nit  
Baber committed
207
    @profile
208
209
    def __init__(
        self,
Lintang Sutawika's avatar
Lintang Sutawika committed
210
        verbosity: Optional[str] = None,
Baber's avatar
Baber committed
211
        include_path: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
212
        include_defaults: bool = True,
Baber Abbasi's avatar
Baber Abbasi committed
213
        metadata: Optional[dict] = None,
214
    ) -> None:
Lintang Sutawika's avatar
Lintang Sutawika committed
215
        if verbosity is not None:
Baber's avatar
nit  
Baber committed
216
            setup_logging(verbosity)
217
        self.include_path = include_path
Baber Abbasi's avatar
Baber Abbasi committed
218
        self.metadata = metadata
219
220
221
        self._task_index = self.initialize_tasks(
            include_path=include_path, include_defaults=include_defaults
        )
222
        self._all_tasks = sorted(list(self._task_index.keys()))
223

224
225
226
227
        self._all_groups = sorted(
            [x for x in self._all_tasks if self._task_index[x]["type"] == "group"]
        )
        self._all_subtasks = sorted(
228
229
230
231
232
            [
                x
                for x in self._all_tasks
                if self._task_index[x]["type"] in ["task", "python_task"]
            ]
233
234
235
236
237
        )
        self._all_tags = sorted(
            [x for x in self._all_tasks if self._task_index[x]["type"] == "tag"]
        )

238
        self.task_group_map = collections.defaultdict(list)
239

Baber's avatar
nit  
Baber committed
240
    @profile
241
242
    def initialize_tasks(
        self,
Baber's avatar
Baber committed
243
        include_path: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
244
        include_defaults: bool = True,
Baber Abbasi's avatar
Baber Abbasi committed
245
246
    ) -> dict[str, dict]:
        """Creates a dictionary of tasks indexes.
247

248
249
250
251
252
        :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
253
        return
254
255
            Dictionary of task names as key and task metadata
        """
256
        if include_defaults:
Baber's avatar
Baber committed
257
            all_paths = [Path(__file__).parent]
258
259
        else:
            all_paths = []
260
        if include_path is not None:
Baber's avatar
Baber committed
261
            if isinstance(include_path, (str, Path)):
262
                include_path = [include_path]
Baber's avatar
Baber committed
263
264
            # Convert all paths to Path objects
            all_paths.extend(Path(p) for p in include_path)
265

266
267
268
269
        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
270

271
272
273
        return task_index

    @property
Baber's avatar
Baber committed
274
    def all_tasks(self) -> List[str]:
275
276
        return self._all_tasks

277
    @property
Baber's avatar
Baber committed
278
    def all_groups(self) -> List[str]:
279
280
281
        return self._all_groups

    @property
Baber's avatar
Baber committed
282
    def all_subtasks(self) -> List[str]:
283
284
285
        return self._all_subtasks

    @property
Baber's avatar
Baber committed
286
    def all_tags(self) -> List[str]:
287
288
        return self._all_tags

289
    @property
Baber's avatar
Baber committed
290
    def task_index(self) -> Dict[str, Dict[str, Union[str, int, List[str]]]]:
291
292
        return self._task_index

293
    def list_all_tasks(
Baber's avatar
Baber committed
294
295
296
297
        self,
        list_groups: bool = True,
        list_tags: bool = True,
        list_subtasks: bool = True,
298
    ) -> str:
299
300
301
302
303
        """
        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.
        """
304
305
        from pytablewriter import MarkdownTableWriter

306
        # ------------------------------------------------------------------ helpers
Baber's avatar
Baber committed
307
        def sanitize_path(path: str) -> str:
308
309
            # print a relative path for anything inside lm_eval/tasks/
            # path_str = str(path)
310
311
            if "lm_eval/tasks/" in path:
                return "lm_eval/tasks/" + path.split("lm_eval/tasks/")[-1]
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
            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
335
336
        group_table = MarkdownTableWriter()
        group_table.headers = ["Group", "Config Location"]
337
338
339
340
341
342
343
344
345
        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
        ]
346

347
        # ---------------------------------------------------------------- TAG table
348
349
350
351
        tag_table = MarkdownTableWriter()
        tag_table.headers = ["Tag"]
        tag_table.value_matrix = [[t] for t in self.all_tags]

352
        # ------------------------------------------------------------ SUBTASK table
353
354
        subtask_table = MarkdownTableWriter()
        subtask_table.headers = ["Task", "Config Location", "Output Type"]
355
356
        st_values: list[list[str]] = []

357
        for t in self.all_subtasks:
358
359
360
361
362
363
            raw_path = self.task_index[t]["yaml_path"]

            if raw_path == -1:
                # python-only task or generated at runtime
                display_path = "---"
                output_type = ""
364
            else:
365
366
367
368
369
370
371
372
373
374
375
376
                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])

377
378
        subtask_table.value_matrix = st_values

379
380
        # ------------------------------------------------------------- final string
        parts: list[str] = ["\n"]
381
        if list_groups:
382
383
            parts.append(group_table.dumps())
            parts.append("\n")
384
        if list_tags:
385
386
            parts.append(tag_table.dumps())
            parts.append("\n")
387
        if list_subtasks:
388
389
390
391
            parts.append(subtask_table.dumps())
            parts.append("\n")

        return "".join(parts)
392

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

Baber Abbasi's avatar
Baber Abbasi committed
396
    def _name_is_registered(self, name: str) -> bool:
397
        return name in self.all_tasks
398

Baber Abbasi's avatar
Baber Abbasi committed
399
    def _name_is_task(self, name: str) -> bool:
400
401
402
        return (
            self._name_is_registered(name) and self.task_index[name]["type"] == "task"
        )
Lintang Sutawika's avatar
Lintang Sutawika committed
403

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

Baber Abbasi's avatar
Baber Abbasi committed
407
    def _name_is_group(self, name: str) -> bool:
408
409
410
        return (
            self._name_is_registered(name) and self.task_index[name]["type"] == "group"
        )
411

Baber Abbasi's avatar
Baber Abbasi committed
412
    def _name_is_python_task(self, name: str) -> bool:
413
414
415
416
        return (
            self._name_is_registered(name)
            and self.task_index[name]["type"] == "python_task"
        )
417

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

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

Baber Abbasi's avatar
Baber Abbasi committed
424
    def _config_is_python_task(self, config: dict) -> bool:
425
426
427
428
        return "class" in config

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

Baber's avatar
Baber committed
430
    def _get_yaml_path(self, name: str) -> Union[str, int]:
431
432
        if name not in self.task_index:
            raise ValueError
433
434
        return self.task_index[name]["yaml_path"]

Baber's avatar
Baber committed
435
    def _get_config(self, name: str) -> Dict:
436
437
        if name not in self.task_index:
            raise ValueError
438
439
440
441
        yaml_path = self._get_yaml_path(name)
        if yaml_path == -1:
            return {}
        else:
442
            return load_yaml_config(Path(yaml_path), mode="full")
443

Baber's avatar
Baber committed
444
    def _get_tasklist(self, name: str) -> Union[List[str], int]:
445
446
        if self._name_is_task(name):
            raise ValueError
447
448
        return self.task_index[name]["task"]

449
450
451
452
453
    def _register_task(
        self,
        task_name: str,
        task_type: str,
        yaml_path: str,
Baber's avatar
Baber committed
454
455
456
457
        tasks_and_groups: Dict[str, Dict],
        config: Optional[Dict] = None,
        populate_tags_fn: Optional[callable] = None,
    ) -> None:
458
459
460
461
462
463
464
        """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:
465
            populate_tags_fn(config, task_name, tasks_and_groups)
466
467

    def _merge_task_configs(
Baber's avatar
Baber committed
468
469
        self, base_config: Dict, task_specific_config: Dict, task_name: str
    ) -> Dict:
470
471
472
473
474
475
476
        """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
477
478
479
    def _process_tag_subtasks(
        self, tag_name: str, update_config: Optional[Dict] = None
    ) -> Dict:
480
481
482
483
484
485
486
487
        """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
488
    def _process_alias(self, config: Dict, group: Optional[str] = None) -> Dict:
489
490
491
492
493
494
495
496
        # 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
497
    def _class_has_config_in_constructor(self, cls) -> bool:
498
499
500
501
502
503
504
        constructor = getattr(cls, "__init__", None)
        return (
            "config" in inspect.signature(constructor).parameters
            if constructor
            else False
        )

505
    def _load_individual_task_or_group(
506
        self,
Baber's avatar
Baber committed
507
        name_or_config: Optional[Union[str, Dict]] = None,
508
        parent_name: Optional[str] = None,
Baber's avatar
Baber committed
509
        update_config: Optional[Dict] = None,
510
    ) -> Mapping:
Baber's avatar
nit  
Baber committed
511
512
        from lm_eval.api.task import ConfigurableTask, Task

Baber's avatar
Baber committed
513
514
        def _load_task(
            config: Dict, task: str, yaml_path: Optional[str] = None
Baber's avatar
nit  
Baber committed
515
        ) -> Dict[str, Union["ConfigurableTask", "Task"]]:
516
            if "include" in config:
517
518
519
                # Store the task name to preserve it after include processing
                original_task_name = config.get("task", task)

520
                config = {
521
522
                    **load_yaml_config(
                        yaml_path=Path(yaml_path),
523
                        yaml_config={"include": config.pop("include")},
524
                        mode="full" if yaml_path else "simple",
525
526
                    ),
                    **config,
527
                    "task": original_task_name,
528
                }
529
530
531
532

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

533
            if self._config_is_python_task(config):
534
535
536
537
538
539
                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?
540
                    task_object.config.task = task
541
            else:
Baber Abbasi's avatar
Baber Abbasi committed
542
543
544
545
                if self.metadata is not None:
                    config["metadata"] = config.get("metadata", {}) | self.metadata
                else:
                    config["metadata"] = config.get("metadata", {})
546
                task_object = ConfigurableTask(config=config)
Lintang Sutawika's avatar
Lintang Sutawika committed
547

548
549
            return {task: task_object}

Baber Abbasi's avatar
Baber Abbasi committed
550
        def _get_group_and_subtask_from_config(
Baber's avatar
Baber committed
551
552
            config: Dict,
        ) -> tuple[ConfigurableGroup, List[str]]:
Baber Abbasi's avatar
Baber Abbasi committed
553
554
            if self.metadata is not None:
                config["metadata"] = config.get("metadata", {}) | self.metadata
Lintang Sutawika's avatar
Lintang Sutawika committed
555
556
557
558
559
560
561
562
563
            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
564
        def _process_group_config(
Baber's avatar
Baber committed
565
566
            config: Dict, update_config: Optional[Dict] = None
        ) -> tuple[Dict, Optional[Dict]]:
Lintang Sutawika's avatar
Lintang Sutawika committed
567
568
569
570
571
572
573
574
575
576
577
            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

578
579
580
581
        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
582
583
584
            elif self._name_is_task(name_or_config) or self._name_is_python_task(
                name_or_config
            ):
585
586
                # Get the yaml_path for this task
                yaml_path = self._get_yaml_path(name_or_config)
587
                task_config = self._get_config(name_or_config)
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616

                # 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}

617
                return _load_task(task_config, task=name_or_config, yaml_path=yaml_path)
618
            else:
619
620
621
                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
622
623
624
625
626
627
                    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):
628
629
630
                        return self._process_tag_subtasks(
                            name_or_config,
                            name_or_config
Lintang Sutawika's avatar
Lintang Sutawika committed
631
632
633
634
635
636
637
                            if isinstance(name_or_config, dict)
                            else None,
                        )
                    else:
                        group_name = ConfigurableGroup(
                            config={"group": name_or_config, "task": subtask_list}
                        )
638

639
640
        if isinstance(name_or_config, dict):
            if self._config_is_task(name_or_config):
Lintang Sutawika's avatar
Lintang Sutawika committed
641
642
643
                name = name_or_config.pop("task")
                if update_config is not None:
                    name_or_config = {**name_or_config, **update_config}
644
645
                # If the name is registered as a group
                if self._name_is_group(name):
Lintang Sutawika's avatar
Lintang Sutawika committed
646
647
648
649
650
651
652
653
654
                    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):
655
                    return self._process_tag_subtasks(name, name_or_config)
656
                else:
657
                    yaml_path = None
658
                    if self._name_is_registered(name):
659
                        yaml_path = self._get_yaml_path(name)
660
661
662
663
                        base_task_config = self._get_config(name)

                        # Check if this is a duplicate.
                        if parent_name is not None:
664
665
666
667
668
669
670
671
                            num_duplicate = len(
                                list(
                                    filter(
                                        lambda x: x.startswith(name),
                                        self.task_group_map[parent_name],
                                    )
                                )
                            )
672
673
674
675
                            if num_duplicate > 0:
                                name = f"{name}-{num_duplicate}"
                            self.task_group_map[parent_name].append(name)

676
677
678
679
                        task_config = {
                            **base_task_config,
                            **name_or_config,
                        }
680
681
                    else:
                        task_config = name_or_config
682
                    return _load_task(task_config, task=name, yaml_path=yaml_path)
683
            else:
Lintang Sutawika's avatar
Lintang Sutawika committed
684
685
686
687
                group_config, update_config = _process_group_config(name_or_config)
                group_name, subtask_list = _get_group_and_subtask_from_config(
                    group_config
                )
688

689
690
691
692
693
        fn = partial(
            self._load_individual_task_or_group,
            parent_name=group_name,
            update_config=update_config,
        )
Lintang Sutawika's avatar
Lintang Sutawika committed
694
695
        return {
            group_name: dict(collections.ChainMap(*map(fn, reversed(subtask_list))))
696
        }
697

Baber's avatar
Baber committed
698
699
700
    def load_task_or_group(
        self, task_list: Optional[Union[str, List[str]]] = None
    ) -> Dict:
701
        """Loads a dictionary of task objects from a list
702

703
704
        :param task_list: Union[str, list] = None
            Single string or list of string of task names to be loaded
705

706
707
708
709
710
        :return
            Dictionary of task objects
        """
        if isinstance(task_list, str):
            task_list = [task_list]
711

712
        all_loaded_tasks = dict(
Baber Abbasi's avatar
Baber Abbasi committed
713
714
715
716
717
718
            collections.ChainMap(
                *map(
                    lambda task: self._load_individual_task_or_group(task),
                    task_list,
                )
            )
719
720
721
        )
        return all_loaded_tasks

Baber's avatar
Baber committed
722
    def load_config(self, config: Dict) -> Mapping:
723
724
        return self._load_individual_task_or_group(config)

Baber's avatar
nit  
Baber committed
725
    @profile
Baber's avatar
Baber committed
726
    def _get_task_and_group(self, task_dir: Union[str, Path]) -> Dict[str, Dict]:
727
        """Creates a dictionary of tasks index with the following metadata,
Lintang Sutawika's avatar
Lintang Sutawika committed
728
        - `type`, that can be either `task`, `python_task`, `group` or `tags`.
729
730
            `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
731
732
            `group` are group configs. `tags` are labels that can be assigned
            to tasks to assist in sorting and calling tasks of certain themes.
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
        - `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
        """
748

Baber's avatar
Baber committed
749
750
751
        def _populate_tags_and_groups(
            config: Dict, task: str, tasks_and_groups: Dict[str, Dict]
        ) -> None:
752
            # TODO: remove group in next release
753
754
755
756
757
758
759
760
761
762
763
764
765
            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
766
                        eval_logger.info(
767
768
                            f"The tag '{tag}' is already registered as a group, this tag will not be registered. "
                            "This may affect tasks you want to call."
769
                        )
770
771
772
                        break
                    else:
                        tasks_and_groups[tag]["task"].append(task)
773

Lintang Sutawika's avatar
Lintang Sutawika committed
774
        # TODO: remove group in next release
775
776
777
778
        # ignore_dirs = [
        #     "__pycache__",
        #     ".ipynb_checkpoints",
        # ]
779
        tasks_and_groups = collections.defaultdict()
Baber's avatar
Baber committed
780
781
        task_dir_path = Path(task_dir)

782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
        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
812

813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
                # # 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"]
838
                        self._register_task(
839
                            task_name,
840
                            "task",
Baber's avatar
Baber committed
841
                            str(yaml_path),
842
843
844
                            tasks_and_groups,
                            config,
                            _populate_tags_and_groups,
845
                        )
846
847
            else:
                eval_logger.debug(f"File {yaml_path} could not be loaded")
848
849

        return tasks_and_groups
lintangsutawika's avatar
lintangsutawika committed
850

851

852
853
854
855
856
857
858
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
859

860

Baber's avatar
nit  
Baber committed
861
def get_task_name_from_object(task_object: Union["ConfigurableTask", "Task"]) -> str:
862
863
    if hasattr(task_object, "config"):
        return task_object._config["task"]
lintangsutawika's avatar
lintangsutawika committed
864
865
866
867
868
869
870
871
872

    # 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__
    )

873

Baber's avatar
Baber committed
874
def _check_duplicates(task_dict: Dict[str, List[str]]) -> None:
Lintang Sutawika's avatar
Lintang Sutawika committed
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
    """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."
        )


Baber's avatar
nit  
Baber committed
901
@profile
902
def get_task_dict(
Baber's avatar
nit  
Baber committed
903
    task_name_list: Union[str, List[Union[str, Dict, "Task"]]],
904
    task_manager: Optional[TaskManager] = None,
Baber's avatar
nit  
Baber committed
905
) -> Dict[str, Union["ConfigurableTask", "Task"]]:
906
    """Creates a dictionary of task objects from either a name of task, config, or prepared Task object.
lintangsutawika's avatar
lintangsutawika committed
907

908
909
910
911
912
913
914
    :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`
915

916
917
918
    :return
        Dictionary of task objects
    """
Baber's avatar
nit  
Baber committed
919
    from lm_eval.api.task import ConfigurableTask, Task
Lintang Sutawika's avatar
Lintang Sutawika committed
920

921
    task_name_from_string_dict = {}
922
923
924
    task_name_from_config_dict = {}
    task_name_from_object_dict = {}

925
    if isinstance(task_name_list, str):
lintangsutawika's avatar
lintangsutawika committed
926
        task_name_list = [task_name_list]
927
928
929
930
931
932
933
934
935
    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
936

937
    string_task_name_list = [task for task in task_name_list if isinstance(task, str)]
938
939
940
    others_task_name_list = [
        task for task in task_name_list if not isinstance(task, str)
    ]
941
942
943
    if len(string_task_name_list) > 0:
        if task_manager is None:
            task_manager = TaskManager()
lintangsutawika's avatar
lintangsutawika committed
944

945
946
947
        task_name_from_string_dict = task_manager.load_task_or_group(
            string_task_name_list
        )
948

949
950
    for task_element in others_task_name_list:
        if isinstance(task_element, dict):
951
952
            task_name_from_config_dict = {
                **task_name_from_config_dict,
953
                **task_manager.load_config(config=task_element),
954
955
956
957
958
            }

        elif isinstance(task_element, Task):
            task_name_from_object_dict = {
                **task_name_from_object_dict,
lintangsutawika's avatar
lintangsutawika committed
959
                get_task_name_from_object(task_element): task_element,
960
            }
lintangsutawika's avatar
lintangsutawika committed
961

962
    if not set(task_name_from_string_dict.keys()).isdisjoint(
lintangsutawika's avatar
lintangsutawika committed
963
        set(task_name_from_object_dict.keys())
964
965
    ):
        raise ValueError
966

Lintang Sutawika's avatar
Lintang Sutawika committed
967
    final_task_dict = {
968
        **task_name_from_string_dict,
lintangsutawika's avatar
lintangsutawika committed
969
        **task_name_from_config_dict,
970
        **task_name_from_object_dict,
lintangsutawika's avatar
lintangsutawika committed
971
    }
Lintang Sutawika's avatar
Lintang Sutawika committed
972
973
974
975
976
977
978
979

    # 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