param ordering with defaults

microsoft · Apr 5, 2023 · 866422a · 866422a
1 parent 0711d50
commit 866422a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/hummingbird/ml/operator_converters/_one_hot_encoder_implementations.py b/hummingbird/ml/operator_converters/_one_hot_encoder_implementations.py
@@ -22,7 +22,7 @@ class OneHotEncoderString(PhysicalOperator, torch.nn.Module):
     Because we are dealing with tensors, strings require additional length information for processing.
     """
 
-    def __init__(self, logical_operator, categories, handle_unknown, device, infrequent=None, extra_config={}):
+    def __init__(self, logical_operator, categories, device, extra_config={}, handle_unknown='error', infrequent=None):
         super(OneHotEncoderString, self).__init__(logical_operator, transformer=True)
 
         self.num_columns = len(categories)
@@ -113,7 +113,7 @@ class OneHotEncoder(PhysicalOperator, torch.nn.Module):
     Class implementing OneHotEncoder operators for ints in PyTorch.
     """
 
-    def __init__(self, logical_operator, categories, handle_unknown, device, infrequent=None):
+    def __init__(self, logical_operator, categories, device, handle_unknown='error', infrequent=None):
         super(OneHotEncoder, self).__init__(logical_operator, transformer=True)
 
         self.num_columns = len(categories)

diff --git a/hummingbird/ml/operator_converters/sklearn/one_hot_encoder.py b/hummingbird/ml/operator_converters/sklearn/one_hot_encoder.py
@@ -49,11 +49,12 @@ def convert_sklearn_one_hot_encoder(operator, device, extra_config):
         ]
     ):
         categories = [[str(x) for x in c.tolist()] for c in operator.raw_operator.categories_]
-        return OneHotEncoderString(operator, categories, operator.raw_operator.handle_unknown,
-                                   device, infrequent, extra_config)
+        return OneHotEncoderString(operator, categories, device, extra_config=extra_config,
+                                    handle_unknown=operator.raw_operator.handle_unknown,
+                                    infrequent=infrequent)
     else:
-        return OneHotEncoder(operator, operator.raw_operator.categories_, operator.raw_operator.handle_unknown,
-                             device, infrequent)
+        return OneHotEncoder(operator, operator.raw_operator.categories_, device,
+                             handle_unknown=operator.raw_operator.handle_unknown, infrequent=infrequent)
 
 
 register_converter("SklearnOneHotEncoder", convert_sklearn_one_hot_encoder)