__init__.py

import os
import abc
import yaml
import collections

from functools import partial, lru_cache
from typing import List, Union, Dict

from lm_eval import utils
from lm_eval import prompts
from lm_eval.api.task import TaskConfig, Task, ConfigurableTask
from lm_eval.api.registry import (
    register_task,
    register_group,
    TASK_REGISTRY,
    GROUP_REGISTRY,
)

import logging

# import python tasks
from .squadv2.task import SQuAD2
from .scrolls.task import (
    QuALITY,
    NarrativeQA,
    ContractNLI,
    GovReport,
    SummScreenFD,
    QMSum,
)

eval_logger = utils.eval_logger


class TaskManager(abc.ABC):

    def __init__(
        self,
        verbosity="INFO",
        include_path=None
        ) -> None:

        self.verbosity = verbosity
        self.include_path = include_path
        self.logger = eval_logger.setLevel(getattr(logging, f"{verbosity}"))

        self.ALL_TASKS = self.initialize_tasks(
            include_path=include_path
            )

    def initialize_tasks(self, include_path=None):

        all_paths = [os.path.dirname(os.path.abspath(__file__)) + "/"]
        if include_path is not None:
            if isinstance(include_path, str):
                include_path = [include_path]
            all_paths.extend(include_path)

        ALL_TASKS = {}
        for task_dir in all_paths:
            tasks = self._get_task_and_group(task_dir)
            ALL_TASKS = {**tasks, **ALL_TASKS}

        return ALL_TASKS

    def all_tasks(self):
        return sorted(list(self.ALL_TASKS.keys()))

    def _name_is_registered(self, name):
        if name in self.ALL_TASKS:
            return True
        return False

    def _name_is_task(self, name):
        if self.ALL_TASKS[name]["type"] == "task":
            return True
        return False

    def _config_is_task(self, config):
        if list(config.keys()) == ["group", "task"]:
            return False
        return True

    def _get_yaml_path(self, name):
        assert name in self.ALL_TASKS
        return self.ALL_TASKS[name]["yaml_path"]

    def _get_config(self, name):
        assert name in self.ALL_TASKS
        yaml_path = self._get_yaml_path(name)
        return utils.load_yaml_config(yaml_path)

    def _get_tasklist(self, name):
        assert self._name_is_task(name) == False
        return self.ALL_TASKS[name]["task"]

    @lru_cache(None)
    def _load_individual_task_or_group(self, name_or_config: Union[str, dict] = None, parent_name: str = None) -> ConfigurableTask:

        def load_task(config, task, group=None):
            task_object = ConfigurableTask(config=config)
            if group is not None:
                task_object = (group, task_object)
            return {task: task_object}

        if isinstance(name_or_config, str):
            if self._name_is_task(name_or_config):
                task_config = self._get_config(name_or_config)
                return load_task(task_config, task=name_or_config, group=parent_name)
            else:
                group_name = name_or_config
                subtask_list = self._get_tasklist(name_or_config)
                if subtask_list == -1:
                    subtask_list = self._get_config(name_or_config)["task"]

        elif isinstance(name_or_config, dict):
            if self._config_is_task(name_or_config):
                task_name = name_or_config["task"]
                if self._name_is_registered(task_name):
                    base_task_config = self._get_config(task_name)
                    task_config={
                            **base_task_config,
                            **name_or_config,
                        }
                else:
                    task_config = name_or_config
                return load_task(task_config, task=name_or_config, group=parent_name)
            else:
                group_name = name_or_config["group"]
                subtask_list = name_or_config["task"]

        if self._get_yaml_path(group_name) == -1:
            all_subtasks = {group_name: (parent_name, None)}
        else:
            all_subtasks = {}

        fn = partial(self._load_individual_task_or_group, parent_name=group_name)
        all_subtasks = {**all_subtasks, **dict(collections.ChainMap(*map(fn, subtask_list)))}
        return all_subtasks


    def load_task_or_group(self, task_list: Union[str, list] = None) -> dict:

        if isinstance(task_list, str):
            task_list = [task_list]

        all_loaded_tasks = dict(
            collections.ChainMap(
                *map(
                    self._load_individual_task_or_group,
                    task_list
                )
            )
        )
        return all_loaded_tasks

    def _get_task_and_group(self, task_dir: str):
        tasks_and_groups = collections.defaultdict()
        for root, _, file_list in os.walk(task_dir):
            for f in file_list:
                if f.endswith(".yaml"):
                    yaml_path = os.path.join(root, f)
                    config = utils.simple_load_yaml_config(yaml_path)
                    if list(config.keys()) == ["group", "task"]:
                        # 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": yaml_path,
                        }
                    else:
                        # This is a task config
                        task = config["task"]
                        tasks_and_groups[task] = {
                            "type": "task",
                            "yaml_path": yaml_path,
                            }

                        if "group" in config:
                            groups = config["group"]
                            if isinstance(config["group"], str):
                                groups = [groups]

                            for group in groups:
                                if group not in tasks_and_groups:
                                    tasks_and_groups[group] = {
                                        "type": "group",
                                        "task": [task],
                                        "yaml_path": -1,
                                    }
                                else:
                                    tasks_and_groups[group]["task"].append(task)

        return tasks_and_groups


