salary-pipeline.py

#!/usr/bin/env python3
"""
Employee Salary Prediction
==========================
Predict current annual salary for Montgomery County employees
using a scikit-learn pipeline with mixed feature types.
"""

try:
    import cuml.accel
    cuml.accel.install()
    _GPU = True
except ImportError:
    _GPU = False

import warnings
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd

print(f"Running on: {'GPU (cuml.accel)' if _GPU else 'CPU (sklearn)'}")

from sklearn.pipeline import Pipeline, make_pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import (
    OneHotEncoder, StandardScaler, FunctionTransformer, SplineTransformer,
)
from sklearn.feature_extraction.text import HashingVectorizer
from sklearn.impute import SimpleImputer
from sklearn.decomposition import PCA, TruncatedSVD
from sklearn.linear_model import Ridge, LinearRegression
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import mean_absolute_error, r2_score

from skrub.datasets import fetch_employee_salaries


# ---------------------------------------------------------------------------
# Load data
# ---------------------------------------------------------------------------
dataset = fetch_employee_salaries()
X, y = dataset.X, dataset.y

X["date_first_hired"] = pd.to_datetime(X["date_first_hired"])

y_values = y.values.ravel().astype(np.float64)

print(f"Dataset shape: {X.shape}")
print(X.dtypes)


# ---------------------------------------------------------------------------
# Column groups
# ---------------------------------------------------------------------------
DATETIME_COLS  = ["date_first_hired"]
NUMERIC_COLS   = ["year_first_hired"]
LOW_CARD_COLS  = ["gender", "department", "department_name", "assignment_category"]
HIGH_CARD_COLS = ["division", "employee_position_title"]


# ---------------------------------------------------------------------------
# Preprocessing sub-pipelines
# ---------------------------------------------------------------------------
low_card_pipe = make_pipeline(
    SimpleImputer(strategy="most_frequent"),
    OneHotEncoder(drop="if_binary", handle_unknown="ignore", sparse_output=False, dtype=np.float32),
)

def _extract_datetime_features(X):
    col = X.iloc[:, 0]
    return pd.DataFrame({
        "year": col.dt.year,
        "month": col.dt.month,
        "day_of_week": col.dt.dayofweek,
        "day_of_year": col.dt.dayofyear,
    })

datetime_pipe = make_pipeline(
    FunctionTransformer(_extract_datetime_features),
    SimpleImputer(strategy="median"),
    SplineTransformer(n_knots=4, degree=3),
)

numeric_pipe = make_pipeline(
    SimpleImputer(strategy="median"),
)

def _to_dense(X):
    return X.toarray() if hasattr(X, "toarray") else np.asarray(X)

def _to_numpy(X):
    return X.get() if hasattr(X, "get") else np.asarray(X)

high_card_pipes = [
    (
        f"hashing_{col}",
        make_pipeline(
            HashingVectorizer(analyzer="char_wb", ngram_range=(3, 4), n_features=2**12),
            FunctionTransformer(_to_dense),
            TruncatedSVD(n_components=30),
            FunctionTransformer(_to_numpy),
        ),
        col,
    )
    for col in HIGH_CARD_COLS
]

column_transformer = ColumnTransformer(
    transformers=[
        ("low_card", low_card_pipe, LOW_CARD_COLS),
        ("datetime", datetime_pipe, DATETIME_COLS),
        ("numeric",  numeric_pipe,  NUMERIC_COLS),
    ] + high_card_pipes,
    remainder="drop",
)


# ---------------------------------------------------------------------------
# Full pipelines
# ---------------------------------------------------------------------------
# Standard ridge regression
pipeline_ridge = Pipeline([
    ("column_transformer", column_transformer),
    ("imputer",            SimpleImputer(strategy="mean", add_indicator=True)),
    ("scaler",             StandardScaler()),
    ("pca",                PCA(n_components=50)),
    ("estimator",          Ridge(alpha=1.0)),
])

# Non-negative ridge (coefficients constrained to be positive)
pipeline_ridge_positive = Pipeline([
    ("column_transformer", column_transformer),
    ("imputer",            SimpleImputer(strategy="mean", add_indicator=True)),
    ("scaler",             StandardScaler()),
    ("pca",                PCA(n_components=50)),
    ("estimator",          Ridge(alpha=1.0, positive=True)),
])

# Non-negative linear regression
pipeline_linreg_positive = Pipeline([
    ("column_transformer", column_transformer),
    ("imputer",            SimpleImputer(strategy="mean", add_indicator=True)),
    ("scaler",             StandardScaler()),
    ("pca",                PCA(n_components=50)),
    ("estimator",          LinearRegression(positive=True)),
])


# ---------------------------------------------------------------------------
# Train / test split
# ---------------------------------------------------------------------------
X_train, X_test, y_train, y_test = train_test_split(
    X, y_values, test_size=0.2, random_state=42
)


# ---------------------------------------------------------------------------
# Fit and evaluate
# ---------------------------------------------------------------------------
pipelines = {
    "Ridge":                  pipeline_ridge,
    "Ridge (positive=True)":  pipeline_ridge_positive,
    "LinearRegression (positive=True)": pipeline_linreg_positive,
}

for name, pipe in pipelines.items():
    pipe.fit(X_train, y_train)
    y_pred = pipe.predict(X_test)
    mae = mean_absolute_error(y_test, y_pred)
    r2  = r2_score(y_test, y_pred)
    print(f"{name:40s}  MAE={mae:>10,.0f}  R²={r2:.4f}")


# ---------------------------------------------------------------------------
# Cross-validation on the main pipeline
# ---------------------------------------------------------------------------
print("\n5-fold cross-validation (Ridge):")
cv_scores = cross_val_score(pipeline_ridge, X, y_values, cv=5, scoring="r2")
print(f"  R² per fold : {cv_scores.round(4)}")
print(f"  Mean R²     : {cv_scores.mean():.4f} ± {cv_scores.std():.4f}")