Explainable Generic ML Pipeline with MLflow | by Mena Wang, PhD | Nov, 2024

Many machine studying algorithms — reminiscent of linear fashions (e.g., linear regression, SVM), distance-based fashions (e.g., KNN, PCA), and gradient-based fashions (e.g., gradient boosting strategies or gradient descent optimization) — are likely to carry out higher with scaled enter options, as a result of scaling prevents options with bigger ranges from dominating the educational course of. Moreover, real-world knowledge usually accommodates lacking values. Subsequently, on this first iteration, we are going to construct a pre-processor that may be educated to scale new knowledge and impute lacking values, making ready it for mannequin consumption.

As soon as this pre-processor is constructed, I’ll then demo the best way to simply plug it into pyfunc ML pipeline. Sounds good? Let’s go. 🤠

class PreProcessor(BaseEstimator, TransformerMixin):
"""
Customized preprocessor for numeric options.
- Handles scaling of numeric knowledge
- Performs imputation of lacking values
Attributes:
transformer (Pipeline): Pipeline for numeric preprocessing
options (Checklist[str]): Names of enter options
"""
def __init__(self):
"""
Initialize preprocessor.
- Creates placeholder for transformer pipeline
"""
self.transformer = None
def match(self, X, y=None):
"""
Suits the transformer on the offered dataset.
- Configures scaling for numeric options
- Units up imputation for lacking values
- Shops characteristic names for later use
Parameters:
X (pd.DataFrame): The enter options to suit the transformer.
y (pd.Sequence, non-compulsory): Goal variable, not used on this technique.
Returns:
PreProcessor: The fitted transformer occasion.
"""
self.options = X.columns.tolist()
if self.options:
self.transformer = Pipeline(steps=[
('imputer', SimpleImputer(strategy='median')),
('scaler', StandardScaler())
])
self.transformer.match(X[self.features])
return self
def rework(self, X):
"""
Remodel enter knowledge utilizing fitted pipeline.
- Applies scaling to numeric options
- Handles lacking values by way of imputation
Parameters:
X (pd.DataFrame): Enter options to remodel
Returns:
pd.DataFrame: Reworked knowledge with scaled and imputed options
"""
X_transformed = pd.DataFrame()
if self.options:
transformed_data = self.transformer.rework(X[self.features])
X_transformed[self.features] = transformed_data
X_transformed.index = X.index
return X_transformed
def fit_transform(self, X, y=None):
"""
Suits the transformer on the enter knowledge after which transforms it.
Parameters:
X (pd.DataFrame): The enter options to suit and rework.
y (pd.Sequence, non-compulsory): Goal variable, not used on this technique.
Returns:
pd.DataFrame: The reworked knowledge.
"""
self.match(X, y)
return self.rework(X)

This pre-processor may be fitted on practice knowledge after which used to course of any new knowledge. It’s going to turn into a component within the ML pipeline under, however after all, we will use or check it independently. Let’s create an artificial dataset and use the pre-processor to remodel it.

# Set parameters for artificial knowledge
n_feature = 10
n_inform = 4
n_redundant = 0
n_samples = 1000
# Generate artificial classification knowledge
X, y = make_classification(
n_samples=n_samples,
n_features=n_feature,
n_informative=n_inform,
n_redundant=n_redundant,
shuffle=False,
random_state=12
)
# Create characteristic names
feat_names = [f'inf_{i+1}' for i in range(n_inform)] + 
[f'rand_{i+1}' for i in range(n_feature - n_inform)]
# Convert to DataFrame with named options
X = pd.DataFrame(X, columns=feat_names)
# Cut up knowledge into practice and check units
X_train, X_test, y_train, y_test = train_test_split(
X, y,
test_size=0.2,
random_state=22
)

Under are screenshots from {sweetViz} reviews earlier than vs after scaling; you possibly can see that scaling didn’t change the underlying form of every characteristic’s distribution however merely rescaled and shifted it. BTW, it takes two traces to generate a reasonably complete EDA report with {sweetViz}, code obtainable within the GitHub repo linked above. 🥂

Now, let’s create an ML pipeline within the mlflow.pyfunc flavour that may encapsulate this preprocessor.

