fix: ensure to reseek before loading segmentation models if needed

iamDecode · Nov 5, 2023 · a4bed07 · a4bed07
1 parent e1c4bde
commit a4bed07
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/sklearn_pmml_model/auto_detect/base.py b/sklearn_pmml_model/auto_detect/base.py
@@ -54,8 +54,12 @@ def parse(file: Iterator, seek=False):
 
         if all(clf is PMMLTreeClassifier or clf is PMMLLogisticRegression for clf in clfs):
           if 'multipleModelMethod="majorityVote"' in line or 'multipleModelMethod="average"' in line:
+            if seek:
+              pmml.seek(0)
             return PMMLForestClassifier(pmml=pmml, **kwargs)
           if 'multipleModelMethod="modelChain"' in line:
+            if seek:
+              pmml.seek(0)
             return PMMLGradientBoostingClassifier(pmml=pmml, **kwargs)
 
         raise Exception('Unsupported PMML classifier: invalid segmentation.')
@@ -94,8 +98,12 @@ def parse(file: Iterator, seek=False):
 
         if all(reg is PMMLTreeRegressor or reg is PMMLLinearRegression for reg in regs):
           if 'multipleModelMethod="majorityVote"' in line or 'multipleModelMethod="average"' in line:
+            if seek:
+              pmml.seek(0)
             return PMMLForestRegressor(pmml=pmml, **kwargs)
           if 'multipleModelMethod="sum"' in line:
+            if seek:
+              pmml.seek(0)
             return PMMLGradientBoostingRegressor(pmml=pmml, **kwargs)
 
         raise Exception('Unsupported PMML regressor: invalid segmentation.')