Merge pull request #56 from mlflow/fix-transform

Adjust for earlier scikit-learn versions
mlflow · Sep 4, 2024 · b2bf60e · b2bf60e
2 parents 5bf79c8 + e608f19
commit b2bf60e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/pylintrc b/pylintrc
@@ -63,15 +63,10 @@ disable=print-statement,
         unpacking-in-except,
         old-raise-syntax,
         backtick,
-        long-suffix,
-        old-ne-operator,
-        old-octal-literal,
         import-star-module-level,
-        non-ascii-bytes-literal,
         raw-checker-failed,
         bad-inline-option,
         locally-disabled,
-        locally-enabled,
         file-ignored,
         suppressed-message,
         useless-suppression,
@@ -115,7 +110,6 @@ disable=print-statement,
         range-builtin-not-iterating,
         filter-builtin-not-iterating,
         using-cmp-argument,
-        eq-without-hash,
         div-method,
         idiv-method,
         rdiv-method,
@@ -165,6 +159,7 @@ disable=print-statement,
         not-an-iterable,
         unsubscriptable-object,
         unsupported-assignment-operation,
+        undefined-variable,
 
 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option

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"],
                         ),
                         (