__init__.py 36.2 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
10
11
12
13
from typing import Any, Callable, Dict, Generator, List, Mapping, Optional, Union

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

15
from lm_eval import utils
Lintang Sutawika's avatar
Lintang Sutawika committed
16
from lm_eval.api.group import ConfigurableGroup, GroupConfig
17
from lm_eval.api.task import ConfigurableTask, Task
Lintang Sutawika's avatar
Lintang Sutawika committed
18
19
20
21
from lm_eval.evaluator_utils import get_subtask_list


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

lintangsutawika's avatar
lintangsutawika committed
24

Lintang Sutawika's avatar
Lintang Sutawika committed
25
eval_logger = logging.getLogger(__name__)
26
27
28
29
30
31
32
33
34
35
36
37
38
39
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
_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
185

186

187
188
189
class TaskManager:
    """TaskManager indexes all tasks from the default `lm_eval/tasks/`
    and an optional directory if provided.
190

191
192
    """

193
194
    def __init__(
        self,
Lintang Sutawika's avatar
Lintang Sutawika committed
195
        verbosity: Optional[str] = None,
Baber's avatar
Baber committed
196
        include_path: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
197
        include_defaults: bool = True,
Baber Abbasi's avatar
Baber Abbasi committed
198
        metadata: Optional[dict] = None,
199
    ) -> None:
Lintang Sutawika's avatar
Lintang Sutawika committed
200
201
        if verbosity is not None:
            utils.setup_logging(verbosity)
202
        self.include_path = include_path
Baber Abbasi's avatar
Baber Abbasi committed
203
        self.metadata = metadata
204
205
206
        self._task_index = self.initialize_tasks(
            include_path=include_path, include_defaults=include_defaults
        )
207
        self._all_tasks = sorted(list(self._task_index.keys()))
208

209
210
211
212
        self._all_groups = sorted(
            [x for x in self._all_tasks if self._task_index[x]["type"] == "group"]
        )
        self._all_subtasks = sorted(
213
214
215
216
217
            [
                x
                for x in self._all_tasks
                if self._task_index[x]["type"] in ["task", "python_task"]
            ]
218
219
220
221
222
        )
        self._all_tags = sorted(
            [x for x in self._all_tasks if self._task_index[x]["type"] == "tag"]
        )

223
        self.task_group_map = collections.defaultdict(list)
224

225
226
    def initialize_tasks(
        self,
Baber's avatar
Baber committed
227
        include_path: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
228
        include_defaults: bool = True,
Baber Abbasi's avatar
Baber Abbasi committed
229
230
    ) -> dict[str, dict]:
        """Creates a dictionary of tasks indexes.
231

232
233
234
235
236
        :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
237
        return
238
239
            Dictionary of task names as key and task metadata
        """
240
        if include_defaults:
Baber's avatar
Baber committed
241
            all_paths = [Path(__file__).parent]
242
243
        else:
            all_paths = []
244
        if include_path is not None:
Baber's avatar
Baber committed
245
            if isinstance(include_path, (str, Path)):
246
                include_path = [include_path]
Baber's avatar
Baber committed
247
248
            # Convert all paths to Path objects
            all_paths.extend(Path(p) for p in include_path)
249

250
251
252
253
        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
254

255
256
257
        return task_index

    @property
Baber's avatar
Baber committed
258
    def all_tasks(self) -> List[str]:
259
260
        return self._all_tasks

261
    @property
Baber's avatar
Baber committed
262
    def all_groups(self) -> List[str]:
263
264
265
        return self._all_groups

    @property
Baber's avatar
Baber committed
266
    def all_subtasks(self) -> List[str]:
267
268
269
        return self._all_subtasks

    @property
Baber's avatar
Baber committed
270
    def all_tags(self) -> List[str]:
271
272
        return self._all_tags

273
    @property
Baber's avatar
Baber committed
274
    def task_index(self) -> Dict[str, Dict[str, Union[str, int, List[str]]]]:
275
276
        return self._task_index

277
    def list_all_tasks(
Baber's avatar
Baber committed
278
279
280
281
        self,
        list_groups: bool = True,
        list_tags: bool = True,
        list_subtasks: bool = True,
282
    ) -> str:
283
284
285
286
287
        """
        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.
        """
288
289
        from pytablewriter import MarkdownTableWriter

290
        # ------------------------------------------------------------------ helpers
Baber's avatar
Baber committed
291
        def sanitize_path(path: str) -> str:
292
293
            # print a relative path for anything inside lm_eval/tasks/
            # path_str = str(path)
294
295
            if "lm_eval/tasks/" in path:
                return "lm_eval/tasks/" + path.split("lm_eval/tasks/")[-1]
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
            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
319
320
        group_table = MarkdownTableWriter()
        group_table.headers = ["Group", "Config Location"]
321
322
323
324
325
326
327
328
329
        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
        ]
330

331
        # ---------------------------------------------------------------- TAG table
332
333
334
335
        tag_table = MarkdownTableWriter()
        tag_table.headers = ["Tag"]
        tag_table.value_matrix = [[t] for t in self.all_tags]

336
        # ------------------------------------------------------------ SUBTASK table
337
338
        subtask_table = MarkdownTableWriter()
        subtask_table.headers = ["Task", "Config Location", "Output Type"]
339
340
        st_values: list[list[str]] = []

341
        for t in self.all_subtasks:
342
343
344
345
346
347
            raw_path = self.task_index[t]["yaml_path"]

            if raw_path == -1:
                # python-only task or generated at runtime
                display_path = "---"
                output_type = ""
348
            else:
349
350
351
352
353
354
355
356
357
358
359
360
                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])

361
362
        subtask_table.value_matrix = st_values

363
364
        # ------------------------------------------------------------- final string
        parts: list[str] = ["\n"]
365
        if list_groups:
366
367
            parts.append(group_table.dumps())
            parts.append("\n")
368
        if list_tags:
369
370
            parts.append(tag_table.dumps())
            parts.append("\n")
371
        if list_subtasks:
372
373
374
375
            parts.append(subtask_table.dumps())
            parts.append("\n")

        return "".join(parts)
376

Baber Abbasi's avatar
Baber Abbasi committed
377
    def match_tasks(self, task_list: list[str]) -> list[str]:
378
        return utils.pattern_match(task_list, self.all_tasks)
379

Baber Abbasi's avatar
Baber Abbasi committed
380
    def _name_is_registered(self, name: str) -> bool:
381
        return name in self.all_tasks
382

Baber Abbasi's avatar
Baber Abbasi committed
383
    def _name_is_task(self, name: str) -> bool:
384
385
386
        return (
            self._name_is_registered(name) and self.task_index[name]["type"] == "task"
        )
Lintang Sutawika's avatar
Lintang Sutawika committed
387

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

Baber Abbasi's avatar
Baber Abbasi committed
391
    def _name_is_group(self, name: str) -> bool:
392
393
394
        return (
            self._name_is_registered(name) and self.task_index[name]["type"] == "group"
        )
395

Baber Abbasi's avatar
Baber Abbasi committed
396
    def _name_is_python_task(self, name: str) -> bool:
397
398
399
400
        return (
            self._name_is_registered(name)
            and self.task_index[name]["type"] == "python_task"
        )
401

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

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

Baber Abbasi's avatar
Baber Abbasi committed
408
    def _config_is_python_task(self, config: dict) -> bool:
409
410
411
412
        return "class" in config

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

Baber's avatar
Baber committed
414
    def _get_yaml_path(self, name: str) -> Union[str, int]:
415
416
        if name not in self.task_index:
            raise ValueError
417
418
        return self.task_index[name]["yaml_path"]

Baber's avatar
Baber committed
419
    def _get_config(self, name: str) -> Dict:
420
421
        if name not in self.task_index:
            raise ValueError
422
423
424
425
        yaml_path = self._get_yaml_path(name)
        if yaml_path == -1:
            return {}
        else:
426
            return load_yaml_config(Path(yaml_path), mode="full")
427

Baber's avatar
Baber committed
428
    def _get_tasklist(self, name: str) -> Union[List[str], int]:
429
430
        if self._name_is_task(name):
            raise ValueError
431
432
        return self.task_index[name]["task"]

