tabular.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""Tabular backends: linear and boosting on CSV."""
 
from __future__ import annotations
 
import pickle
from pathlib import Path
from typing import Any
 
import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import HistGradientBoostingClassifier, HistGradientBoostingRegressor
from sklearn.impute import SimpleImputer
from sklearn.inspection import permutation_importance
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.metrics import accuracy_score, f1_score, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder
 
from aquin.recipe.schema import Checkpoint, EvalGate, Recipe, RecipeError
from aquin.run import Run
 
 
def _task(recipe: Recipe, y: pd.Series) -> str:
    explicit = str(recipe.train.get("task") or "").lower()
    if explicit in ("classification", "regression"):
        return explicit
    if y.dtype == object or str(y.dtype) in ("string", "bool", "boolean"):
        return "classification"
    if y.nunique() <= 20 and y.dtype.kind in "iu":
        return "classification"
    return "regression"
 
 
def _build_pipeline(method: str, task: str, X: pd.DataFrame, train_cfg: dict[str, Any]) -> Pipeline:
    num_cols = [c for c in X.columns if pd.api.types.is_numeric_dtype(X[c])]
    cat_cols = [c for c in X.columns if c not in num_cols]
    transformers: list[tuple[str, Any, list[str]]] = []
    if num_cols:
        transformers.append(("num", SimpleImputer(strategy="median"), num_cols))
    if cat_cols:
        transformers.append(
            (
                "cat",
                Pipeline(
                    [
                        ("impute", SimpleImputer(strategy="most_frequent")),
                        ("oh", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
                    ]
                ),
                cat_cols,
            )
        )
    if not transformers:
        raise RecipeError("No usable columns in CSV.")
    pre = ColumnTransformer(transformers, remainder="drop")
    if task == "classification":
        est: Any = (
            LogisticRegression(max_iter=int(train_cfg.get("max_iter") or 400))
            if method == "linear"
            else HistGradientBoostingClassifier(
                max_depth=int(train_cfg.get("max_depth") or 6),
                min_samples_leaf=int(train_cfg.get("min_samples_leaf") or 2),
            )
        )
    else:
        est = (
            Ridge(alpha=float(train_cfg.get("alpha") or 1.0))
            if method == "linear"
            else HistGradientBoostingRegressor(
                max_depth=int(train_cfg.get("max_depth") or 6),
                min_samples_leaf=int(train_cfg.get("min_samples_leaf") or 2),
            )
        )
    return Pipeline([("pre", pre), ("est", est)])
 
 
def _scores(task: str, y_true: Any, y_pred: Any) -> dict[str, float]:
    if task == "classification":
        return {
            "accuracy": float(accuracy_score(y_true, y_pred)),
            "f1_macro": float(f1_score(y_true, y_pred, average="macro", zero_division=0)),
        }
    rmse = float(np.sqrt(mean_squared_error(y_true, y_pred)))
    return {"rmse": rmse, "r2": float(r2_score(y_true, y_pred))}
 
 
def _gate_score(metric: str, scores: dict[str, float], task: str) -> float:
    if metric in scores:
        return scores[metric]
    return scores["accuracy"] if task == "classification" else scores["rmse"]
 
 
def train_tabular(recipe: Recipe, recorder: Run, out_dir: Path) -> dict[str, Any]:
    path = Path(str(recipe.data["path"]))
    target = str(recipe.data["target"])
    df = pd.read_csv(path)
    if target not in df.columns:
        raise RecipeError(f"target column {target!r} not in {list(df.columns)}")
    y = df[target]
    X = df.drop(columns=[target])
    if X.empty:
        raise RecipeError("CSV has no feature columns.")
    task = _task(recipe, y)
    method = str(recipe.train["method"]).lower()
    seed = int(recipe.train.get("seed") or 0)
    test_size = float(recipe.train.get("test_size") or 0.2)
 
    strat = y if task == "classification" and y.nunique() > 1 else None
    try:
        X_tr, X_te, y_tr, y_te = train_test_split(
            X, y, test_size=test_size, random_state=seed, stratify=strat
        )
    except ValueError:
        X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=test_size, random_state=seed)
 
    pipe = _build_pipeline(method, task, X_tr, recipe.train)
    pipe.fit(X_tr, y_tr)
    pred = pipe.predict(X_te)
    scores = _scores(task, y_te, pred)
    recorder.log(1, loss=float(scores.get("rmse") or 1.0 - scores.get("accuracy", 0.0)))
 
    ckpt_path = out_dir / "checkpoints" / "model.pkl"
    ckpt_path.parent.mkdir(parents=True, exist_ok=True)
    with ckpt_path.open("wb") as fh:
        pickle.dump({"pipeline": pipe, "task": task, "target": target, "columns": list(X.columns)}, fh)
    recorder.signal(1, f"checkpoint saved ({ckpt_path.name})")
 
    inspect: dict[str, Any] = {"task": task, "features": list(X.columns)}
    try:
        perm = permutation_importance(pipe, X_te, y_te, n_repeats=5, random_state=seed)
        ranking = sorted(
            zip(X.columns, perm.importances_mean.tolist()),
            key=lambda x: abs(x[1]),
            reverse=True,
        )
        inspect["permutation_importance"] = [
            {"feature": str(name), "importance": float(val)} for name, val in ranking
        ]
    except Exception as exc:
        inspect["permutation_importance_error"] = str(exc)
 
    metric = str(recipe.eval.get("metric") or ("accuracy" if task == "classification" else "rmse"))
    score = _gate_score(metric, scores, task)
    min_score = recipe.eval.get("min_score")
    if min_score is None:
        min_score = recipe.eval.get("max_score") if metric == "rmse" else None
    lower_is_better = metric in ("rmse", "mae", "loss")
    if min_score is None:
        passed = True
        skipped = True
    else:
        skipped = False
        thresh = float(min_score)
        passed = score <= thresh if lower_is_better else score >= thresh
 
    gate = EvalGate(
        passed=passed,
        metric=metric,
        score=round(float(score), 6),
        min_score=None if min_score is None else float(min_score),
        skipped=skipped,
        details={"scores": scores, "n_train": int(len(X_tr)), "n_eval": int(len(X_te))},
    )
    return {
        "checkpoint": Checkpoint(path=str(ckpt_path), kind="sklearn", step=1, extra={"task": task}),
        "gate": gate,
        "metrics": scores,
        "inspect": inspect,
    }