__init__.py 841 Bytes
Newer Older
Jason Phang's avatar
gpt3  
Jason Phang committed
1
2
import importlib
import os
Jason Phang's avatar
lib  
Jason Phang committed
3
from lm_eval.base import Registry
Jason Phang's avatar
gpt3  
Jason Phang committed
4
5
6
7
8
9
10
11
12
13
14
15

TASK_REGISTRY = Registry(registry_name="tasks")
# Load all modules in models directory to populate registry
tasks_dir = os.path.dirname(__file__)
for file in os.listdir(tasks_dir):
    path = os.path.join(tasks_dir, file)
    if (
        not file.startswith('_')
        and not file.startswith('.')
        and (file.endswith('.py') or os.path.isdir(path))
    ):
        module_name = file[:file.find('.py')] if file.endswith('.py') else file
Jason Phang's avatar
lib  
Jason Phang committed
16
        module = importlib.import_module('lm_eval.tasks.' + module_name)
Jason Phang's avatar
gpt3  
Jason Phang committed
17
18


Jason Phang's avatar
Jason Phang committed
19
20
21
ALL_TASKS = sorted(list(TASK_REGISTRY.registry))


Jason Phang's avatar
cleanup  
Jason Phang committed
22
23
24
25
26
27
28
29
30
def get_task(task_name):
    return TASK_REGISTRY.registry[task_name]


def get_task_dict(task_name_list):
    return {
        task_name: get_task(task_name)()
        for task_name in task_name_list
    }