Adjust for earlier scikit-learn versions

Signed-off-by: Ben Wilson <[email protected]>
mlflow · Sep 3, 2024 · d1f5fca · d1f5fca
1 parent 5bf79c8
commit d1f5fca
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/regression/steps/transform.py b/regression/steps/transform.py
@@ -41,6 +41,12 @@ def transformer_fn():
         else {"feature_names_out": "one-to-one"}
     )
 
+    onehot_params = (
+        {"sparse_output": False}
+        if Version(sklearn.__version__) >= Version("1.2")
+        else {"sparse": False}
+    )
+
     return Pipeline(
         steps=[
             (
@@ -53,12 +59,12 @@ def transformer_fn():
                     transformers=[
                         (
                             "hour_encoder",
-                            OneHotEncoder(categories="auto", sparse_output=False),
+                            OneHotEncoder(categories="auto", **onehot_params),
                             ["pickup_hour"],
                         ),
                         (
                             "day_encoder",
-                            OneHotEncoder(categories="auto", sparse_output=False),
+                            OneHotEncoder(categories="auto", **onehot_params),
                             ["pickup_dow"],
                         ),
                         (