class ML_PIPELINE(mlflow.pyfunc.PythonModel):
"""
Customized ML pipeline for classification and regression.
- work with any scikit-learn appropriate mannequin
- Combines preprocessing and mannequin coaching
- Handles mannequin predictions
- Appropriate with MLflow monitoring
- Helps MLflow deployment
Attributes:
mannequin (BaseEstimator or None): A scikit-learn appropriate mannequin occasion
preprocessor (Any or None): Knowledge preprocessing pipeline
config (Any or None): Non-obligatory config for mannequin settings 
job(str): Sort of ML job ('classification' or 'regression')
"""
def __init__(self, mannequin=None, preprocessor=None, config=None):
"""
Initialize the ML_PIPELINE.
Parameters:
mannequin (BaseEstimator, non-compulsory): 
- Scikit-learn appropriate mannequin
- Defaults to None
preprocessor (Any, non-compulsory):
- Transformer or pipeline for knowledge preprocessing
- Defaults to None
config (Any, non-compulsory):
- Extra mannequin settings
- Defaults to None
"""
self.mannequin = mannequin
self.preprocessor = preprocessor
self.config = config
self.job = "classification" if hasattr(self.mannequin, "predict_proba") else "regression"
def match(self, X_train: pd.DataFrame, y_train: pd.Sequence):
"""
Practice the mannequin on offered knowledge.
- Applies preprocessing to options
- Suits mannequin on reworked knowledge
Parameters:
X_train (pd.DataFrame): Coaching options
y_train (pd.Sequence): Goal values
"""
X_train_preprocessed = self.preprocessor.fit_transform(X_train.copy())
self.mannequin.match(X_train_preprocessed, y_train)
def predict(
self, context: Any, model_input: pd.DataFrame
) -> np.ndarray:
"""
Generate predictions utilizing educated mannequin.
- Applies preprocessing to new knowledge
- Makes use of mannequin to make predictions
Parameters:
context (Any): Non-obligatory context info offered 
by MLflow through the prediction part
model_input (pd.DataFrame): Enter options
Returns:
Any: Mannequin predictions or possibilities
"""
processed_model_input = self.preprocessor.rework(model_input.copy())
if self.job == "classification":
prediction = self.mannequin.predict_proba(processed_model_input)[:,1]
elif self.job == "regression":
prediction = self.mannequin.predict(processed_model_input)
return prediction

The ML pipeline outlined above takes the preprocessor and ML algorithm as parameters. Utilization instance under

# outline the ML pipeline occasion with lightGBM classifier
ml_pipeline = ML_PIPELINE(mannequin = lgb.LGBMClassifier(),
preprocessor = PreProcessor())

It is so simple as that! 🎉 If you wish to experiment with one other algorithm, simply swap it like proven under. As a wrapper, it may well encapsulate each regression and classification algorithms. For the latter, predicted possibilities are returned, as proven within the instance above.

# outline the ML pipeline occasion with random forest regressor
ml_pipeline = ML_PIPELINE(mannequin = RandomForestRegressor(),
preprocessor = PreProcessor())

As you possibly can see from the code chunk under, passing hyperparameters to the algorithms is straightforward, making this ML pipeline an ideal instrument for hyperparameter tuning. I’ll elaborate on this matter within the following articles.

params = {
'n_estimators': 100,
'max_depth': 6,
'learning_rate': 0.1 
}
mannequin = xgb.XGBClassifier(**params)
ml_pipeline = ML_PIPELINE(mannequin = mannequin,
preprocessor = PreProcessor())

As a result of this ml pipeline is constructed within the mlflow.pyfunc flavour. We are able to log it with wealthy metadata saved robotically by mlflow for downstream use. When deployed, we will feed the metadata as context for the mannequin within the predict perform as proven under. Extra information and demos can be found in my earlier article, which is linked at first.

# practice the ML pipeline
ml_pipeline.match(X_train, y_train)
# use the educated pipeline for prediction
y_prob = ml_pipeline.predict(
context=None, # present metadata for mannequin in manufacturing
model_input=X_test
)
auc = roc_auc_score(y_test, y_prob)
print(f"auc: {auc:.3f}")

The above pre-processor has labored nicely to this point, however let’s enhance it in two methods under after which exhibit the best way to swap between pre-processors simply.

1. Permit customers to customise the pre-processing course of. For example, to specify the impute technique.
2. Increase pre-processor capability to deal with categorical options.

    class PreProcessor_v2(BaseEstimator, TransformerMixin):
