__init__.py 5.46 KB
Newer Older
1
import os
lintangsutawika's avatar
lintangsutawika committed
2
from typing import List, Union
&'s avatar
& committed
3

4
5
from .gsm8k import *
from .triviaqa import *
haileyschoelkopf's avatar
haileyschoelkopf committed
6
from .glue import *
lintangsutawika's avatar
lintangsutawika committed
7

8
from lm_eval import utils
lintangsutawika's avatar
lintangsutawika committed
9
from lm_eval.logger import eval_logger
10
from lm_eval.api.task import TaskConfig, Task, ConfigurableTask
lintangsutawika's avatar
lintangsutawika committed
11
from lm_eval.api.registry import (
12
13
    register_task,
    register_group,
lintangsutawika's avatar
lintangsutawika committed
14
15
    TASK_REGISTRY,
    GROUP_REGISTRY,
haileyschoelkopf's avatar
haileyschoelkopf committed
16
    ALL_TASKS,
lintangsutawika's avatar
lintangsutawika committed
17
)
lintangsutawika's avatar
lintangsutawika committed
18

19

20
def get_task_name_from_config(task_config):
21
    return "{dataset_path}_{dataset_name}".format(**task_config)
22
23


24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def include_task_folder(task_dir):
    """
    Calling this function
    """
    for root, subdirs, file_list in os.walk(task_dir):
        if (subdirs == []) and (len(file_list) > 0):
            for f in file_list:
                if f.endswith(".yaml"):
                    yaml_path = os.path.join(root, f)
                    try:
                        config = utils.load_yaml_config(yaml_path)

                        SubClass = type(
                            config["task"] + "ConfigurableTask",
                            (ConfigurableTask,),
                            {"CONFIG": TaskConfig(**config)},
                        )
41

42
                        if "task" in config:
Lintang Sutawika's avatar
Lintang Sutawika committed
43
44
45
                            # task_name = "{}:{}".format(
                            #     get_task_name_from_config(config), config["task"]
                            # )
46
47
                            task_name = "{}".format(config["task"])
                            register_task(task_name)(SubClass)
48

49
50
51
52
53
54
55
56
57
58
                        if "group" in config:
                            for group in config["group"]:
                                register_group(group)(SubClass)
                    except Exception as error:
                        eval_logger.warning(
                            "Failed to load config in\n"
                            f"                                 {yaml_path}\n"
                            "                                 Config will not be added to registry\n"
                            f"                                 Error: {error}"
                        )
59

60

61
62
task_dir = os.path.dirname(os.path.abspath(__file__)) + "/"
include_task_folder(task_dir)
lintangsutawika's avatar
lintangsutawika committed
63

lintangsutawika's avatar
lintangsutawika committed
64

65
66
def get_task(task_name, config):
    try:
67
        return TASK_REGISTRY[task_name](config=config)
68
    except KeyError:
lintangsutawika's avatar
lintangsutawika committed
69
        eval_logger.info("Available tasks:")
70
        eval_logger.info(list(TASK_REGISTRY) + list(GROUP_REGISTRY))
lintangsutawika's avatar
lintangsutawika committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
        raise KeyError(f"Missing task {task_name}")


def get_task_name_from_object(task_object):
    for name, class_ in TASK_REGISTRY.items():
        if class_ is task_object:
            return name

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


# TODO: pass num_fewshot and other cmdline overrides in a better way
lintangsutawika's avatar
lintangsutawika committed
89
90
91
def get_task_dict(task_name_list: List[Union[str, dict, Task]], **kwargs):

    config = {**kwargs}
92
93
94
95
96
97
98
99
100
101
102
103
104

    task_name_from_registry_dict = {}
    task_name_from_config_dict = {}
    task_name_from_object_dict = {}

    for task_element in task_name_list:
        if isinstance(task_element, str):

            if task_element in GROUP_REGISTRY:
                for task_name in GROUP_REGISTRY[task_element]:
                    if task_name not in task_name_from_registry_dict:
                        task_name_from_registry_dict = {
                            **task_name_from_registry_dict,
lintangsutawika's avatar
lintangsutawika committed
105
106
                            task_name: get_task(task_name=task_name, config=config),
                        }
107
            else:
108
                task_name = task_element
109
110
111
                if task_name not in task_name_from_registry_dict:
                    task_name_from_registry_dict = {
                        **task_name_from_registry_dict,
lintangsutawika's avatar
lintangsutawika committed
112
113
                        task_name: get_task(task_name=task_element, config=config),
                    }
114
115

        elif isinstance(task_element, dict):
116
            task_element.update(config)
117
118
119
            task_name_from_config_dict = {
                **task_name_from_config_dict,
                get_task_name_from_config(task_element): ConfigurableTask(
120
                    config=task_element
lintangsutawika's avatar
lintangsutawika committed
121
                ),
122
123
124
125
126
127
            }

        elif isinstance(task_element, Task):

            task_name_from_object_dict = {
                **task_name_from_object_dict,
lintangsutawika's avatar
lintangsutawika committed
128
                get_task_name_from_object(task_element): task_element,
129
            }
lintangsutawika's avatar
lintangsutawika committed
130

131
132
133
134
135
    # task_name_from_registry_dict = {
    #     task_name: get_task(
    #         task_name=task_name,
    #         task_config=config
    #     )
lintangsutawika's avatar
lintangsutawika committed
136
    #     for group_name in task_name_list for task_name in GROUP_REGISTRY[group_name]
137
138
139
140
141
142
143
144
145
146
    #     if (isinstance(group_name, str)) and (group_name in GROUP_REGISTRY)
    # }
    # task_name_from_config_dict = {
    #     get_task_name_from_config(task_config): ConfigurableTask(
    #         config=task_config
    #     )
    #     for task_config in task_name_list
    #     if isinstance(task_config, dict)
    # }
    # # TODO: Do we still need this?
lintangsutawika's avatar
lintangsutawika committed
147
148
149
150
151
    # task_name_from_object_dict = {
    #     get_task_name_from_object(task_object): task_object
    #     for task_object in task_name_list
    #     if isinstance(task_object, Task)
    # }
152

lintangsutawika's avatar
lintangsutawika committed
153
154
155
    assert set(task_name_from_registry_dict.keys()).isdisjoint(
        set(task_name_from_object_dict.keys())
    )
lintangsutawika's avatar
lintangsutawika committed
156
157
158
    return {
        **task_name_from_registry_dict,
        **task_name_from_config_dict,
159
        **task_name_from_object_dict,
lintangsutawika's avatar
lintangsutawika committed
160
    }