eval-framework 0.2.0__py3-none-any.whl

eval_framework/tasks/registry.py

@@ -0,0 +1,186 @@
+import contextlib
+import importlib.util
+import re
+from collections.abc import Generator, Iterator, Sequence
+from typing import Annotated, Any
+
+import pydantic
+from pydantic import AfterValidator
+
+from eval_framework.tasks.base import BaseTask
+from eval_framework.utils.packaging import is_extra_installed, validate_package_extras
+
+__all__ = [
+    "register_task",
+    "register_lazy_task",
+    "Registry",
+    "with_registry",
+    "get_task",
+    "registered_tasks_iter",
+    "is_registered",
+    "validate_task_name",
+    "registered_task_names",
+]
+
+
+def validate_import_path(import_path: str) -> str:
+    if importlib.util.find_spec(import_path) is None:
+        raise ValueError(f"Invalid import path: {import_path}")
+    return import_path
+
+
+class TaskPlaceholder(pydantic.BaseModel, extra="forbid", frozen=True):
+    name: Annotated[
+        str,
+        "The name of the Task class that we want to import",
+    ]
+    module: Annotated[
+        str,
+        "The module from where to import the task",
+        validate_import_path,
+    ]
+    extras: Annotated[
+        tuple[str, ...],
+        "Extra dependencies that are required for the task",
+        AfterValidator(validate_package_extras),
+    ] = ()
+
+    def load(self) -> type[BaseTask]:
+        for extra in self.extras:
+            if not is_extra_installed(extra):
+                raise ImportError(f"The required package eval_framework[{extra}] is not installed.")
+        module = importlib.import_module(self.module)
+        return getattr(module, self.name)
+
+
+class Registry:
+    """A registry for tasks with support for lazy loading.
+
+    Task names are hashed based on the upper-case name, to avoid issues with
+    ambiguous naming.
+    """
+
+    def __init__(self) -> None:
+        # TODO: Lookup only with upper names
+        self._registry: dict[str, tuple[str, type[BaseTask] | TaskPlaceholder]] = dict()
+
+    def __iter__(self) -> Iterator[str]:
+        for name, _ in self._registry.values():
+            yield name
+
+    @staticmethod
+    def _task_key(name: str, /) -> str:
+        name = re.sub(r"[\s\-_]+", "", name).upper()
+        if not name.isalnum():
+            raise ValueError(
+                f"Task name '{name}' contains invalid characters. Only alphanumeric characters are allowed."
+            )
+        return name
+
+    def __contains__(self, name: str) -> bool:
+        task_key = self._task_key(name)
+        return task_key in self._registry
+
+    def __getitem__(self, name: str, /) -> type[BaseTask]:
+        task_key = self._task_key(name)
+        try:
+            name, task = self._registry[task_key]
+        except KeyError:
+            raise KeyError(f"Task not found: {name}")
+
+        if isinstance(task, TaskPlaceholder):
+            task = task.load()
+            self._registry[task_key] = (name, task)
+        return task
+
+    def add(self, task: type[BaseTask]) -> None:
+        task_key = self._task_key(task.NAME)
+        self._registry[task_key] = (task.NAME, task)
+
+    def __setitem__(self, name: str, task: type[BaseTask] | TaskPlaceholder) -> None:
+        task_key = self._task_key(name)
+        if task_key in self._registry:
+            raise ValueError(f"Cannot register duplicate task with key: {task_key}")
+
+        self._registry[task_key] = (name, task)
+
+
+_REGISTRY = Registry()
+
+
+@contextlib.contextmanager
+def with_registry(registry: Registry) -> Generator[None, Any, None]:
+    """Contextmanager to change the current registry."""
+    global _REGISTRY
+    old_registry = _REGISTRY
+    try:
+        _REGISTRY = registry
+        yield
+    finally:
+        _REGISTRY = old_registry
+
+
+def registered_task_names() -> list[str]:
+    """Return the names of all registered tasks."""
+    return list(_REGISTRY)
+
+
+def is_registered(name: str, /) -> bool:
+    """Return True if a task is registered."""
+    return name in _REGISTRY
+
+
+def validate_task_name(name: str) -> str:
+    """Pydantic-style validator for task names."""
+    if not is_registered(name):
+        raise ValueError(f"Task not registered: {name}")
+    return name
+
+
+def registered_tasks_iter() -> Iterator[tuple[str, type[BaseTask]]]:
+    """Iterate over the names and classes of all registered tasks.
+
+    Note: This method will import any lazily registered task.
+    """
+    for name in registered_task_names():
+        yield name, get_task(name)
+
+
+def get_task(name: str, /) -> type[BaseTask]:
+    """Return a registered task for a given name.
+
+    Note: This method will import any lazily registered task.
+    """
+    return _REGISTRY[name]
+
+
+def register_task(task: type[BaseTask]) -> str:
+    """The class name is used as the task name."""
+    if not issubclass(task, BaseTask):
+        raise ValueError(f"Can only register subclasses of BaseTask, got {task}")
+    name = task.__name__
+    _REGISTRY[name] = task
+    return name
+
+
+def register_lazy_task(class_path: str, /, *, extras: Sequence[str] = ()) -> None:
+    """Register a task without importing it.
+
+    Lazily register a task without importing the module.
+
+    Args:
+        class_path: The full path to the task class. For example,
+            `eval_framework.tasks.benchmarks.truthfulqa.TRUTHFULQA`.
+        extras: Any extra dependencies of `eval_framework` that need to be installed for this task.
+    """
+    if isinstance(extras, str):
+        extras = [extras]
+    if "." not in class_path:
+        raise ValueError(
+            f"Invalid class path `{class_path}`. This needs to be a global path like "
+            "`eval_framework.tasks.benchmarks.truthfulqa.TRUTHFULQA`): "
+        )
+
+    base_module, class_name = class_path.rsplit(".", maxsplit=1)
+    placeholder = TaskPlaceholder(name=class_name, module=base_module, extras=extras)
+    _REGISTRY[class_name] = placeholder