"""
Customized transformer for knowledge preprocessing.
- Scales numeric options
- Encodes categorical options
- Handles lacking values by way of imputation
- Appropriate with scikit-learn pipeline
Attributes:
num_impute_strategy (str): Numeric imputation technique
cat_impute_strategy (str): Categorical imputation technique
num_transformer (Pipeline): Numeric preprocessing pipeline
cat_transformer (Pipeline): Categorical preprocessing pipeline
transformed_cat_cols (Checklist[str]): One-hot encoded column names
num_features (Checklist[str]): Numeric characteristic names
cat_features (Checklist[str]): Categorical characteristic names
"""
def __init__(self, num_impute_strategy='median', 
cat_impute_strategy='most_frequent'):
"""
Initialize the transformer.
- Units up numeric knowledge transformer
- Units up categorical knowledge transformer
- Configures imputation methods
Parameters:
num_impute_strategy (str): Technique for numeric lacking values
cat_impute_strategy (str): Technique for categorical lacking values
"""
self.num_impute_strategy = num_impute_strategy
self.cat_impute_strategy = cat_impute_strategy
def match(self, X, y=None):
"""
Match transformer on enter knowledge.
- Identifies characteristic varieties
- Configures characteristic scaling
- Units up encoding
- Suits imputation methods
Parameters:
X (pd.DataFrame): Enter options
y (pd.Sequence, non-compulsory): Goal variable, not used
Returns:
CustomTransformer: Fitted transformer
"""
self.num_features = X.select_dtypes(embody=np.quantity).columns.tolist()
self.cat_features = X.select_dtypes(exclude=np.quantity).columns.tolist()
if self.num_features:
self.num_transformer = Pipeline(steps=[
('imputer', SimpleImputer(strategy=self.num_impute_strategy)),
('scaler', StandardScaler())
])
self.num_transformer.match(X[self.num_features])
if self.cat_features:
self.cat_transformer = Pipeline(steps=[
('imputer', SimpleImputer(strategy=self.cat_impute_strategy)),
('encoder', OneHotEncoder(handle_unknown='ignore'))
])
self.cat_transformer.match(X[self.cat_features])
return self
def get_transformed_cat_cols(self):
"""
Get reworked categorical column names.
- Creates names after one-hot encoding
- Combines class with encoded values
Returns:
Checklist[str]: One-hot encoded column names
"""
cat_cols = []
cats = self.cat_features
cat_values = self.cat_transformer['encoder'].categories_
for cat, values in zip(cats, cat_values):
cat_cols += [f'{cat}_{value}' for value in values]
return cat_cols
def rework(self, X):
"""
Remodel enter knowledge.
- Applies fitted scaling
- Applies fitted encoding
- Handles numeric and categorical options
Parameters:
X (pd.DataFrame): Enter options
Returns:
pd.DataFrame: Reworked knowledge
"""
X_transformed = pd.DataFrame()
if self.num_features:
transformed_num_data = self.num_transformer.rework(X[self.num_features])
X_transformed[self.num_features] = transformed_num_data
if self.cat_features:
transformed_cat_data = self.cat_transformer.rework(X[self.cat_features]).toarray()
self.transformed_cat_cols = self.get_transformed_cat_cols()
transformed_cat_df = pd.DataFrame(transformed_cat_data, columns=self.transformed_cat_cols)
X_transformed = pd.concat([X_transformed, transformed_cat_df], axis=1)
X_transformed.index = X.index
return X_transformed
def fit_transform(self, X, y=None):
"""
Match and rework enter knowledge.
- Suits transformer to knowledge
- Applies transformation
- Combines each operations
Parameters:
X (pd.DataFrame): Enter options
y (pd.Sequence, non-compulsory): Goal variable, not used
Returns:
pd.DataFrame: Reworked knowledge
"""
self.match(X, y)
return self.rework(X)

There you’ve it: a brand new preprocessor that’s 1) extra customizable and a couple of) handles each numerical and categorical options. Let’s outline an ML pipeline occasion with it.

# Outline a PreProcessor (V2) occasion whereas specifying impute technique
preprocessor = PreProcessor_v2(
num_impute_strategy = 'imply'
)
# Outline an ML Pipeline occasion with this preprocessor
ml_pipeline = ML_PIPELINE(
mannequin = xgb.XGBClassifier(), # swap ML algorithms
preprocessor = PreProcessor # swap pre-processors 
)

Let’s check this new ML pipeline occasion with one other artificial dataset containing each numerical and categorical options.