433
434
435
436
437
    def _register_task(
        self,
        task_name: str,
        task_type: str,
        yaml_path: str,
Baber's avatar
Baber committed
438
439
440
441
        tasks_and_groups: Dict[str, Dict],
        config: Optional[Dict] = None,
        populate_tags_fn: Optional[callable] = None,
    ) -> None:
442
443
444
445
446
447
448
        """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:
449
            populate_tags_fn(config, task_name, tasks_and_groups)
450
451

    def _merge_task_configs(
Baber's avatar
Baber committed
452
453
        self, base_config: Dict, task_specific_config: Dict, task_name: str
    ) -> Dict:
454
455
456
457
458
459
460
        """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
461
462
463
    def _process_tag_subtasks(
        self, tag_name: str, update_config: Optional[Dict] = None
    ) -> Dict:
464
465
466
467
468
469
470
471
        """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
472
    def _process_alias(self, config: Dict, group: Optional[str] = None) -> Dict:
473
474
475
476
477
478
479
480
        # 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
481
    def _class_has_config_in_constructor(self, cls) -> bool:
482
483
484
485
486
487
488
        constructor = getattr(cls, "__init__", None)
        return (
            "config" in inspect.signature(constructor).parameters
            if constructor
            else False
        )

489
    def _load_individual_task_or_group(
490
        self,
Baber's avatar
Baber committed
491
        name_or_config: Optional[Union[str, Dict]] = None,
492
        parent_name: Optional[str] = None,
Baber's avatar
Baber committed
493
        update_config: Optional[Dict] = None,
494
    ) -> Mapping:
Baber's avatar
Baber committed
495
496
497
        def _load_task(
            config: Dict, task: str, yaml_path: Optional[str] = None
        ) -> Dict[str, Union[ConfigurableTask, Task]]:
498
            if "include" in config:
499
500
501
                # Store the task name to preserve it after include processing
                original_task_name = config.get("task", task)

502
                config = {
503
504
                    **load_yaml_config(
                        yaml_path=Path(yaml_path),
505
                        yaml_config={"include": config.pop("include")},
506
                        mode="full" if yaml_path else "simple",
507
508
                    ),
                    **config,
509
                    "task": original_task_name,
510
                }
511
512
513
514

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

515
            if self._config_is_python_task(config):
516
517
518
519
520
521
                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?
522
                    task_object.config.task = task
523
            else:
Baber Abbasi's avatar
Baber Abbasi committed
524
525
526
527
                if self.metadata is not None:
                    config["metadata"] = config.get("metadata", {}) | self.metadata
                else:
                    config["metadata"] = config.get("metadata", {})
528
                task_object = ConfigurableTask(config=config)
Lintang Sutawika's avatar
Lintang Sutawika committed
529

530
531
            return {task: task_object}

Baber Abbasi's avatar
Baber Abbasi committed
532
        def _get_group_and_subtask_from_config(
Baber's avatar
Baber committed
533
534
            config: Dict,
        ) -> tuple[ConfigurableGroup, List[str]]:
Baber Abbasi's avatar
Baber Abbasi committed
535
536
            if self.metadata is not None:
                config["metadata"] = config.get("metadata", {}) | self.metadata
Lintang Sutawika's avatar
Lintang Sutawika committed
537
538
539
540
541
542
543
544
545
            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
546
        def _process_group_config(
Baber's avatar
Baber committed
547
548
            config: Dict, update_config: Optional[Dict] = None
        ) -> tuple[Dict, Optional[Dict]]:
Lintang Sutawika's avatar
Lintang Sutawika committed
549
550
551
552
553
554
555
556
557
558
559
            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

560
561
562
563
        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
564
565
566
            elif self._name_is_task(name_or_config) or self._name_is_python_task(
                name_or_config
            ):
567
568
                # Get the yaml_path for this task
                yaml_path = self._get_yaml_path(name_or_config)
569
                task_config = self._get_config(name_or_config)
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598

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

599
                return _load_task(task_config, task=name_or_config, yaml_path=yaml_path)
600
            else:
601
602
603
                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
604
605
606
607
608
609
                    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):
610
611
612
                        return self._process_tag_subtasks(
                            name_or_config,
                            name_or_config
Lintang Sutawika's avatar
Lintang Sutawika committed
613
614
615
616
617
618
619
                            if isinstance(name_or_config, dict)
                            else None,
                        )
                    else:
                        group_name = ConfigurableGroup(
                            config={"group": name_or_config, "task": subtask_list}
                        )
620

621
622
        if isinstance(name_or_config, dict):
            if self._config_is_task(name_or_config):
Lintang Sutawika's avatar
Lintang Sutawika committed
623
624
625
                name = name_or_config.pop("task")
                if update_config is not None:
                    name_or_config = {**name_or_config, **update_config}
626
627
                # If the name is registered as a group
                if self._name_is_group(name):
Lintang Sutawika's avatar
Lintang Sutawika committed
628
629
630
631
632
633
634
635
636
                    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):
637
                    return self._process_tag_subtasks(name, name_or_config)
638
                else:
639
                    yaml_path = None
640
                    if self._name_is_registered(name):
641
                        yaml_path = self._get_yaml_path(name)
642
643
644
645
                        base_task_config = self._get_config(name)

                        # Check if this is a duplicate.
                        if parent_name is not None:
646
647
648
649
650
651
652
653
                            num_duplicate = len(
                                list(
                                    filter(
                                        lambda x: x.startswith(name),
                                        self.task_group_map[parent_name],
                                    )
                                )
                            )
654
655
656
657
                            if num_duplicate > 0:
                                name = f"{name}-{num_duplicate}"
                            self.task_group_map[parent_name].append(name)

658
659
660
661
                        task_config = {
                            **base_task_config,
                            **name_or_config,
                        }
662
663
                    else:
                        task_config = name_or_config
664
                    return _load_task(task_config, task=name, yaml_path=yaml_path)
665
            else:
Lintang Sutawika's avatar
Lintang Sutawika committed
666
667
668
669
                group_config, update_config = _process_group_config(name_or_config)
                group_name, subtask_list = _get_group_and_subtask_from_config(
                    group_config
                )
670

671
672
673
674
675
        fn = partial(
            self._load_individual_task_or_group,
            parent_name=group_name,
            update_config=update_config,
        )
Lintang Sutawika's avatar
Lintang Sutawika committed
676
677
        return {
            group_name: dict(collections.ChainMap(*map(fn, reversed(subtask_list))))
678
        }
679

Baber's avatar
Baber committed
680
681
682
    def load_task_or_group(
        self, task_list: Optional[Union[str, List[str]]] = None
    ) -> Dict:
683
        """Loads a dictionary of task objects from a list
