__init__.py 37 KB
Newer Older
1
import collections
2
3
import functools
import importlib.util
4
import inspect
5
import logging
6
7
import re
import sys
8
from functools import partial
9
from glob import iglob
Baber's avatar
Baber committed
10
from pathlib import Path
11
12
13
14
15
from typing import Any, Callable, Dict, Generator, List, Mapping, Optional, Union

import yaml
from jinja2 import BaseLoader, Environment, StrictUndefined
from yaml import YAMLError
&'s avatar
& committed
16

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


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

lintangsutawika's avatar
lintangsutawika committed
26

Lintang Sutawika's avatar
Lintang Sutawika committed
27
eval_logger = logging.getLogger(__name__)
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
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
_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 regex_replace(string, pattern, repl, count: int = 0):
    """Implements the `re.sub` function as a custom Jinja filter."""
    return re.sub(pattern, repl, string, count=count)


@functools.lru_cache(maxsize=256)
def _compile_tpl(src: str):
    return apply_template._env.from_string(src)


def apply_template(template: str, doc: dict) -> str:
    if not hasattr(apply_template, "_env"):
        apply_template._env = Environment(
            loader=BaseLoader(),
            undefined=StrictUndefined,
            keep_trailing_newline=True,
        )
        apply_template._env.filters["regex_replace"] = regex_replace

    return _compile_tpl(template).render(**doc)


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
209

210

211
212
213
class TaskManager:
    """TaskManager indexes all tasks from the default `lm_eval/tasks/`
    and an optional directory if provided.
214

215
216
    """

217
218
    def __init__(
        self,
Lintang Sutawika's avatar
Lintang Sutawika committed
219
        verbosity: Optional[str] = None,
Baber's avatar
Baber committed
220
        include_path: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
221
        include_defaults: bool = True,
Baber Abbasi's avatar
Baber Abbasi committed
222
        metadata: Optional[dict] = None,
223
    ) -> None:
Lintang Sutawika's avatar
Lintang Sutawika committed
224
225
        if verbosity is not None:
            utils.setup_logging(verbosity)
226
        self.include_path = include_path
Baber Abbasi's avatar
Baber Abbasi committed
227
        self.metadata = metadata
228
229
230
        self._task_index = self.initialize_tasks(
            include_path=include_path, include_defaults=include_defaults
        )
231
        self._all_tasks = sorted(list(self._task_index.keys()))
232

233
234
235
236
        self._all_groups = sorted(
            [x for x in self._all_tasks if self._task_index[x]["type"] == "group"]
        )
        self._all_subtasks = sorted(
237
238
239
240
241
            [
                x
                for x in self._all_tasks
                if self._task_index[x]["type"] in ["task", "python_task"]
            ]
242
243
244
245
246
        )
        self._all_tags = sorted(
            [x for x in self._all_tasks if self._task_index[x]["type"] == "tag"]
        )

247
        self.task_group_map = collections.defaultdict(list)
248