# add missings
np.random.seed(42) 
missing_rate = 0.20 
n_missing = int(np.flooring(missing_rate * X.dimension))
rows = np.random.randint(0, X.form[0], n_missing)
cols = np.random.randint(0, X.form[1], n_missing)
X.values[rows, cols] = np.nan
actual_missing_rate = X.isna().sum().sum() / X.dimension
print(f"Goal lacking fee: {missing_rate:.2%}")
print(f"Precise lacking fee: {actual_missing_rate:.2%}")
# change X['inf_1] to categorical
percentiles = [0, 0.1, 0.5, 0.9, 1]
labels = ['bottom', 'lower-mid', 'upper-mid', 'top']
X['inf_1'] = pd.qcut(X['inf_1'], q=percentiles, labels=labels)

There you’ve it—the ML pipeline runs easily with the brand new knowledge. As anticipated, nevertheless, if we outline the ML pipeline with the earlier preprocessor after which run it on this dataset, we are going to encounter errors as a result of the earlier preprocessor was not designed to deal with categorical options.

# create an ML pipeline occasion with PreProcessor v1
ml_pipeline = ML_PIPELINE(
mannequin = lgb.LGBMClassifier(verbose = -1), 
preprocessor = PreProcessor()
)
attempt:
ml_pipeline.match(X_train, y_train)
besides Exception as e:
print(f"Error: {e}")

Error: Can not use median technique with non-numeric knowledge:
couldn't convert string to drift: 'lower-mid'

Including an explainer to an ML pipeline may be tremendous useful in a number of methods:

1. Mannequin Choice: It helps us choose the most effective mannequin by evaluating the soundness of its reasoning. Two algorithms could carry out equally on metrics like AUC or precision, however the important thing options they depend on could differ. Reviewing mannequin reasoning with area specialists to debate which mannequin makes extra sense in such situations is a good suggestion.
2. Troubleshooting: One useful technique for mannequin enchancment is to research the reasoning behind errors. For instance, in classification issues, we will establish false positives the place the mannequin was most assured (i.e., produced the best predicted prospects) and examine what went flawed within the reasoning and what key options contributed to the errors.
3. Mannequin Monitoring: In addition to the everyday monitoring components reminiscent of knowledge drift and efficiency metrics, it’s informative to watch mannequin reasoning as nicely. If there’s a vital shift in key options that drive the choices made by a mannequin in manufacturing, I need to be alerted.
4. Mannequin Implementation: In some situations, supplying mannequin reasoning together with mannequin predictions may be extremely helpful to our finish customers. For instance, to assist a customer support agent greatest retain a churning buyer, we will present the churn rating alongside the client options that contributed to this rating.

As a result of our ML pipeline is algorithm agnostic, it’s crucial that the explainer can even work throughout algorithms.

SHAP (SHapley Additive exPlanations) values are a wonderful alternative for our function as a result of they supply theoretically sturdy explanations based mostly on recreation principle. They’re designed to work constantly throughout algorithms, together with each tree-based and non-tree-based fashions, with some approximations for the latter. Moreover, SHAP presents wealthy visualization capabilities and is broadly thought to be an trade commonplace.

Within the notebooks under, I’ve dug into the similarities and variations between SHAP implementations for varied ML algorithms.

To create a generic explainer for our ML pipeline, the important thing variations to deal with are

1. Whether or not the mannequin is instantly supported by shap.Explainer

The model-specific SHAP explainers are considerably extra environment friendly than the model-agnostic ones. Subsequently, the method we take right here is

• first makes an attempt to make use of the direct SHAP explainer for the mannequin kind,
• If that fails, falls again to a model-agnostic explainer utilizing the predict perform.

2. The form of SHAP values

For binary classification issues, SHAP values can are available two codecs/shapes.

• Format 1: Solely exhibits impression on constructive class

form = (n_samples, n_features) # second array

• Format 2: Exhibits impression on each lessons

form = (n_samples, n_features, n_classes) # 3d array

• The explainer implementation under at all times exhibits the impression on the constructive class. When the impression on each lessons is out there in SHAP values, it selects those on the constructive class.

Please see the code under for the implementation of the method mentioned above.

class ML_PIPELINE(mlflow.pyfunc.PythonModel):
"""
Customized ML pipeline for classification and regression.
- Works with scikit-learn appropriate fashions
- Handles knowledge preprocessing
- Manages mannequin coaching and predictions
- Present international and native mannequin rationalization
- Appropriate with MLflow monitoring
- Helps MLflow deployment
Attributes:
mannequin (BaseEstimator or None): A scikit-learn appropriate mannequin occasion
preprocessor (Any or None): Knowledge preprocessing pipeline
config (Any or None): Non-obligatory config for mannequin settings 
job(str): Sort of ML job ('classification' or 'regression')
both_class (bool): Whether or not SHAP values embody each lessons
shap_values (shap.Clarification): SHAP values for mannequin rationalization
X_explain (pd.DataFrame): Processed options for SHAP rationalization
"""
# ------- similar code as above ---------
def explain_model(self,X):
"""
Generate SHAP values and plots for mannequin interpretation. 
This technique:
1. Transforms the enter knowledge utilizing the fitted preprocessor
2. Creates a SHAP explainer applicable for the mannequin kind
3. Calculates SHAP values for characteristic significance
4. Generates a abstract plot of characteristic significance
Parameters:
X : pd.DataFrame
Enter options to generate explanations for. 
Returns: None
The strategy shops the next attributes within the class:
- self.X_explain : pd.DataFrame
Reworked knowledge with authentic numeric values for interpretation
- self.shap_values : shap.Clarification
SHAP values for every prediction
- self.both_class : bool
Whether or not the mannequin outputs possibilities for each lessons          
"""
X_transformed = self.preprocessor.rework(X.copy())
self.X_explain = X_transformed.copy()
# get pre-transformed values for numeric options
self.X_explain[self.preprocessor.num_features] = X[self.preprocessor.num_features]
self.X_explain.reset_index(drop=True)
attempt:
# Try to create an explainer that instantly helps the mannequin
explainer = shap.Explainer(self.mannequin)
besides:
# Fallback for fashions or shap variations the place direct help could also be restricted
explainer = shap.Explainer(self.mannequin.predict, X_transformed)
self.shap_values = explainer(X_transformed)  
# get the form of shap values and extract accordingly
self.both_class = len(self.shap_values.values.form) == 3
if self.both_class:
shap.summary_plot(self.shap_values[:,:,1])
elif self.both_class == False:
shap.summary_plot(self.shap_values)
def explain_case(self,n):
"""
Generate SHAP waterfall plot for one particular case.
- Exhibits characteristic contributions
- Begins from base worth
- Ends at closing prediction
- Exhibits authentic characteristic values for higher interpretability
Parameters:
n (int): Case index (1-based)
e.g., n=1 explains the primary case.
Returns:
None: Shows SHAP waterfall plot
Notes:
- Requires explain_model() first
- Exhibits constructive class for binary duties
"""
if self.shap_values is None:
print("""
Please clarify mannequin first by operating
`explain_model()` utilizing a particular dataset
""")
else:
self.shap_values.knowledge = self.X_explain
if self.both_class:
shap.plots.waterfall(self.shap_values[:,:,1][n-1])
elif self.both_class == False:
shap.plots.waterfall(self.shap_values[n-1]) 

Now, the up to date ML pipeline occasion can create explanatory graphs for you in only one line of code. 😎

After all, you possibly can log a educated ML pipeline utilizing mlflow and revel in all of the metadata for mannequin deployment and reproducibility. Within the screenshot under, you possibly can see that along with the pickled pyfunc mannequin itself, the Python surroundings, metrics, and hyperparameters have all been logged in only a few traces of code under. To be taught extra, please consult with my earlier article on mlflow.pyfunc, which is linked at first.

# Log the mannequin with MLflow
with mlflow.start_run() as run:
# Log the customized mannequin with auto-captured conda surroundings
model_info = mlflow.pyfunc.log_model(
artifact_path="mannequin",
python_model=ml_pipeline,
conda_env=mlflow.sklearn.get_default_conda_env()
)
# Log mannequin parameters
mlflow.log_params(ml_pipeline.mannequin.get_params())
# Log metrics
mlflow.log_metric("rmse", rmse)
# Get the run ID
run_id = run.information.run_id

That is it, a generic and explainable ML pipeline that works for each classification and regression algorithms. Take the code and prolong it to fit your use case. 🤗 In the event you discover this handy, please give me a clap 👏🥰

To additional our journey on the mlflow.pyfunc sequence, under are some matters I’m contemplating. Be at liberty to depart a remark and let me know what you want to see. 🥰

• Characteristic choice
• Hyperparameter tuning
• If as a substitute of selecting between off-the-shelf algorithms, one decides to ensemble a number of algorithms or have extremely personalized options, they will nonetheless get pleasure from a generic mannequin illustration and seamless migration by way of mlflow.pyfunc.

Keep tuned and comply with me on Medium. 😁

💼LinkedIn | 😺GitHub | 🕊️Twitter/X