684

685
686
        :param task_list: Union[str, list] = None
            Single string or list of string of task names to be loaded
687

688
689
690
691
692
        :return
            Dictionary of task objects
        """
        if isinstance(task_list, str):
            task_list = [task_list]
693

694
        all_loaded_tasks = dict(
Baber Abbasi's avatar
Baber Abbasi committed
695
696
697
698
699
700
            collections.ChainMap(
                *map(
                    lambda task: self._load_individual_task_or_group(task),
                    task_list,
                )
            )
701
702
703
        )
        return all_loaded_tasks

Baber's avatar
Baber committed
704
    def load_config(self, config: Dict) -> Mapping:
705
706
        return self._load_individual_task_or_group(config)

Baber's avatar
Baber committed
707
    def _get_task_and_group(self, task_dir: Union[str, Path]) -> Dict[str, Dict]:
708
        """Creates a dictionary of tasks index with the following metadata,
Lintang Sutawika's avatar
Lintang Sutawika committed
709
        - `type`, that can be either `task`, `python_task`, `group` or `tags`.
710
711
            `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
712
713
            `group` are group configs. `tags` are labels that can be assigned
            to tasks to assist in sorting and calling tasks of certain themes.
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
        - `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
        """
729

Baber's avatar
Baber committed
730
731
732
        def _populate_tags_and_groups(
            config: Dict, task: str, tasks_and_groups: Dict[str, Dict]
        ) -> None:
733
            # TODO: remove group in next release
734
735
736
737
738
739
740
741
742
743
744
745
746
            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
747
                        eval_logger.info(
748
749
                            f"The tag '{tag}' is already registered as a group, this tag will not be registered. "
                            "This may affect tasks you want to call."
750
                        )
751
752
753
                        break
                    else:
                        tasks_and_groups[tag]["task"].append(task)
754

Lintang Sutawika's avatar
Lintang Sutawika committed
755
        # TODO: remove group in next release
756
757
758
759
        # ignore_dirs = [
        #     "__pycache__",
        #     ".ipynb_checkpoints",
        # ]
760
        tasks_and_groups = collections.defaultdict()
Baber's avatar
Baber committed
761
762
        task_dir_path = Path(task_dir)

763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
        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
793

794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
                # # 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"]
819
                        self._register_task(
820
                            task_name,
821
                            "task",
Baber's avatar
Baber committed
822
                            str(yaml_path),
823
824
825
                            tasks_and_groups,
                            config,
                            _populate_tags_and_groups,
826
                        )
827
828
            else:
                eval_logger.debug(f"File {yaml_path} could not be loaded")
829
830

        return tasks_and_groups
lintangsutawika's avatar
lintangsutawika committed
831

832

833
834
835
836
837
838
839
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
840

841

Baber's avatar
Baber committed
842
def get_task_name_from_object(task_object: Union[ConfigurableTask, Task]) -> str:
843
844
    if hasattr(task_object, "config"):
        return task_object._config["task"]
lintangsutawika's avatar
lintangsutawika committed
845
846
847
848
849
850
851
852
853

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

854

Baber's avatar
Baber committed
855
def _check_duplicates(task_dict: Dict[str, List[str]]) -> None:
Lintang Sutawika's avatar
Lintang Sutawika committed
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
    """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."
        )