249
250
    def initialize_tasks(
        self,
Baber's avatar
Baber committed
251
        include_path: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
252
        include_defaults: bool = True,
Baber Abbasi's avatar
Baber Abbasi committed
253
254
    ) -> dict[str, dict]:
        """Creates a dictionary of tasks indexes.
255

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

274
275
276
277
        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
278

279
280
281
        return task_index

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

285
    @property
Baber's avatar
Baber committed
286
    def all_groups(self) -> List[str]:
287
288
289
        return self._all_groups

    @property
Baber's avatar
Baber committed
290
    def all_subtasks(self) -> List[str]:
291
292
293
        return self._all_subtasks

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

297
    @property
Baber's avatar
Baber committed
298
    def task_index(self) -> Dict[str, Dict[str, Union[str, int, List[str]]]]:
299
300
        return self._task_index

301
    def list_all_tasks(
Baber's avatar
Baber committed
302
303
304
305
        self,
        list_groups: bool = True,
        list_tags: bool = True,
        list_subtasks: bool = True,
306
    ) -> str:
307
308
309
310
311
        """
        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.
        """
312
313
        from pytablewriter import MarkdownTableWriter

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

355
        # ---------------------------------------------------------------- TAG table
356
357
358
359
        tag_table = MarkdownTableWriter()
        tag_table.headers = ["Tag"]
        tag_table.value_matrix = [[t] for t in self.all_tags]

360
        # ------------------------------------------------------------ SUBTASK table
361
362
        subtask_table = MarkdownTableWriter()
        subtask_table.headers = ["Task", "Config Location", "Output Type"]
363
364
        st_values: list[list[str]] = []

365
        for t in self.all_subtasks:
366
367
368
369
370
371
            raw_path = self.task_index[t]["yaml_path"]

            if raw_path == -1:
                # python-only task or generated at runtime
                display_path = "---"
                output_type = ""
372
            else:
373
374
375
376
377
378
379
380
381
382
383
384
                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])

385
386
        subtask_table.value_matrix = st_values

387
388
        # ------------------------------------------------------------- final string
        parts: list[str] = ["\n"]
389
        if list_groups:
390
391
            parts.append(group_table.dumps())
            parts.append("\n")
392
        if list_tags:
393
394
            parts.append(tag_table.dumps())
            parts.append("\n")
395
        if list_subtasks:
396
397
398
399
            parts.append(subtask_table.dumps())
            parts.append("\n")

        return "".join(parts)
400

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

Baber Abbasi's avatar
Baber Abbasi committed
404
    def _name_is_registered(self, name: str) -> bool:
405
        return name in self.all_tasks
406

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

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

Baber Abbasi's avatar
Baber Abbasi committed
415
    def _name_is_group(self, name: str) -> bool:
416
417
418
        return (
            self._name_is_registered(name) and self.task_index[name]["type"] == "group"
        )
419

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

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

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

Baber Abbasi's avatar
Baber Abbasi committed
432
    def _config_is_python_task(self, config: dict) -> bool:
433
434
435
436
        return "class" in config

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

Baber's avatar
Baber committed
438
    def _get_yaml_path(self, name: str) -> Union[str, int]:
439
440
        if name not in self.task_index:
            raise ValueError
441
442
        return self.task_index[name]["yaml_path"]

Baber's avatar
Baber committed
443
    def _get_config(self, name: str) -> Dict:
444
445
        if name not in self.task_index:
            raise ValueError
446
447
448
449
        yaml_path = self._get_yaml_path(name)
        if yaml_path == -1:
            return {}
        else:
450
            return load_yaml_config(Path(yaml_path), mode="full")
451

Baber's avatar
Baber committed
452
    def _get_tasklist(self, name: str) -> Union[List[str], int]:
453
454
        if self._name_is_task(name):
            raise ValueError
455
456
        return self.task_index[name]["task"]

457
458
459
460
461
    def _register_task(
        self,
        task_name: str,
        task_type: str,
        yaml_path: str,
Baber's avatar
Baber committed
462
463
464
465
        tasks_and_groups: Dict[str, Dict],
        config: Optional[Dict] = None,
        populate_tags_fn: Optional[callable] = None,
    ) -> None:
466
467
468
469
470
471
472
        """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:
473
            populate_tags_fn(config, task_name, tasks_and_groups)
474
475

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

513
    def _load_individual_task_or_group(
514
        self,
Baber's avatar
Baber committed
515
        name_or_config: Optional[Union[str, Dict]] = None,
516
        parent_name: Optional[str] = None,
Baber's avatar
Baber committed
517
        update_config: Optional[Dict] = None,
518
    ) -> Mapping:
Baber's avatar
Baber committed
519
520
521
        def _load_task(
            config: Dict, task: str, yaml_path: Optional[str] = None
        ) -> Dict[str, Union[ConfigurableTask, Task]]:
522
            if "include" in config:
523
524
525
                # Store the task name to preserve it after include processing
                original_task_name = config.get("task", task)

526
                config = {
527
528
                    **load_yaml_config(
                        yaml_path=Path(yaml_path),
529
                        yaml_config={"include": config.pop("include")},
530
                        mode="full" if yaml_path else "simple",
531
532
                    ),
                    **config,
533
                    "task": original_task_name,
534
                }
535
536
537
538

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

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

554
555
            return {task: task_object}

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

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

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

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

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

                        # Check if this is a duplicate.
                        if parent_name is not None:
670
671
672
673
674
675
676
677
                            num_duplicate = len(
                                list(
                                    filter(
                                        lambda x: x.startswith(name),
                                        self.task_group_map[parent_name],
                                    )
                                )
                            )
