__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
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Generator,
    Mapping,
    Optional,
    Union,
)
19
20
21

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

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

Baber's avatar
nit  
Baber committed
31
eval_logger = logging.getLogger(__name__)
Lintang Sutawika's avatar
Lintang Sutawika committed
32
33

GROUP_ONLY_KEYS = list(GroupConfig().to_dict().keys())
Baber's avatar
nit  
Baber committed
34
35
36
37
38
_Base = yaml.CLoader if getattr(yaml, "__with_libyaml__", False) else yaml.FullLoader
_IGNORE_DIRS = (
    "__pycache__",
    ".ipynb_checkpoints",
)
39

lintangsutawika's avatar
lintangsutawika committed
40

Baber's avatar
nit  
Baber committed
41
42
def ignore_constructor(loader: yaml.Loader, node: yaml.Node) -> None:
    return None
43
44


Baber's avatar
nit  
Baber committed
45
@functools.lru_cache(maxsize=2048)  # ← reuse per (directory, simple) pair
46
47
48
49
50
51
52
53
54
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
Baber's avatar
nit  
Baber committed
55
56
        If True we ignore !function completely (used by `mode="simple"`),
        used on TaskManager init to index.
57
58
59
60
61
62
63
    """

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

    # Register (or stub) the !function constructor **for this Loader only**
    if simple:
Baber's avatar
nit  
Baber committed
64
        yaml.add_constructor("!function", ignore_constructor, Loader=Loader)
65
66
67
68
69
70
71
72
73
74
75
76
77
78
    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
79
@functools.lru_cache(maxsize=None)  # ← cache module objects
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
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)


Baber's avatar
nit  
Baber committed
96
@functools.lru_cache(maxsize=4096)  #
97
98
99
100
101
102
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
103
@functools.lru_cache(maxsize=4096)
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
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


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
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
170
171
172
        # ---------- fast-path: use LRU cached function ----------
        if yaml_config is None and resolve_includes:
            return _get_cached_config(yaml_path, mode)
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
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213

        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]:
Baber's avatar
nit  
Baber committed
214
215
    for p in iglob("**/*.yaml", root_dir=root, recursive=True):
        # ignore check
Baber's avatar
nit  
Baber committed
216
        if Path(p).parts[0] in _IGNORE_DIRS:
217
            continue
Baber's avatar
nit  
Baber committed
218
        yield root / p
Lintang Sutawika's avatar
Lintang Sutawika committed
219

220

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

225
226
    """

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

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

257
        self.task_group_map = collections.defaultdict(list)
258

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

Baber's avatar
nit  
Baber committed
266
        :param include_path: Union[str, list] = None
267
268
269
270
            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
271
        return
Baber's avatar
nit  
Baber committed
272
            dictionary of task names as key and task metadata
273
        """
274
        if include_defaults:
Baber's avatar
Baber committed
275
            all_paths = [Path(__file__).parent]
276
277
        else:
            all_paths = []
278
        if include_path is not None:
Baber's avatar
Baber committed
279
            if isinstance(include_path, (str, Path)):
280
                include_path = [include_path]
Baber's avatar
Baber committed
281
282
            # Convert all paths to Path objects
            all_paths.extend(Path(p) for p in include_path)
283

284
285
286
287
        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
288

289
290
291
        return task_index

    @property
Baber's avatar
nit  
Baber committed
292
    def all_tasks(self) -> list[str]:
293
294
        return self._all_tasks

295
    @property
Baber's avatar
nit  
Baber committed
296
    def all_groups(self) -> list[str]:
297
298
299
        return self._all_groups

    @property
Baber's avatar
nit  
Baber committed
300
    def all_subtasks(self) -> list[str]:
301
302
303
        return self._all_subtasks

    @property
Baber's avatar
nit  
Baber committed
304
    def all_tags(self) -> list[str]:
305
306
        return self._all_tags

307
    @property
Baber's avatar
nit  
Baber committed
308
    def task_index(self) -> dict[str, dict[str, Union[str, int, list[str]]]]:
309
310
        return self._task_index

311
    def list_all_tasks(
Baber's avatar
Baber committed
312
313
314
315
        self,
        list_groups: bool = True,
        list_tags: bool = True,
        list_subtasks: bool = True,
316
    ) -> str:
317
318
319
320
321
        """
        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.
        """
322
323
        from pytablewriter import MarkdownTableWriter

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

        # -------------------------------------------------------------- GROUP table
354
355
        group_table = MarkdownTableWriter()
        group_table.headers = ["Group", "Config Location"]
356
357
358
359
360
361
362
363
364
        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
        ]
365

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

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

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

            if raw_path == -1:
                # python-only task or generated at runtime
                display_path = "---"
                output_type = ""
383
            else:
384
385
386
387
388
389
390
391
392
393
394
395
                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])

396
397
        subtask_table.value_matrix = st_values

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

        return "".join(parts)
411

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

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

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

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

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

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

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

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

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

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

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

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

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

468
469
470
471
472
    def _register_task(
        self,
        task_name: str,
        task_type: str,
        yaml_path: str,
Baber's avatar
nit  
Baber committed
473
474
        tasks_and_groups: dict[str, dict],
        config: Optional[dict] = None,
Baber's avatar
Baber committed
475
476
        populate_tags_fn: Optional[callable] = None,
    ) -> None:
477
478
479
480
481
482
483
        """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:
