core/task_detector.py

"""Automatic ML task and metric detection from the target column."""

import logging
from enum import Enum

import pandas as pd

logger = logging.getLogger(__name__)


class TaskType(str, Enum):
    BINARY = "binary_classification"
    MULTICLASS = "multiclass_classification"
    REGRESSION = "regression"


DEFAULT_METRICS: dict[TaskType, str] = {
    TaskType.BINARY: "roc_auc",
    TaskType.MULTICLASS: "f1_macro",
    TaskType.REGRESSION: "rmse",
}


def detect_task(df: pd.DataFrame, target: str) -> tuple[TaskType, str]:
    """Infer the ML task type and best default metric from the target column.

    Args:
        df: The full DataFrame (after loading).
        target: Name of the target column.

    Returns:
        A tuple of (TaskType, metric_name).

    Raises:
        KeyError: If the target column is not present in the DataFrame.
        ValueError: If the target column is entirely null.
    """
    if target not in df.columns:
        raise KeyError(
            f"Target column '{target}' not found. Available columns: {list(df.columns)}"
        )

    series = df[target].dropna()
    if series.empty:
        raise ValueError(f"Target column '{target}' has no non-null values.")

    n_unique = series.nunique()
    dtype = series.dtype

    if n_unique == 2:
        task = TaskType.BINARY
    elif n_unique <= 20 and (
        pd.api.types.is_object_dtype(dtype) or pd.api.types.is_integer_dtype(dtype)
    ):
        task = TaskType.MULTICLASS
    elif pd.api.types.is_float_dtype(dtype) or n_unique > 20:
        task = TaskType.REGRESSION
    else:
        task = TaskType.MULTICLASS

    metric = DEFAULT_METRICS[task]
    logger.info(
        "Detected task=%s  |  unique target values=%d  |  default metric=%s",
        task.value,
        n_unique,
        metric,
    )
    return task, metric