Skip to content

GitLab

Commit 5e632643 authored by Baber's avatar Baber
Browse files

add docs

parent 68b3cddc
Hide whitespace changes
Inline Side-by-side
Showing with 477 additions and 47 deletions
lm_eval/tasks/__init__.py
View file @ 5e632643
"""
Task Management Module for LM Evaluation Harness.
This module provides comprehensive task discovery, loading, and management functionality
for the LM Evaluation Harness. It handles YAML configuration parsing with include support,
dynamic function importing, and task indexing across multiple directories.
Key Components:
- TaskManager: Main class for task discovery and management
- YAML configuration loading with !function tag support
- Task, group, and tag indexing
- Include resolution with cycle detection
- Caching for performance optimization
Example:
Basic usage::
task_manager = TaskManager()
all_tasks = task_manager.all_tasks
task_config = task_manager._get_config("hellaswag")
Custom task paths::
task_manager = TaskManager(
include_path="/path/to/custom/tasks",
include_defaults=True
)
"""
import collections
import functools
import importlib.util
......@@ -30,8 +59,13 @@ if TYPE_CHECKING:
eval_logger = logging.getLogger(__name__)
#: List of configuration keys that are specific to groups only
GROUP_ONLY_KEYS = list(GroupConfig().to_dict().keys())
#: Base YAML loader class - uses C loader if available for performance
_Base = yaml.CLoader if getattr(yaml, "__with_libyaml__", False) else yaml.FullLoader
#: Directory names to ignore during task discovery
_IGNORE_DIRS = (
"__pycache__",
".ipynb_checkpoints",
......@@ -39,6 +73,19 @@ _IGNORE_DIRS = (
def ignore_constructor(loader: yaml.Loader, node: yaml.Node) -> None:
"""
YAML constructor that ignores !function tags during simple parsing.
This is used when mode="simple" to skip function resolution for
faster indexing operations.
Args:
loader: YAML loader instance
node: YAML node being processed
Returns:
None
"""
return None
......@@ -78,6 +125,27 @@ def _make_loader(yaml_dir: Path, simple: bool = False) -> type[yaml.Loader]:
@functools.lru_cache(maxsize=None) # ← cache module objects
def _import_function(qualname: str, *, base_path: Path) -> Callable:
"""
Dynamically import a function from a Python module relative to base_path.
This function enables YAML files to reference Python functions using
the !function tag. It supports dot notation for nested modules and
caches imported modules for performance.
Args:
qualname: Qualified function name like "my_module.my_function"
base_path: Base directory for resolving relative module paths
Returns:
The imported callable function
Raises:
ValueError: If qualname doesn't contain a module part
Example:
>>> func = _import_function("utils.custom_metric", base_path=Path("/tasks"))
>>> result = func(predictions, references)
"""
mod_path, _, func_name = qualname.rpartition(".")
if not mod_path:
raise ValueError(f"{qualname!r} has no module part")
......@@ -93,8 +161,18 @@ def _import_function(qualname: str, *, base_path: Path) -> Callable:
return getattr(mod, func_name)
@functools.lru_cache(maxsize=4096) #
@functools.lru_cache(maxsize=4096)
def _parse_yaml_file(path: Path, mode: str) -> dict:
"""
Parse a single YAML file with the appropriate loader.
Args:
path: Path to the YAML file
mode: Parsing mode ("full" or "simple")
Returns:
Parsed YAML configuration as dictionary
"""
loader_cls = _make_loader(path.parent, simple=(mode == "simple"))
with path.open("rb") as fh:
return yaml.load(fh, Loader=loader_cls)
......@@ -211,6 +289,21 @@ def load_yaml_config(
def iter_yaml_files(root: Path) -> Generator[Path, Any, None]:
"""
Recursively iterate over all YAML files in a directory tree.
Excludes files in ignored directories like __pycache__ and .ipynb_checkpoints.
Args:
root: Root directory to search for YAML files
Yields:
Path objects for each discovered YAML file
Example:
>>> for yaml_file in iter_yaml_files(Path("tasks")):
... print(f"Found task config: {yaml_file}")
"""
for p in iglob("**/*.yaml", root_dir=root, recursive=True):
# ignore check
if Path(p).parts[0] in _IGNORE_DIRS:
......@@ -219,9 +312,39 @@ def iter_yaml_files(root: Path) -> Generator[Path, Any, None]:
class TaskManager:
"""TaskManager indexes all tasks from the default `lm_eval/tasks/`
and an optional directory if provided.
"""
Central manager for task discovery, indexing, and loading.
TaskManager scans directories for YAML task configurations and maintains
an index of all available tasks, groups, and tags. It provides methods
for listing, filtering, and loading tasks with their configurations.
The manager supports:
- Automatic discovery from default lm_eval/tasks/ directory
- Custom task directories via include_path
- Task grouping and tagging
- Configuration inheritance via YAML includes
- Caching for performance
Attributes:
include_path: Additional directories to search for tasks
metadata: Global metadata to inject into all task configs
task_group_map: Mapping of tasks to their parent groups
Example:
Basic usage::
tm = TaskManager()
print(f"Found {len(tm.all_tasks)} tasks")
hellaswag_config = tm._get_config("hellaswag")
With custom tasks::
tm = TaskManager(
include_path="/my/custom/tasks",
verbosity="INFO"
)
custom_tasks = [t for t in tm.all_tasks if "custom" in t]
"""
def __init__(
......@@ -231,6 +354,16 @@ class TaskManager:
include_defaults: bool = True,
metadata: Optional[dict] = None,
) -> None:
"""
Initialize the TaskManager.
Args:
verbosity: Logging verbosity level (DEBUG, INFO, WARNING, ERROR)
include_path: Additional path(s) to search for tasks. Can be a single
path or list of paths.
include_defaults: Whether to include default tasks from lm_eval/tasks/
metadata: Global metadata dictionary to inject into all task configs
"""
if verbosity is not None:
setup_logging(verbosity)
self.include_path = include_path
......@@ -290,22 +423,27 @@ class TaskManager:
@property
def all_tasks(self) -> list[str]:
"""Get sorted list of all task names (tasks, groups, and tags)."""
return self._all_tasks
@property
def all_groups(self) -> list[str]:
"""Get sorted list of all group names."""
return self._all_groups
@property
def all_subtasks(self) -> list[str]:
"""Get sorted list of all individual task names (excludes groups and tags)."""
return self._all_subtasks
@property
def all_tags(self) -> list[str]:
"""Get sorted list of all tag names."""
return self._all_tags
@property
def task_index(self) -> dict[str, dict[str, Union[str, int, list[str]]]]:
"""Get the complete task index with metadata for all tasks."""
return self._task_index
def list_all_tasks(
......@@ -410,48 +548,96 @@ class TaskManager:
return "".join(parts)
def match_tasks(self, task_list: list[str]) -> list[str]:
"""
Match task names using pattern matching.
Supports glob-style patterns and returns all matching task names.
Args:
task_list: List of task name patterns to match
Returns:
List of matching task names
Example:
>>> tm.match_tasks(["hella*", "arc_*"])
['hellaswag', 'arc_easy', 'arc_challenge']
"""
return pattern_match(task_list, self.all_tasks)
def _name_is_registered(self, name: str) -> bool:
"""Check if a name is registered in the task index."""
return name in self.all_tasks
def _name_is_task(self, name: str) -> bool:
"""Check if a name refers to an individual task (not group or tag)."""
return (
self._name_is_registered(name) and self.task_index[name]["type"] == "task"
)
def _name_is_tag(self, name: str) -> bool:
"""Check if a name refers to a tag."""
return self._name_is_registered(name) and self.task_index[name]["type"] == "tag"
def _name_is_group(self, name: str) -> bool:
"""Check if a name refers to a group."""
return (
self._name_is_registered(name) and self.task_index[name]["type"] == "group"
)
def _name_is_python_task(self, name: str) -> bool:
"""Check if a name refers to a Python-defined task."""
return (
self._name_is_registered(name)
and self.task_index[name]["type"] == "python_task"
)
def _config_is_task(self, config: dict) -> bool:
"""Check if a config dictionary defines a single task."""
return "task" in config and isinstance(config["task"], str)
def _config_is_group(self, config: dict) -> bool:
"""Check if a config dictionary defines a group of tasks."""
return "task" in config and isinstance(config["task"], list)
def _config_is_python_task(self, config: dict) -> bool:
"""Check if a config dictionary defines a Python class-based task."""
return "class" in config
def _config_is_task_list(self, config: dict) -> bool:
"""Check if a config dictionary defines a task list."""
return "task_list" in config and isinstance(config["task_list"], list)
def _get_yaml_path(self, name: str) -> Union[str, int]:
"""
Get the YAML file path for a registered task.
Args:
name: Task name
Returns:
Path to YAML file, or -1 for Python-only tasks
Raises:
ValueError: If task name is not registered
"""
if name not in self.task_index:
raise ValueError
return self.task_index[name]["yaml_path"]
def _get_config(self, name: str) -> dict:
"""
Load the full configuration for a registered task.
Args:
name: Task name
Returns:
Complete task configuration dictionary
Raises:
ValueError: If task name is not registered
"""
if name not in self.task_index:
raise ValueError
yaml_path = self._get_yaml_path(name)
......@@ -461,6 +647,18 @@ class TaskManager:
return load_yaml_config(Path(yaml_path), mode="full")
def _get_tasklist(self, name: str) -> Union[list[str], int]:
"""
Get the task list for a group or tag.
Args:
name: Group or tag name
Returns:
List of task names in the group/tag
Raises:
ValueError: If name refers to an individual task
"""
if self._name_is_task(name):
raise ValueError
return self.task_index[name]["task"]
......@@ -505,15 +703,34 @@ class TaskManager:
return dict(collections.ChainMap(*map(fn, reversed(subtask_list))))
def _process_alias(self, config: dict, group: Optional[str] = None) -> dict:
# 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.
"""
Process group alias configuration.
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.
Args:
config: Task configuration dictionary
group: Group name to validate against
Returns:
Modified configuration with processed aliases
"""
if ("group_alias" in config) and ("group" in config) and group is not None:
if config["group"] != group:
config["group_alias"] = None
return config
def _class_has_config_in_constructor(self, cls) -> bool:
"""
Check if a class constructor accepts a 'config' parameter.
Args:
cls: Class to inspect
Returns:
True if constructor has 'config' parameter, False otherwise
"""
constructor = getattr(cls, "__init__", None)
return (
"config" in inspect.signature(constructor).parameters
......@@ -527,11 +744,65 @@ class TaskManager:
parent_name: Optional[str] = None,
update_config: Optional[dict] = None,
) -> Mapping:
"""
Load a single task or group with all its configurations and dependencies.
This is the core method for instantiating task objects from either task names
or configuration dictionaries. It handles complex scenarios including:
- Individual tasks and Python class-based tasks
- Groups and their constituent subtasks
- Tags and their associated tasks
- Configuration merging and inheritance
- Duplicate detection and name resolution
- Include processing and YAML inheritance
Args:
name_or_config: Either a task name (str) or configuration dict.
If str, looks up the task in the index.
If dict, processes as inline configuration.
parent_name: Name of parent group (for duplicate detection)
update_config: Additional configuration to merge into task configs
Returns:
Mapping of task/group names to instantiated task objects.
For individual tasks: {task_name: ConfigurableTask}
For groups: {group_name: {subtask1: Task1, subtask2: Task2, ...}}
Example:
Load individual task::
task_dict = tm._load_individual_task_or_group("hellaswag")
# Returns: {"hellaswag": ConfigurableTask(...)}
Load with config override::
task_dict = tm._load_individual_task_or_group(
{"task": "hellaswag", "num_fewshot": 5}
)
Load a group::
group_dict = tm._load_individual_task_or_group("arc_group")
# Returns: {"arc_group": {"arc_easy": Task1, "arc_challenge": Task2}}
"""
from lm_eval.api.task import ConfigurableTask, Task
def _load_task(
config: dict, task: str, yaml_path: Optional[str] = None
) -> dict[str, Union["ConfigurableTask", "Task"]]:
"""
Create a single task object from configuration.
Handles include processing, Python class instantiation, and metadata injection.
Args:
config: Task configuration dictionary
task: Task name
yaml_path: Path to source YAML file (for include resolution)
Returns:
Dictionary mapping task name to instantiated task object
"""
if "include" in config:
# Store the task name to preserve it after include processing
original_task_name = config.get("task", task)
......@@ -569,6 +840,17 @@ class TaskManager:
def _get_group_and_subtask_from_config(
config: dict,
) -> tuple[ConfigurableGroup, list[str]]:
"""
Extract group object and subtask list from group configuration.
Expands any tags in the task list to their constituent tasks.
Args:
config: Group configuration dictionary
Returns:
Tuple of (ConfigurableGroup, list of subtask names)
"""
if self.metadata is not None:
config["metadata"] = config.get("metadata", {}) | self.metadata
group_name = ConfigurableGroup(config=config)
......@@ -583,6 +865,19 @@ class TaskManager:
def _process_group_config(
config: dict, update_config: Optional[dict] = None
) -> tuple[dict, Optional[dict]]:
"""
Separate group-specific config from task-level config overrides.
Group-only keys (like 'group', 'aggregate') stay with the group,
while other keys become config overrides for constituent tasks.
Args:
config: Full configuration dictionary
update_config: Additional config to merge
Returns:
Tuple of (group_config, task_update_config)
"""
if update_config is not None:
config = {**config, **update_config}
_update_config = {
......@@ -717,13 +1012,31 @@ class TaskManager:
def load_task_or_group(
self, task_list: Optional[Union[str, list[str]]] = None
) -> dict:
"""Loads a dictionary of task objects from a list
"""
Load multiple tasks or groups from a list of names.
This is the main entry point for loading tasks. It handles lists
of task names and delegates to _load_individual_task_or_group for
each item, then merges the results.
Args:
task_list: Single task name or list of task names to load.
Can include individual tasks, groups, and tags.
:param task_list: Union[str, list] = None
Single string or list of string of task names to be loaded
Returns:
Dictionary mapping task/group names to loaded task objects.
Results from all requested items are merged into a single dict.
:return
dictionary of task objects
Example:
Load multiple tasks::
tasks = tm.load_task_or_group(["hellaswag", "arc_easy"])
# Returns: {"hellaswag": Task1, "arc_easy": Task2}
Load a group::
tasks = tm.load_task_or_group("arc_group")
# Returns: {"arc_group": {"arc_easy": Task1, "arc_challenge": Task2}}
"""
if isinstance(task_list, str):
task_list = [task_list]
......@@ -739,34 +1052,63 @@ class TaskManager:
return all_loaded_tasks
def load_config(self, config: dict) -> Mapping:
"""
Load a task from an inline configuration dictionary.
Args:
config: Configuration dictionary defining the task
Returns:
Mapping of task name to loaded task object
Example:
>>> config = {"task": "hellaswag", "num_fewshot": 5}
>>> task_dict = tm.load_config(config)
"""
return self._load_individual_task_or_group(config)
def _get_task_and_group(self, task_dir: Union[str, Path]) -> dict[str, dict]:
"""Creates a dictionary of tasks index with the following metadata,
- `type`, that can be either `task`, `python_task`, `group` or `tags`.
`task` refer to regular task configs, `python_task` are special
yaml files that only consists of `task` and `class` parameters.
`group` are group configs. `tags` are labels that can be assigned
to tasks to assist in sorting and calling tasks of certain themes.
- `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
"""
Scan a directory for task configurations and build an index.
Creates a dictionary of task metadata by recursively scanning for
YAML files and parsing their configurations. This method handles:
- Regular task configs with 'task' key
- Python class-based tasks with 'class' key
- Group configs with 'group' key
- Task list configs with 'task_list' key
- Tag extraction and registration
Args:
task_dir: Directory path to scan for YAML task configurations
Returns:
Dictionary mapping task names to metadata dictionaries.
Each metadata dict contains:
- 'type': One of 'task', 'python_task', 'group', 'tag'
- 'yaml_path': Path to source YAML file (or -1 for generated entries)
- 'task': For groups/tags, list of constituent task names
Note:
This method is called during TaskManager initialization to build
the master task index. It uses 'simple' parsing mode for performance.
"""
def _populate_tags_and_groups(
config: dict, task: str, tasks_and_groups: dict[str, dict]
) -> None:
"""
Extract and register tags from a task configuration.
Tags allow grouping tasks by theme or category. This function
processes the 'tag' field in task configs and maintains tag
indices for quick lookup.
Args:
config: Task configuration dictionary
task: Name of the task being processed
tasks_and_groups: Master index to update with tag information
"""
# TODO: remove group in next release
if "tag" in config:
attr_list = config["tag"]
......@@ -868,6 +1210,27 @@ class TaskManager:
def get_task_name_from_config(task_config: dict[str, str]) -> str:
"""
Extract a task name from a configuration dictionary.
Determines the canonical name for a task based on its configuration,
with fallback strategies for different config formats.
Args:
task_config: Task configuration dictionary
Returns:
String name for the task
Example:
>>> config = {"task": "hellaswag", "num_fewshot": 5}
>>> get_task_name_from_config(config)
'hellaswag'
>>> config = {"dataset_path": "custom", "dataset_name": "mytask"}
>>> get_task_name_from_config(config)
'custom_mytask'
"""
if "task" in task_config:
return task_config["task"]
if "dataset_name" in task_config:
......@@ -877,6 +1240,23 @@ def get_task_name_from_config(task_config: dict[str, str]) -> str:
def get_task_name_from_object(task_object: Union["ConfigurableTask", "Task"]) -> str:
"""
Extract the name from an instantiated task object.
Handles both ConfigurableTask and legacy Task objects with different
attribute conventions for storing the task name.
Args:
task_object: An instantiated task object
Returns:
String name of the task
Example:
>>> task = ConfigurableTask(config={"task": "hellaswag"})
>>> get_task_name_from_object(task)
'hellaswag'
"""
if hasattr(task_object, "config"):
return task_object._config["task"]
......@@ -890,10 +1270,25 @@ def get_task_name_from_object(task_object: Union["ConfigurableTask", "Task"]) ->
def _check_duplicates(task_dict: dict[str, list[str]]) -> None:
"""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.
"""
Validate that no tasks appear in multiple groups simultaneously.
Helper function used to prevent conflicts when multiple groups claim
the same constituent task. This could lead to ambiguous configuration
like conflicting num_fewshot values.
Args:
task_dict: Dictionary mapping group names to lists of subtask names
Raises:
ValueError: If any tasks appear in multiple groups
Example:
>>> task_dict = {
... "group1": ["task_a", "task_b"],
... "group2": ["task_b", "task_c"] # task_b appears twice!
... }
>>> _check_duplicates(task_dict) # Raises ValueError
"""
subtask_names = []
for key, value in task_dict.items():
......@@ -920,18 +1315,53 @@ def get_task_dict(
task_name_list: Union[str, list[Union[str, dict, "Task"]]],
task_manager: Optional[TaskManager] = None,
) -> dict[str, Union["ConfigurableTask", "Task"]]:
"""Creates a dictionary of task objects from either a name of task, config, or prepared Task object.
: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`
:return
dictionary of task objects
"""
Create a dictionary of task objects from mixed input types.
This is the main public API for loading tasks. It accepts various input
formats (names, configs, objects) and returns a unified dictionary of
instantiated task objects ready for evaluation.
The function handles:
- String task names (looked up via TaskManager)
- Configuration dictionaries (processed as inline configs)
- Pre-instantiated Task objects (used as-is)
- Validation to prevent conflicting group memberships
Args:
task_name_list: Mixed list of task specifications:
- str: Task name to look up
- dict: Inline task configuration
- Task: Pre-instantiated task object
task_manager: TaskManager instance for name resolution.
If None, creates a default TaskManager.
Returns:
Dictionary mapping task names to instantiated task objects.
All tasks are ready for evaluation.
Raises:
TypeError: If task_name_list contains unsupported types
ValueError: If there are conflicting group memberships
Example:
Mixed input types::
tasks = get_task_dict([
"hellaswag", # lookup by name
{"task": "arc_easy", "num_fewshot": 5}, # inline config
pre_existing_task_object # direct object
])
Simple case::
tasks = get_task_dict("hellaswag")
# Returns: {"hellaswag": ConfigurableTask(...)}
With custom TaskManager::
tm = TaskManager(include_path="/custom/tasks")
tasks = get_task_dict(["custom_task"], task_manager=tm)
"""
from lm_eval.api.task import ConfigurableTask, Task
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment