typestar

sk_churn_model.py in Python

An end-to-end model: split, preprocess by column type, grid search, report.

import numpy as np
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import RandomForestClassifier
from sklearn.impute import SimpleImputer
from sklearn.metrics import classification_report, roc_auc_score
from sklearn.model_selection import GridSearchCV, train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

NUMERIC = [0, 1, 2]
CATEGORICAL = [3]


def make_data(rows=600, seed=0):
    """Tenure, monthly spend, support calls, and a plan name."""
    rng = np.random.default_rng(seed)
    tenure = rng.integers(1, 72, size=rows)
    spend = rng.normal(70, 25, size=rows).clip(10, None)
    calls = rng.poisson(1.2, size=rows)
    plan = rng.choice(["basic", "plus", "pro"], size=rows)

    risk = (0.03 * calls + 0.02 * (spend / 50) - 0.02 * (tenure / 12)
            + (plan == "basic") * 0.15)
    churn = (rng.random(rows) < risk.clip(0.02, 0.9)).astype(int)

    X = np.empty((rows, 4), dtype=object)
    X[:, 0], X[:, 1], X[:, 2], X[:, 3] = tenure, spend, calls, plan
    return X, churn


def build_pipeline():
    numeric = Pipeline([
        ("impute", SimpleImputer(strategy="median")),
        ("scale", StandardScaler()),
    ])
    prepare = ColumnTransformer([
        ("numeric", numeric, NUMERIC),
        ("categorical", OneHotEncoder(handle_unknown="ignore"), CATEGORICAL),
    ])
    return Pipeline([
        ("prepare", prepare),
        ("model", RandomForestClassifier(random_state=0)),
    ])


def main():
    X, y = make_data()
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.25, random_state=0, stratify=y)

    grid = {
        "model__n_estimators": [100, 300],
        "model__max_depth": [4, 8, None],
        "model__min_samples_leaf": [1, 5],
    }
    search = GridSearchCV(build_pipeline(), grid, cv=4, scoring="roc_auc",
                          n_jobs=1)
    search.fit(X_train, y_train)

    print("best parameters")
    for key, value in sorted(search.best_params_.items()):
        print(f"  {key:28} {value}")
    print(f"cross-validated auc {search.best_score_:.3f}")

    best = search.best_estimator_
    probabilities = best.predict_proba(X_test)[:, 1]
    print(f"held-out auc        {roc_auc_score(y_test, probabilities):.3f}")
    print(classification_report(y_test, best.predict(X_test), digits=3))

    names = best.named_steps["prepare"].get_feature_names_out()
    weights = best.named_steps["model"].feature_importances_
    order = np.argsort(weights)[::-1]
    print("what mattered")
    for i in order[:5]:
        print(f"  {names[i]:28} {weights[i]:.3f}")


if __name__ == "__main__":
    main()

How it works

  1. The ColumnTransformer keeps numeric and categorical paths separate.
  2. GridSearchCV cross-validates every combination inside the pipeline.
  3. The held-out set is scored once, at the end, and never tuned against.

Keywords and builtins used here

The run, in numbers

Lines
79
Characters to type
2394
Tokens
666
Three-star pace
115 tpm

At the three-star pace of 115 tokens a minute, this run takes about 347 seconds.

Type this snippet

Step 1 of 1 in Encore, step 25 of 25 in Machine learning with scikit-learn.

← Previous