484
            populate_tags_fn(config, task_name, tasks_and_groups)
485
486

    def _merge_task_configs(
Baber's avatar
nit  
Baber committed
487
488
        self, base_config: dict, task_specific_config: dict, task_name: str
    ) -> dict:
489
490
491
492
493
494
495
        """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
496
    def _process_tag_subtasks(
Baber's avatar
nit  
Baber committed
497
498
        self, tag_name: str, update_config: Optional[dict] = None
    ) -> dict:
499
500
501
502
503
504
505
506
        """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
nit  
Baber committed
507
    def _process_alias(self, config: dict, group: Optional[str] = None) -> dict:
508
509
510
511
512
513
514
515
        # 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
516
    def _class_has_config_in_constructor(self, cls) -> bool:
517
518
519
520
521
522
523
        constructor = getattr(cls, "__init__", None)
        return (
            "config" in inspect.signature(constructor).parameters
            if constructor
            else False
        )

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

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

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

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

552
            if self._config_is_python_task(config):
553
554
555
556
557
558
                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?
559
                    task_object.config.task = task
560
            else:
Baber Abbasi's avatar
Baber Abbasi committed
561
562
563
564
                if self.metadata is not None:
                    config["metadata"] = config.get("metadata", {}) | self.metadata
                else:
                    config["metadata"] = config.get("metadata", {})
565
                task_object = ConfigurableTask(config=config)
Lintang Sutawika's avatar
Lintang Sutawika committed
566

567
568
            return {task: task_object}

Baber Abbasi's avatar
Baber Abbasi committed
569
        def _get_group_and_subtask_from_config(
Baber's avatar
nit  
Baber committed
570
571
            config: dict,
        ) -> tuple[ConfigurableGroup, list[str]]:
Baber Abbasi's avatar
Baber Abbasi committed
572
573
            if self.metadata is not None:
                config["metadata"] = config.get("metadata", {}) | self.metadata
Lintang Sutawika's avatar
Lintang Sutawika committed
574
575
576
577
578
579
580
581
582
            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
583
        def _process_group_config(
Baber's avatar
nit  
Baber committed
584
585
            config: dict, update_config: Optional[dict] = None
        ) -> tuple[dict, Optional[dict]]:
Lintang Sutawika's avatar
Lintang Sutawika committed
586
587
588
589
590
591
592
593
594
595
596
            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

597
598
599
600
        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
601
602
603
            elif self._name_is_task(name_or_config) or self._name_is_python_task(
                name_or_config
            ):
604
605
                # Get the yaml_path for this task
                yaml_path = self._get_yaml_path(name_or_config)
606
                task_config = self._get_config(name_or_config)
607
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

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

636
                return _load_task(task_config, task=name_or_config, yaml_path=yaml_path)
637
            else:
638
639
640
                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