882
def get_task_dict(
883
    task_name_list: Union[str, List[Union[str, Dict, Task]]],
884
    task_manager: Optional[TaskManager] = None,
Baber's avatar
Baber committed
885
) -> Dict[str, Union[ConfigurableTask, Task]]:
886
    """Creates a dictionary of task objects from either a name of task, config, or prepared Task object.
lintangsutawika's avatar
lintangsutawika committed
887

888
889
890
891
892
893
894
    :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`
895

896
897
898
    :return
        Dictionary of task objects
    """
Lintang Sutawika's avatar
Lintang Sutawika committed
899

900
    task_name_from_string_dict = {}
901
902
903
    task_name_from_config_dict = {}
    task_name_from_object_dict = {}

904
    if isinstance(task_name_list, str):
lintangsutawika's avatar
lintangsutawika committed
905
        task_name_list = [task_name_list]
906
907
908
909
910
911
912
913
914
    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
915

916
    string_task_name_list = [task for task in task_name_list if isinstance(task, str)]
917
918
919
    others_task_name_list = [
        task for task in task_name_list if not isinstance(task, str)
    ]
920
921
922
    if len(string_task_name_list) > 0:
        if task_manager is None:
            task_manager = TaskManager()
lintangsutawika's avatar
lintangsutawika committed
923

924
925
926
        task_name_from_string_dict = task_manager.load_task_or_group(
            string_task_name_list
        )
927

928
929
    for task_element in others_task_name_list:
        if isinstance(task_element, dict):
930
931
            task_name_from_config_dict = {
                **task_name_from_config_dict,
932
                **task_manager.load_config(config=task_element),
933
934
935
936
937
            }

        elif isinstance(task_element, Task):
            task_name_from_object_dict = {
                **task_name_from_object_dict,
lintangsutawika's avatar
lintangsutawika committed
938
                get_task_name_from_object(task_element): task_element,
939
            }
lintangsutawika's avatar
lintangsutawika committed
940

941
    if not set(task_name_from_string_dict.keys()).isdisjoint(
lintangsutawika's avatar
lintangsutawika committed
942
        set(task_name_from_object_dict.keys())
943
944
    ):
        raise ValueError
945

Lintang Sutawika's avatar
Lintang Sutawika committed
946
    final_task_dict = {
947
        **task_name_from_string_dict,
lintangsutawika's avatar
lintangsutawika committed
948
        **task_name_from_config_dict,
949
        **task_name_from_object_dict,
lintangsutawika's avatar
lintangsutawika committed
950
    }
Lintang Sutawika's avatar
Lintang Sutawika committed
951
952
953
954
955
956
957
958

    # 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