def check_prompt_config(
    config: Dict[str, str], yaml_path: str = None
) -> List[Dict[str, str]]:
    all_configs = []
    if "use_prompt" in config:
        prompt_list = prompts.load_prompt_list(
            use_prompt=config["use_prompt"],
            dataset_name=config["dataset_path"],
            subset_name=config["dataset_name"] if "dataset_name" in config else None,
            yaml_path=yaml_path,
        )
        for idx, prompt_variation in enumerate(prompt_list):
            all_configs.append(
                {
                    **config,
                    **{"use_prompt": prompt_variation},
                    **{
                        "task": "_".join(
                            [
                                config["task"]
                                if "task" in config
                                else get_task_name_from_config(config),
                                prompt_variation.split("/")[-1]
                                if ".yaml" in prompt_variation
                                else prompt_variation,
                            ]
                        )
                    },
                    **{"output_type": "generate_until"},
                }
            )
    else:
        all_configs.append(config)
    return all_configs

# # TODO: pass num_fewshot and other cmdline overrides in a better way
# def get_task_dict(task_name_list: List[Union[str, Dict, Task]], **kwargs):
#     config = {**kwargs}

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

#     if type(task_name_list) != list:
#         task_name_list = [task_name_list]

#     for task_element in task_name_list:
#         if isinstance(task_element, str):
#             if task_element in GROUP_REGISTRY:
#                 group_name = task_element
#                 for task_name in GROUP_REGISTRY[task_element]:
#                     if task_name not in task_name_from_registry_dict:
#                         task_obj = get_task_dict(task_name)
#                         if task_name in task_obj.keys():
#                             task_dict = {
#                                 task_name: (group_name, task_obj[task_name]),
#                             }
#                         else:
#                             task_dict = {
#                                 task_name: (group_name, None),
#                                 **task_obj,
#                             }

#                         task_name_from_registry_dict = {
#                             **task_name_from_registry_dict,
#                             **task_dict,
#                         }
#             else:
#                 task_name = task_element
#                 if task_name not in task_name_from_registry_dict:
#                     task_name_from_registry_dict = {
#                         **task_name_from_registry_dict,
#                         task_name: get_task(task_name=task_element, config=config),
#                     }

#         elif isinstance(task_element, dict):
#             task_element.update(config)
#             task_name_from_config_dict = {
#                 **task_name_from_config_dict,
#                 get_task_name_from_config(task_element): ConfigurableTask(
#                     config=task_element
#                 ),
#             }

#         elif isinstance(task_element, Task):
#             task_name_from_object_dict = {
#                 **task_name_from_object_dict,
#                 get_task_name_from_object(task_element): task_element,
#             }

#     assert set(task_name_from_registry_dict.keys()).isdisjoint(
#         set(task_name_from_object_dict.keys())
#     )
#     return {
#         **task_name_from_registry_dict,
#         **task_name_from_config_dict,
#         **task_name_from_object_dict,
#     }