From 4ca02fd106160c35bd67150f1137be7123b580ad Mon Sep 17 00:00:00 2001 From: Jan-Benedikt Jagusch Date: Fri, 30 Jul 2021 13:53:17 +0300 Subject: [PATCH] clean up startswith code Signed-off-by: Jan-Benedikt Jagusch --- onnxmltools/convert/lightgbm/_parse.py | 4 +--- onnxmltools/convert/lightgbm/operator_converters/LightGbm.py | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/onnxmltools/convert/lightgbm/_parse.py b/onnxmltools/convert/lightgbm/_parse.py index e527faf7..d711abe0 100644 --- a/onnxmltools/convert/lightgbm/_parse.py +++ b/onnxmltools/convert/lightgbm/_parse.py @@ -27,9 +27,7 @@ def __init__(self, booster): if (_model_dict['objective'].startswith('binary') or _model_dict['objective'].startswith('multiclass')): self.operator_name = 'LgbmClassifier' - elif (_model_dict['objective'].startswith('regression') or - _model_dict['objective'].startswith('poisson') or - _model_dict['objective'].startswith('gamma')): + elif _model_dict['objective'].startswith(('regression', 'poisson', 'gamma')): self.operator_name = 'LgbmRegressor' else: # Other objectives are not supported. diff --git a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py index c2ad360c..242c2c1d 100644 --- a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py +++ b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py @@ -234,7 +234,7 @@ def convert_lightgbm(scope, operator, container): n_classes = 1 # Regressor has only one output variable attrs['post_transform'] = 'NONE' attrs['n_targets'] = n_classes - elif gbm_text['objective'].startswith('poisson') or gbm_text['objective'].startswith('gamma'): + elif gbm_text['objective'].startswith(('poisson', 'gamma')): n_classes = 1 # Regressor has only one output variable attrs['n_targets'] = n_classes # 'Exp' is not a supported post_transform value in the ONNX spec yet,