678
679
680
681
                            if num_duplicate > 0:
                                name = f"{name}-{num_duplicate}"
                            self.task_group_map[parent_name].append(name)

682
683
684
685
                        task_config = {
                            **base_task_config,
                            **name_or_config,
                        }
686
687
                    else:
                        task_config = name_or_config
688
                    return _load_task(task_config, task=name, yaml_path=yaml_path)
689
            else:
Lintang Sutawika's avatar
Lintang Sutawika committed
690
691
692
693
                group_config, update_config = _process_group_config(name_or_config)
                group_name, subtask_list = _get_group_and_subtask_from_config(
                    group_config
                )
694

695
696
697
698
699
        fn = partial(
            self._load_individual_task_or_group,
            parent_name=group_name,
            update_config=update_config,
        )
Lintang Sutawika's avatar
Lintang Sutawika committed
700
701
        return {
            group_name: dict(collections.ChainMap(*map(fn, reversed(subtask_list))))
702
        }
703

Baber's avatar
Baber committed
704
705
706
    def load_task_or_group(
        self, task_list: Optional[Union[str, List[str]]] = None
    ) -> Dict:
707
        """Loads a dictionary of task objects from a list
708

709
710
        :param task_list: Union[str, list] = None
            Single string or list of string of task names to be loaded
711

712
713
714
715
716
        :return
            Dictionary of task objects
        """
        if isinstance(task_list, str):
            task_list = [task_list]
717

718
        all_loaded_tasks = dict(
Baber Abbasi's avatar
Baber Abbasi committed
719
720
721
722
723
724
            collections.ChainMap(
                *map(
                    lambda task: self._load_individual_task_or_group(task),
                    task_list,
                )
            )
725
726
727
        )
        return all_loaded_tasks

Baber's avatar
Baber committed
728
    def load_config(self, config: Dict) -> Mapping:
729
730
        return self._load_individual_task_or_group(config)

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

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

Lintang Sutawika's avatar
Lintang Sutawika committed
779
        # TODO: remove group in next release
780
781
782
783
        # ignore_dirs = [
        #     "__pycache__",
        #     ".ipynb_checkpoints",
        # ]
784
        tasks_and_groups = collections.defaultdict()
Baber's avatar
Baber committed
785
786
        task_dir_path = Path(task_dir)

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
812
813
814
815
816
        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
817

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

        return tasks_and_groups
lintangsutawika's avatar
lintangsutawika committed
855

856

857
858
859
860
861
862
863
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
864

865

Baber's avatar
Baber committed
866
def get_task_name_from_object(task_object: Union[ConfigurableTask, Task]) -> str:
867
868
    if hasattr(task_object, "config"):
        return task_object._config["task"]
lintangsutawika's avatar
lintangsutawika committed
869
870
871
872
873
874
875
876
877

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

878

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


906
def get_task_dict(
907
    task_name_list: Union[str, List[Union[str, Dict, Task]]],
908
    task_manager: Optional[TaskManager] = None,
Baber's avatar
Baber committed
909
) -> Dict[str, Union[ConfigurableTask, Task]]:
910
    """Creates a dictionary of task objects from either a name of task, config, or prepared Task object.
lintangsutawika's avatar
lintangsutawika committed
911

912
913
914
915
916
917
918
    :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`
919

920
921
922
    :return
        Dictionary of task objects
    """
Lintang Sutawika's avatar
Lintang Sutawika committed
923

924
    task_name_from_string_dict = {}
925
926
927
    task_name_from_config_dict = {}
    task_name_from_object_dict = {}

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

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

948
949
950
        task_name_from_string_dict = task_manager.load_task_or_group(
            string_task_name_list
        )
951

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

        elif isinstance(task_element, Task):
            task_name_from_object_dict = {
                **task_name_from_object_dict,
lintangsutawika's avatar
lintangsutawika committed
962
                get_task_name_from_object(task_element): task_element,
963
            }
lintangsutawika's avatar
lintangsutawika committed
964

965
    if not set(task_name_from_string_dict.keys()).isdisjoint(
lintangsutawika's avatar
lintangsutawika committed
966
        set(task_name_from_object_dict.keys())
967
968
    ):
        raise ValueError
969

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

    # 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