eval_framework/tasks/task_loader.py

@@ -0,0 +1,80 @@
+import importlib.util
+import inspect
+import logging
+import os
+from collections.abc import Sequence
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+from eval_framework.tasks.base import BaseTask
+from eval_framework.tasks.registry import is_registered, register_task
+
+logger = logging.getLogger(__name__)
+
+
+def find_all_python_files(*module_paths: str | os.PathLike) -> set[Path]:
+    """Recursively walk through all paths and return all Python files."""
+    all_files: set[Path] = set()
+    for path in module_paths:
+        path = Path(path).resolve()
+
+        if not path.exists():
+            raise FileNotFoundError(f"[User Task Loader] Path does not exist: {path}")
+        if path.is_dir():
+            all_files.update(path.glob("**/*.py"))
+        elif path.is_file():
+            if path.suffix != ".py":
+                raise ValueError(f"The provided path {path} is not a Python file.")
+            all_files.add(path)
+        else:
+            raise ValueError(f"Path is not a .py file or directory: {path}")
+    return all_files
+
+
+def import_file(f: str | os.PathLike, /) -> Any:
+    """Import a file as a Python module."""
+    file_path = Path(f)
+    try:
+        spec = importlib.util.spec_from_file_location("user_task_module", file_path)
+        if spec is None or spec.loader is None:
+            raise RuntimeError(f"Could not create a module spec for {file_path}")
+        user_module: ModuleType = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(user_module)
+    except Exception as e:
+        raise ImportError(f"Failed to import {file_path}: {e}") from e
+    return user_module
+
+
+def load_extra_tasks(module_paths: Sequence[str | os.PathLike]) -> None:
+    """Dynamically load and register user-defined tasks from a list of files or directories.
+
+    Each .py file found will be imported, and any BaseTask subclass will be registered
+    in the TaskName enum for use by name.
+    Provides clear error messages for missing/invalid files or import errors.
+    """
+    for file_path in find_all_python_files(*module_paths):
+        user_module = import_file(file_path)
+
+        for name, obj in inspect.getmembers(user_module):
+            if inspect.isclass(obj) and issubclass(obj, BaseTask) and obj is not BaseTask:
+                # check unique name of BaseTask subclasses (when they have a NAME attribute)
+                if not hasattr(obj, "NAME"):
+                    logger.info(f"[User Task Loader] Skipping {obj.__module__} - no NAME attribute present.")
+                else:
+                    if is_registered(obj.NAME):
+                        # two classes with the same NAME attribute
+                        logger.info(obj.__module__)
+
+                        # if it comes from eval_framework's built-in tasks then will just skip (no need to register
+                        # again; this can happen if a class is imported so that a new task can be derived from it)
+                        # but if it is one of the user defined task then raise a duplicated name error
+                        if "eval_framework.tasks.benchmarks" not in obj.__module__:
+                            # skip if import comes from eval_framework's built-in tasks
+                            raise ValueError(f"Duplicate user task name found (case-insensitive): {obj.NAME}")
+
+                    else:
+                        # if there is no duplicate name conflict then register the new task
+                        class_obj = getattr(user_module, name)
+                        register_task(class_obj)
+                        logger.info(f"[User Task Loader] Registered task: {class_obj.NAME}")

eval_framework/tasks/task_names.py

@@ -0,0 +1,138 @@
+import logging
+import random
+import time
+from enum import Enum
+
+from eval_framework.tasks.base import BaseTask
+from eval_framework.tasks.registry import register_lazy_task, registered_tasks_iter
+
+logger = logging.getLogger(__name__)
+
+
+class TaskNameEnum(Enum):
+    @property
+    def value(self) -> type[BaseTask]:
+        return super().value
+
+
+def register_all_tasks() -> None:
+    """Register all the benchmark tasks with the eval framework."""
+    register_lazy_task("eval_framework.tasks.benchmarks.math_reasoning.AIME2024")
+    register_lazy_task("eval_framework.tasks.benchmarks.arc.ARC")
+    register_lazy_task("eval_framework.tasks.benchmarks.arc_de.ARC_DE")
+    register_lazy_task("eval_framework.tasks.benchmarks.opengptx_eu20.ARC_EU20_DE")
+    register_lazy_task("eval_framework.tasks.benchmarks.opengptx_eu20.ARC_EU20_FR")
+    register_lazy_task("eval_framework.tasks.benchmarks.arc_fi.ARC_FI")
+    register_lazy_task("eval_framework.tasks.benchmarks.belebele.BELEBELE")
+    register_lazy_task("eval_framework.tasks.benchmarks.bigcodebench.BigCodeBench")
+    register_lazy_task("eval_framework.tasks.benchmarks.bigcodebench.BigCodeBenchInstruct")
+    register_lazy_task("eval_framework.tasks.benchmarks.bigcodebench.BigCodeBenchHard")
+    register_lazy_task("eval_framework.tasks.benchmarks.bigcodebench.BigCodeBenchHardInstruct")
+    register_lazy_task("eval_framework.tasks.benchmarks.casehold.CASEHOLD")
+    register_lazy_task("eval_framework.tasks.benchmarks.chembench.ChemBench")
+    register_lazy_task("eval_framework.tasks.benchmarks.copa.COPA")
+    register_lazy_task("eval_framework.tasks.benchmarks.duc.DUC_ABSTRACTIVE")
+    register_lazy_task("eval_framework.tasks.benchmarks.duc.DUC_EXTRACTIVE")
+    register_lazy_task("eval_framework.tasks.benchmarks.flores200.Flores200")
+    register_lazy_task("eval_framework.tasks.benchmarks.flores_plus.FloresPlus", extras=["comet"])
+    register_lazy_task("eval_framework.tasks.benchmarks.gpqa.GPQA")
+    register_lazy_task("eval_framework.tasks.benchmarks.gpqa.GPQA_COT")
+    register_lazy_task("eval_framework.tasks.benchmarks.gsm8k.GSM8K")
+    register_lazy_task("eval_framework.tasks.benchmarks.gsm8k.GSM8KLlamaVersion")
+    register_lazy_task("eval_framework.tasks.benchmarks.math_reasoning.GSM8KReasoning")
+    register_lazy_task("eval_framework.tasks.benchmarks.opengptx_eu20.GSM8K_EU20_DE")
+    register_lazy_task("eval_framework.tasks.benchmarks.opengptx_eu20.GSM8K_EU20_FR")
+    register_lazy_task("eval_framework.tasks.benchmarks.hellaswag.HELLASWAG")
+    register_lazy_task("eval_framework.tasks.benchmarks.hellaswag_de.HELLASWAG_DE")
+    register_lazy_task("eval_framework.tasks.benchmarks.opengptx_eu20.HELLASWAG_EU20_DE")
+    register_lazy_task("eval_framework.tasks.benchmarks.opengptx_eu20.HELLASWAG_EU20_FR")
+    register_lazy_task("eval_framework.tasks.benchmarks.humaneval.HumanEval")
+    register_lazy_task("eval_framework.tasks.benchmarks.humaneval.HumanEvalInstruct")
+    register_lazy_task("eval_framework.tasks.benchmarks.ifeval.IFEval")
+    register_lazy_task("eval_framework.tasks.benchmarks.ifeval.IFEvalDe")
+    register_lazy_task("eval_framework.tasks.benchmarks.ifeval.IFEvalFiSv")
+    register_lazy_task("eval_framework.tasks.benchmarks.include.INCLUDE")
+    register_lazy_task("eval_framework.tasks.benchmarks.infinitebench.InfiniteBench_CodeDebug")
+    register_lazy_task("eval_framework.tasks.benchmarks.infinitebench.InfiniteBench_CodeRun")
+    register_lazy_task("eval_framework.tasks.benchmarks.infinitebench.InfiniteBench_EnDia")
+    register_lazy_task("eval_framework.tasks.benchmarks.infinitebench.InfiniteBench_EnMC")
+    register_lazy_task("eval_framework.tasks.benchmarks.infinitebench.InfiniteBench_EnQA")
+    register_lazy_task("eval_framework.tasks.benchmarks.infinitebench.InfiniteBench_MathFind")
+    register_lazy_task("eval_framework.tasks.benchmarks.infinitebench.InfiniteBench_RetrieveKV2")
+    register_lazy_task("eval_framework.tasks.benchmarks.infinitebench.InfiniteBench_RetrieveNumber")
+    register_lazy_task("eval_framework.tasks.benchmarks.infinitebench.InfiniteBench_RetrievePassKey1")
+    register_lazy_task("eval_framework.tasks.benchmarks.math_reasoning.MATH")
+    register_lazy_task("eval_framework.tasks.benchmarks.math_reasoning.MATHLvl5")
+    register_lazy_task("eval_framework.tasks.benchmarks.math_reasoning.MATH500")
+    register_lazy_task("eval_framework.tasks.benchmarks.mbpp.MBPP")
+    register_lazy_task("eval_framework.tasks.benchmarks.mbpp.MBPP_SANITIZED")
+    register_lazy_task("eval_framework.tasks.benchmarks.mbpp.MBPP_PROMPT_WITHOUT_TESTS")
+    register_lazy_task("eval_framework.tasks.benchmarks.mbpp.MBPP_PROMPT_WITHOUT_TESTS_SANITIZED")
+    register_lazy_task("eval_framework.tasks.benchmarks.mmlu.MMLU")
+    register_lazy_task("eval_framework.tasks.benchmarks.mmlu.FullTextMMLU")
+    register_lazy_task("eval_framework.tasks.benchmarks.opengptx_eu20.MMLU_EU20_DE")
+    register_lazy_task("eval_framework.tasks.benchmarks.opengptx_eu20.MMLU_EU20_FR")
+    register_lazy_task("eval_framework.tasks.benchmarks.mmlu_de.MMLU_DE")
+    register_lazy_task("eval_framework.tasks.benchmarks.mmlu_pro.MMLU_PRO")
+    register_lazy_task("eval_framework.tasks.benchmarks.mmlu_pro.MMLU_PRO_COT")
+    register_lazy_task("eval_framework.tasks.benchmarks.mmlu.MMLU_COT")
+    register_lazy_task("eval_framework.tasks.benchmarks.mmmlu.MMMLU")
+    register_lazy_task("eval_framework.tasks.benchmarks.mmmlu.MMMLU_GERMAN_COT")
+    register_lazy_task("eval_framework.tasks.benchmarks.pawsx.PAWSX")
+    register_lazy_task("eval_framework.tasks.benchmarks.piqa.PIQA")
+    register_lazy_task("eval_framework.tasks.benchmarks.openbookqa.OPENBOOKQA")
+    register_lazy_task("eval_framework.tasks.benchmarks.sciq.SCIQ")
+    register_lazy_task("eval_framework.tasks.benchmarks.squad.SQUAD")
+    register_lazy_task("eval_framework.tasks.benchmarks.squad.SQUAD2")
+    register_lazy_task("eval_framework.tasks.benchmarks.tablebench.TableBench")
+    register_lazy_task("eval_framework.tasks.benchmarks.triviaqa.TRIVIAQA")
+    register_lazy_task("eval_framework.tasks.benchmarks.truthfulqa.TRUTHFULQA")
+    register_lazy_task("eval_framework.tasks.benchmarks.opengptx_eu20.TRUTHFULQA_EU20_DE")
+    register_lazy_task("eval_framework.tasks.benchmarks.opengptx_eu20.TRUTHFULQA_EU20_FR")
+    register_lazy_task("eval_framework.tasks.benchmarks.winogender.WINOGENDER")
+    register_lazy_task("eval_framework.tasks.benchmarks.winogrande.WINOGRANDE")
+    register_lazy_task("eval_framework.tasks.benchmarks.winox.WINOX_DE")
+    register_lazy_task("eval_framework.tasks.benchmarks.winox.WINOX_FR")
+    register_lazy_task("eval_framework.tasks.benchmarks.wmt.WMT14")
+    register_lazy_task("eval_framework.tasks.benchmarks.wmt.WMT16")
+    register_lazy_task("eval_framework.tasks.benchmarks.wmt.WMT20")
+    register_lazy_task("eval_framework.tasks.benchmarks.wmt.WMT14_INSTRUCT")
+    register_lazy_task("eval_framework.tasks.benchmarks.wmt.WMT16_INSTRUCT")
+    register_lazy_task("eval_framework.tasks.benchmarks.wmt.WMT20_INSTRUCT")
+    register_lazy_task("eval_framework.tasks.benchmarks.zero_scrolls.ZERO_SCROLLS_QUALITY")
+    register_lazy_task("eval_framework.tasks.benchmarks.zero_scrolls.ZERO_SCROLLS_SQUALITY")
+    register_lazy_task("eval_framework.tasks.benchmarks.zero_scrolls.ZERO_SCROLLS_QMSUM")
+    register_lazy_task("eval_framework.tasks.benchmarks.zero_scrolls.ZERO_SCROLLS_QASPER")
+    register_lazy_task("eval_framework.tasks.benchmarks.zero_scrolls.ZERO_SCROLLS_GOV_REPORT")
+    register_lazy_task("eval_framework.tasks.benchmarks.zero_scrolls.ZERO_SCROLLS_NARRATIVEQA")
+    register_lazy_task("eval_framework.tasks.benchmarks.zero_scrolls.ZERO_SCROLLS_MUSIQUE")
+    register_lazy_task("eval_framework.tasks.benchmarks.zero_scrolls.ZERO_SCROLLS_SPACE_DIGEST")
+    register_lazy_task("eval_framework.tasks.benchmarks.quality.QUALITY")
+    register_lazy_task("eval_framework.tasks.benchmarks.sphyr.SPHYR")
+    register_lazy_task("eval_framework.tasks.benchmarks.struct_eval.StructEval")
+    register_lazy_task("eval_framework.tasks.benchmarks.struct_eval.RenderableStructEval")
+
+    try:
+        # Importing the companion registers the additional tasks from the module.
+        # This is mostly for convenience for internal use-cases
+        import eval_framework_companion  # noqa
+    except ImportError:
+        pass
+
+
+def make_sure_all_hf_datasets_are_in_cache() -> None:
+    for task_name, task_class in registered_tasks_iter():
+        task = task_class()
+        for attempt in range(10):
+            try:
+                for _ in task.iterate_samples(num_samples=1):
+                    pass
+                break
+            except Exception as e:
+                logger.info(f"{e} Will retry loading {task_name} in a few seconds, attempt #{attempt + 1}.")
+                time.sleep(random.randint(1, 5))
+        logger.info(f"Processed {task_name}")
+
+
+if __name__ == "__main__":
+    print(list(registered_tasks_iter()))