641
642
643
644
645
646
                    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):
647
648
649
                        return self._process_tag_subtasks(
                            name_or_config,
                            name_or_config
Lintang Sutawika's avatar
Lintang Sutawika committed
650
651
652
653
654
655
656
                            if isinstance(name_or_config, dict)
                            else None,
                        )
                    else:
                        group_name = ConfigurableGroup(
                            config={"group": name_or_config, "task": subtask_list}
                        )
657

658
659
        if isinstance(name_or_config, dict):
            if self._config_is_task(name_or_config):
Lintang Sutawika's avatar
Lintang Sutawika committed
660
661
662
                name = name_or_config.pop("task")
                if update_config is not None:
                    name_or_config = {**name_or_config, **update_config}
663
664
                # If the name is registered as a group
                if self._name_is_group(name):
Lintang Sutawika's avatar
Lintang Sutawika committed
665
666
667
668
669
670
671
672
673
                    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):
674
                    return self._process_tag_subtasks(name, name_or_config)
675
                else:
676
                    yaml_path = None
677
                    if self._name_is_registered(name):
678
                        yaml_path = self._get_yaml_path(name)
679
680
681
682
                        base_task_config = self._get_config(name)

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

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

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

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

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

725
        :return
Baber's avatar
nit  
Baber committed
726
            dictionary of task objects
727
728
729
        """
        if isinstance(task_list, str):
            task_list = [task_list]
730

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

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

Baber's avatar
nit  
Baber committed
744
    def _get_task_and_group(self, task_dir: Union[str, Path]) -> dict[str, dict]:
745
        """Creates a dictionary of tasks index with the following metadata,
Lintang Sutawika's avatar
Lintang Sutawika committed
746
        - `type`, that can be either `task`, `python_task`, `group` or `tags`.
747
748
            `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
749
750
            `group` are group configs. `tags` are labels that can be assigned
            to tasks to assist in sorting and calling tasks of certain themes.
751
752
753
754
755
756
757
758
759
760
761
762
763
        - `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
Baber's avatar
nit  
Baber committed
764
            dictionary of task names as key and task metadata
765
        """
766

Baber's avatar
Baber committed
767
        def _populate_tags_and_groups(
Baber's avatar
nit  
Baber committed
768
            config: dict, task: str, tasks_and_groups: dict[str, dict]
Baber's avatar
Baber committed
769
        ) -> None:
770
            # TODO: remove group in next release
771
772
773
774
775
776
777
778
779
780
781
782
783
            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
784
                        eval_logger.info(
785
786
                            f"The tag '{tag}' is already registered as a group, this tag will not be registered. "
                            "This may affect tasks you want to call."
787
                        )
788
789
790
                        break
                    else:
                        tasks_and_groups[tag]["task"].append(task)
791

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

800
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
        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
830

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
                # # 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"]
856
                        self._register_task(
857
                            task_name,
858
                            "task",
Baber's avatar
Baber committed
859
                            str(yaml_path),
860
861
862
                            tasks_and_groups,
                            config,
                            _populate_tags_and_groups,
863
                        )
864
865
            else:
                eval_logger.debug(f"File {yaml_path} could not be loaded")
866
867

        return tasks_and_groups
lintangsutawika's avatar
lintangsutawika committed
868

869

Baber's avatar
nit  
Baber committed
870
def get_task_name_from_config(task_config: dict[str, str]) -> str:
871
872
873
874
875
876
    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
877

878

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

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

891

Baber's avatar
nit  
Baber committed
892
def _check_duplicates(task_dict: dict[str, list[str]]) -> None:
Lintang Sutawika's avatar
Lintang Sutawika committed
893
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
    """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."
        )


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

Baber's avatar
nit  
Baber committed
925
    :param task_name_list: list[Union[str, dict, Task]]
926
927
928
929
930
931
        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`
932

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

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

942
    if isinstance(task_name_list, str):
lintangsutawika's avatar
lintangsutawika committed
943
        task_name_list = [task_name_list]
944
945
946
947
948
949
950
951
952
    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
953

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

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

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

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

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

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

    # 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