SNOW-1967393: Snowflake ML Registry for Model Explainability: `ValueError: Model type <class 'NoneType'> is not supported` when logging a Pipeline

[Masoud-Ghodrati]

I am training an XGBClassifier model using Snowflake ML and attempting to register it using Snowflake’s model registry. The training and evaluation steps complete successfully, but when I try to log the model using log_model(), I get the following error:

Training Accuracy: 0.9610
Evaluation Accuracy: 0.8730
Traceback (most recent call last):
  File "C:\Mas\rla_projects\Claims-AI--lodgement\code\python\src\explain\train_model.py", line 56, in <module>
    model_version = native_registry.log_model(
  File "C:\Mas\py39_sf_ml\lib\site-packages\snowflake\ml\_internal\telemetry.py", line 542, in wrap
    return ctx.run(execute_func_with_statement_params)
  ...
  File "C:\Mas\py39_sf_ml\lib\site-packages\snowflake\ml\model\_packager\model_task\model_task_utils.py", line 149, in _get_model_task    
    raise ValueError(f"Model type {type(model)} is not supported")
ValueError: Model type <class 'NoneType'> is not supported

Code Snippet:

from snowflake.snowpark import Session
from snowflake.ml.registry import registry
from snowflake.ml.modeling.preprocessing import StandardScaler
from snowflake.ml.modeling.impute import SimpleImputer
from snowflake.ml.modeling.pipeline import Pipeline
from snowflake.ml.modeling.xgboost import XGBClassifier

# Snowflake connection parameters
conn_params = {
    "user": "<...>",
    "account": "<...>",
    "warehouse": "<...>",
    "database": "<...>",
    "schema": "<...>",
    "role": "<...>",
    "authenticator": "externalbrowser",
}

# Create session
session = Session.builder.configs(conn_params).create()

# Load and prepare data
all_data = session.sql("SELECT *, IFF(CLASS = 'g', 1.0, 0.0) AS LABEL FROM Gamma_Telescope_Data").drop("CLASS")
train_data, test_data = all_data.random_split(weights=[0.9, 0.1], seed=0)

# Define feature and label columns
FEATURE_COLS = [c for c in train_data.columns if c != "LABEL"]
LABEL_COLS = ["LABEL"]

# Construct pipeline
pipeline = Pipeline(steps=[
    ("impute", SimpleImputer(input_cols=FEATURE_COLS, output_cols=FEATURE_COLS)),
    ("scaler", StandardScaler(input_cols=FEATURE_COLS, output_cols=FEATURE_COLS)),
    ("model", XGBClassifier(input_cols=FEATURE_COLS, label_cols=LABEL_COLS))
])

# Train the pipeline
pipeline.fit(train_data)

# Register model
native_registry = registry.Registry(
    session=session, 
    database_name=session.get_current_database(), 
    schema_name=session.get_current_schema()
)

model_name = "Gamma_test"
version = "V8"

model_version = native_registry.log_model(
    model=pipeline,  # <-- This line triggers the error
    model_name=model_name,
    version_name=version,
    sample_input_data=test_data,
    comment="Gamma test",
    conda_dependencies=["snowflake-ml-python==1.7.4", "snowflake-snowpark-python==1.28.0"],
    options={"enable_explainability": True}
)

Observations & Debugging Attempts:

1. Pipeline Training Works – pipeline.fit(train_data) runs without errors.
2. Pipeline Predictions Work – Predictions on training and test data succeed.
3. Model Explanation Works Without Pipeline – If I train an XGBClassifier without a pipeline, I can successfully generate predictions and explanations.
4. Session is Active – session.get_current_database() and session.get_current_schema() return valid values.
5. Feature & Label Columns Look Correct – FEATURE_COLS and LABEL_COLS contain expected values.

Additional Context:

• Environment:
  • Win 10
  • Python 3.9
  • snowflake-connector-python 3.14.0
  • snowflake-ml-python 1.7.4
  • snowflake-snowpark-python 1.28.0
• This example is based on Snowflake’s official documentation:
  Feature Preprocessing and Training on Non-Synthetic Data
• The documentation suggests using a Pipeline, but it does not provide an example of registering a Pipeline-based model that has explainability.
• The error message suggests that the model is somehow being treated as NoneType when passed to log_model().
• As mentioned the whole pipeline can be registered and retrieved when I remove this options={"enable_explainability": True} from .log_model.
• This options={"enable_explainability": True} works if I don't have a pipeline: https://stackoverflow.com/questions/79488423/snowflake-ml-there-is-no-method-with-name-explain-available-in-the-model

Question:

Is it possible to register a pipeline with options={"enable_explainability": True}?

[Masoud-Ghodrati]

guys I undestand that you added the following in the 1.7.5, but what you added I think is buggy:

• Support Python 3.12.
• Explainability: Support native and snowml sklearn pipeline

Now when you run predict, it passes target column in it too. And if you check on Snowsight, the list of columns in predict and predict_pra has target/label column in it. If you don't provide the target column, it will fail.

[sfc-gh-rroy]

Hi @Masoud-Ghodrati yes, we added pipeline support in 1.7.5.

Are you still seeing issues logging the model or afterwards when you call predict? Additionally, are you also seeing an issue while calling explain?

[sfc-gh-hayu]

guys I undestand that you added the following in the 1.7.5, but what you added I think is buggy:

• Support Python 3.12.
• Explainability: Support native and snowml sklearn pipeline

Now when you run predict, it passes target column in it too. And if you check on Snowsight, the list of columns in predict and predict_pra has target/label column in it. If you don't provide the target column, it will fail.

Hi @Masoud-Ghodrati the function signatures are inferred with sample_input_data=test_data in log_model, and test_data includes the label column. To infer the signatures correctly (w/o LABEL) you can do sample_input_data=test_data[FEATURE_COLS].

[Masoud-Ghodrati]

Thanks @sfc-gh-hayu, is this part of the new version (1.7.5)? Because it was working before, and there was no target column in Snowsight in version 1.7.4 (even with sample_input_data=test_data)?

Second question: Do we need to do the same for the predict method? Like mv.run(test_data[FEATURE_COLS], function_name="predict")

[sfc-gh-hayu]

@Masoud-Ghodrati Thanks for spotting this. This is unexpected. The workarounds is either using sample_input_data=test_data[FEATURE_COLS] or switching to open source sklearn and xgboost pipeline/models from the snowflake.ml.modeling ones which is the case we prioritize.

No, mv.run(test_data, function_name="predict") should be good.

[Masoud-Ghodrati]

Thanks @sfc-gh-hayu
It seems that a few things happening.

1. So, I upgraded to 1.7.5 and followed your suggestion about sample_input_data=test_data[FEATURE_COLS] and now both predict and explain work for this example:

# Construct pipeline
pipeline = Pipeline(steps=[
    ("impute", SimpleImputer(input_cols=FEATURE_COLS, output_cols=FEATURE_COLS)),
    ("scaler", StandardScaler(input_cols=FEATURE_COLS, output_cols=FEATURE_COLS)),
    (
        "grid_search_xgb",
        GridSearchCV(
            estimator=XGBClassifier(),
            param_grid={
                "n_estimators": [150],
                "learning_rate": [0.1],
                "max_depth": [5],
                "subsample": [1],
            },
            n_jobs=-1,
            scoring="accuracy",
            input_cols=FEATURE_COLS,
            label_cols=LABEL_COLS,
            output_cols=["OUTPUT_LABEL"],
            drop_input_cols=True,
            cv=RepeatedStratifiedKFold(n_splits=5, n_repeats=3, random_state=1),
        ),
    ),
])
# Train model
pipeline.fit(train_data)
# Model registration
native_registry = registry.Registry(session=session, database_name=session.get_current_database(), schema_name=session.get_current_schema())
model_name = "Gamma_test"
version = "V10"
model_version = native_registry.log_model(
    model=pipeline,
    model_name=model_name,
    version_name=version,
    sample_input_data=test_data[FEATURE_COLS],
    comment="Gamma test",
    conda_dependencies=["snowflake-ml-python==1.7.5", "snowflake-snowpark-python==1.28.0"],
    options={"enable_explainability": True}
)
# Run predictions and explanations
model_version = native_registry.get_model(model_name).version(version)
model_version.run(test_data[FEATURE_COLS], function_name="predict").show()
model_version.run(test_data[FEATURE_COLS], function_name="explain").show()

However, for our actual model, we get some error only for explain (prediction works). This is our model which I am working on:

# Define the ML pipeline with hyperparameter tuning
tuning_pipeline = Pipeline(
        steps=[
		("imputer", SimpleImputer(strategy="median", input_cols=num_columns, output_cols=num_columns)),
                ("scaler", StandardScaler(input_cols=num_columns, output_cols=num_columns)),
		("onehot", OneHotEncoder(input_cols=cat_columns, output_cols=cat_columns, drop_input_cols=True, handle_unknown="ignore")),
		("grid_search_xgb",
			GridSearchCV(
				estimator=XGBClassifier(
					random_state=123,
					deterministic=True
				),
				param_grid={
					"n_estimators": [150],
					"learning_rate": [0.1],
					"max_depth": [5],
					"subsample": [1],
					"random_state": [123]
				},
				n_jobs=-1,
				scoring="accuracy",
				label_cols=["TARGET"],
				output_cols=["PREDICTION"],
				drop_input_cols=True,
				cv=RepeatedStratifiedKFold(
					n_splits=5, n_repeats=3, random_state=1
				),
			),
        ),
    ]
)
# Train the pipeline
tuning_pipeline.fit(train_df)
native_registry = registry.Registry(
	session=session,
	database_name=session.get_current_database(),
	schema_name=session.get_current_schema()
)
model_name = "MY_AI_SNOWFLAKE_ML"
version_name = set_model_version(
	native_registry, model_name,
	major_version=True
)
model_description = "AI model trained using Snowflake ML."
model_version = native_registry.log_model(
	model=tuning_pipeline,
	model_name=model_name,
	version_name=version_name,
	sample_input_data=train_df.drop("TARGET").sample(n=1000),
	comment=model_description,
	options={"enable_explainability": True},
	conda_dependencies=[
		"snowflake-ml-python==1.7.5",
		"snowflake-snowpark-python==1.28.0"
	],
	task=type_hints.Task.TABULAR_MULTI_CLASSIFICATION,
)

and this is the is the error:

Traceback (most recent call last):
  File "/home/udf/73121818409322/snowflake-ml-python.zip/snowflake/ml/model/_packager/model_handlers/snowmlmodel.py", line 305, in explain_fn
    explainer = shap.Explainer(base_model, masker=background_data)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/shap/explainers/_explainer.py", line 174, in __init__
    raise TypeError("The passed model is not callable and cannot be analyzed directly with the given masker! Model: " + str(model))
TypeError: The passed model is not callable and cannot be analyzed directly with the given masker! Model: Pipeline(steps=[('imputer',
                 ColumnTransformer(remainder='passthrough',
                                   transformers=[('imputer',
                                                  SimpleImputer(strategy='median'),
                                                  [8, 11, 23, 24, 25, 12, 13,
                                                   14, 16])])),
                ('scaler',
                 ColumnTransformer(remainder='passthrough',
                                   transformers=[('scaler', StandardScaler(),
                                                  [0, 1, 2, 3, 4, 5, 6, 7,
                                                   8])])),
                ('onehot',
                 ColumnTransformer(remainder='passthrough',
                                   transformers=[('onehot...
                                                      max_cat_to_onehot=None,
                                                      max_delta_step=None,
                                                      max_depth=None,
                                                      max_leaves=None,
                                                      min_child_weight=None,
                                                      missing=nan,
                                                      monotone_constraints=None,
                                                      multi_strategy=None,
                                                      n_estimators=None,
                                                      n_jobs=3,
                                                      num_parallel_tree=None, ...),
                              n_jobs=-1,
                              param_grid={'learning_rate': [0.1],
                                          'max_depth': [5],
                                          'n_estimators': [150],
                                          'random_state': [123],
                                          'subsample': [1]},
                              scoring='accuracy'))])
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "/home/udf/73121818409322/explain.py", line 50, in process
    return runner(input_df[input_cols])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/udf/73121818409322/snowflake-ml-python.zip/snowflake/ml/model/_packager/model_handlers/snowmlmodel.py", line 315, in explain_fn
    X = X.astype(dtype_map, copy=False)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/pandas/core/generic.py", line 6620, in astype
    res_col = col.astype(dtype=cdt, copy=copy, errors=errors)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/pandas/core/generic.py", line 6643, in astype
    new_data = self._mgr.astype(dtype=dtype, copy=copy, errors=errors)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/pandas/core/internals/managers.py", line 430, in astype
    return self.apply(
           ^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/pandas/core/internals/managers.py", line 363, in apply
    applied = getattr(b, f)(**kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/pandas/core/internals/blocks.py", line 758, in astype
    new_values = astype_array_safe(values, dtype, copy=copy, errors=errors)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/pandas/core/dtypes/astype.py", line 237, in astype_array_safe
    new_values = astype_array(values, dtype, copy=copy)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/pandas/core/dtypes/astype.py", line 179, in astype_array
    values = values.astype(dtype, copy=copy)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/pandas/core/arrays/masked.py", line 586, in astype
    raise ValueError("cannot convert NA to integer")
ValueError: cannot convert NA to integer: Error while type casting for column 'OCC_COGNITIVE'
 in function V24_0_0.EXPLAIN with handler explain.infer

2. Second question is about df.random_split(weights=[0.9, 0.1], seed=0). My assumption is that the following assret should pass considering the seed. But it fails.

all_data = session.sql("SELECT *, IFF(CLASS = 'g', 1.0, 0.0) AS LABEL FROM Gamma_Telescope_Data").drop("CLASS")
train_data, test_data = all_data.random_split(weights=[0.9, 0.1], seed=0)
# assert split
train_data_1, test_data_1 = all_data.random_split(weights=[0.9, 0.1], seed=0)
train_data_2, test_data_2 = all_data.random_split(weights=[0.9, 0.1], seed=0)
assert train_data_1.collect == train_data_2.collect()
assert test_data_1.collect == test_data_2.collect()

[sfc-gh-hayu]

@Masoud-Ghodrati For 2, row orders are not guaranteed in Snowpark dataframes unless sorted. @sfc-gh-rroy Could you please help with 1?

[sfc-gh-rroy]

@Masoud-Ghodrati you have nulls or NaNs in the data you tried to run explain on? Looks like we enforce a data type conversion which is failing- you might need to filter those rows out.

[sfc-gh-hayu]

Update: sample_input_data is only enabled for modeling models in >=1.7.5. Prior to that it's actually ignored, and thus the behavior change above occurred.

[Masoud-Ghodrati]

Thanks @sfc-gh-rroy for your reply. However, this does not address the issue, as evidenced by the error message. To isolate the problem, I selected a row from my data source that contains no null/NaN values, yet the error persisted. The error indicates that the model cannot be found, leading me to suspect that Snowflake ML might have trouble handling one-hot encoding in the pipeline. Could you please take a closer look at this?

Additionally, thanks @sfc-gh-hayu for your reply regarding random_split. Could you clarify your suggestion? Are you recommending that we sort the data before performing the random split?

Error for explainability:

Traceback (most recent call last):
  File "/home/udf/73121818409322/snowflake-ml-python.zip/snowflake/ml/model/_packager/model_handlers/snowmlmodel.py", line 305, in explain_fn
    explainer = shap.Explainer(base_model, masker=background_data)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/shap/explainers/_explainer.py", line 174, in __init__
    raise TypeError("The passed model is not callable and cannot be analyzed directly with the given masker! Model: " + str(model))
TypeError: The passed model is not callable and cannot be analyzed directly with the given masker! Model: Pipeline(steps=[('imputer',
                 ColumnTransformer(remainder='passthrough',
                                   transformers=[('imputer',
                                                  SimpleImputer(strategy='median'),
                                                  [8, 11, 23, 24, 25, 12, 13,
                                                   14, 16])])),
                ('scaler',
                 ColumnTransformer(remainder='passthrough',
                                   transformers=[('scaler', StandardScaler(),
                                                  [0, 1, 2, 3, 4, 5, 6, 7,
                                                   8])])),
                ('onehot',
                 ColumnTransformer(remainder='passthrough',
                                   transformers=[('onehot...
                                                      max_cat_to_onehot=None,
                                                      max_delta_step=None,
                                                      max_depth=None,
                                                      max_leaves=None,
                                                      min_child_weight=None,
                                                      missing=nan,
                                                      monotone_constraints=None,
                                                      multi_strategy=None,
                                                      n_estimators=None,
                                                      n_jobs=3,
                                                      num_parallel_tree=None, ...),
                              n_jobs=-1,
                              param_grid={'learning_rate': [0.1],
                                          'max_depth': [5],
                                          'n_estimators': [150],
                                          'random_state': [123],
                                          'subsample': [1]},
                              scoring='accuracy'))])
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "/home/udf/73121818409322/snowflake-ml-python.zip/snowflake/ml/model/_packager/model_handlers/snowmlmodel.py", line 333, in explain_fn
    df = handlers_utils.convert_explanations_to_2D_df(base_model, explainer(X).values)
                                                                  ^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/shap/explainers/_permutation.py", line 77, in __call__
    return super().__call__(
           ^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/shap/explainers/_explainer.py", line 266, in __call__
    row_result = self.explain_row(
                 ^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/shap/explainers/_permutation.py", line 85, in explain_row
    fm = MaskedModel(self.model, self.masker, self.link, self.linearize_link, *row_args)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/shap/utils/_masked_model.py", line 31, in __init__
    self._variants = ~self.masker.invariants(*args)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/shap/maskers/_tabular.py", line 150, in invariants
    return np.isclose(x, self.data)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python_udf/0b4bf6a6d44e8b7c0fe9e188813a348a7ab4818b4c72b879863a554529a9fc3c/lib/python3.11/site-packages/numpy/core/numeric.py", line 2348, in isclose
    xfin = isfinite(x)
           ^^^^^^^^^^^
TypeError: ufunc 'isfinite' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "/home/udf/73121818409322/explain.py", line 50, in process
    return runner(input_df[input_cols])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/udf/73121818409322/snowflake-ml-python.zip/snowflake/ml/model/_packager/model_handlers/snowmlmodel.py", line 335, in explain_fn
    raise ValueError(f"Explanation for this model type not supported yet: {str(e)}")
ValueError: Explanation for this model type not supported yet: ufunc 'isfinite' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''
 in function V24_0_0.EXPLAIN with handler explain.infer

[sfc-gh-hayu]

@Masoud-Ghodrati If you need consistent random_split results then you'll have to sort before the split. Otherwise you are splitting a dataframe whose order is not guaranteed, so each time you split, the train/test data can be different.

[sfc-gh-rroy]

Hi @Masoud-Ghodrati, even though the explain call failed again, your error stacktrace is different here- this is the one where you only sent a single row to explain? Looks like something is wrong with the masker in this case- any chance your sample_input_data that nulls/NaNs? Can you try filtering those and passing the sample_input_data in that case?

[Masoud-Ghodrati]

I have tested with public data which does not have any missing values. This is not the problem of NaNs/Nulls/Missing values. The problem is that the Pipeline in Snowflake ML has some problems with data transform with one hot coding. As soon as you remove one-hot from the Pipeline everything works.

This example works:

pipeline = Pipeline(steps=[
    ("impute", SimpleImputer(input_cols=NUMERIC_FEATURE_COLS, output_cols=NUMERIC_FEATURE_COLS)),
    ("scaler", StandardScaler(input_cols=NUMERIC_FEATURE_COLS, output_cols=NUMERIC_FEATURE_COLS)),
    (
        "grid_search_xgb",
        GridSearchCV(
            estimator=XGBClassifier(),
            param_grid={
                "n_estimators": [150],
                "learning_rate": [0.1],
                "max_depth": [5],
                "subsample": [1],
            },
            n_jobs=-1,
            scoring="accuracy",
            label_cols=LABEL_COLS,
            output_cols=["OUTPUT_LABEL"],
            drop_input_cols=True,
            cv=RepeatedStratifiedKFold(n_splits=5, n_repeats=3, random_state=1),
        ),
    ),
])

While this does not (note that I added a dummy cat column to the data):

pipeline = Pipeline(steps=[
    ("impute", SimpleImputer(input_cols=NUMERIC_FEATURE_COLS, output_cols=NUMERIC_FEATURE_COLS)),
    ("scaler", StandardScaler(input_cols=NUMERIC_FEATURE_COLS, output_cols=NUMERIC_FEATURE_COLS)),
    ("onehot", OneHotEncoder(input_cols=CATEGORICAL_COLS, output_cols=CATEGORICAL_COLS, drop_input_cols=True, handle_unknown="ignore")),  # One-hot encoding step
    (
        "grid_search_xgb",
        GridSearchCV(
            estimator=XGBClassifier(),
            param_grid={
                "n_estimators": [150],
                "learning_rate": [0.1],
                "max_depth": [5],
                "subsample": [1],
            },
            n_jobs=-1,
            scoring="accuracy",
            label_cols=LABEL_COLS,
            output_cols=["OUTPUT_LABEL"],
            drop_input_cols=True,
            cv=RepeatedStratifiedKFold(n_splits=5, n_repeats=3, random_state=1),
        ),
    ),
])

[sfc-gh-rroy]

When you skip the OHE- you are only using numeric features right? My hypothesis is that this specific codepath in the shap library cannot handle categorical data- we will verify this.