From 88c26fbe4ad5cc7d422d01f5624804cb2a03deed Mon Sep 17 00:00:00 2001 From: Fabio Date: Mon, 18 Mar 2024 14:57:59 +0100 Subject: [PATCH] Add exception when unpickling --- hiclass/LocalClassifierPerLevel.py | 4 ++-- hiclass/LocalClassifierPerNode.py | 4 ++-- hiclass/LocalClassifierPerParentNode.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/hiclass/LocalClassifierPerLevel.py b/hiclass/LocalClassifierPerLevel.py index db819d78..896828cd 100644 --- a/hiclass/LocalClassifierPerLevel.py +++ b/hiclass/LocalClassifierPerLevel.py @@ -264,8 +264,8 @@ def _fit_classifier(self, level, separator): f"Loaded trained model for local classifier {level} from file {filename}" ) return classifier - except pickle.UnpicklingError: - self.logger_.warning(f"Could not load model from file {filename}") + except (pickle.UnpicklingError, EOFError): + self.logger_.error(f"Could not load model from file {filename}") self.logger_.info(f"Training local classifier {level}") X, y, sample_weight = self._remove_empty_leaves( separator, self.X_, self.y_[:, level], self.sample_weight_ diff --git a/hiclass/LocalClassifierPerNode.py b/hiclass/LocalClassifierPerNode.py index 7e76862e..65c14113 100644 --- a/hiclass/LocalClassifierPerNode.py +++ b/hiclass/LocalClassifierPerNode.py @@ -255,8 +255,8 @@ def _fit_classifier(self, node): f"Loaded trained model for local classifier {node.split(self.separator_)[-1]} from file {filename}" ) return classifier - except pickle.UnpicklingError: - self.logger_.warning(f"Could not load model from file {filename}") + except (pickle.UnpicklingError, EOFError): + self.logger_.error(f"Could not load model from file {filename}") self.logger_.info( f"Training local classifier {str(node).split(self.separator_)[-1]}" ) diff --git a/hiclass/LocalClassifierPerParentNode.py b/hiclass/LocalClassifierPerParentNode.py index 439c5e75..5873b52a 100644 --- a/hiclass/LocalClassifierPerParentNode.py +++ b/hiclass/LocalClassifierPerParentNode.py @@ -224,8 +224,8 @@ def _fit_classifier(self, node): f"Loaded trained model for local classifier {node.split(self.separator_)[-1]} from file {filename}" ) return classifier - except pickle.UnpicklingError: - self.logger_.warning(f"Could not load model from file {filename}") + except (pickle.UnpicklingError, EOFError): + self.logger_.error(f"Could not load model from file {filename}") self.logger_.info( f"Training local classifier {str(node).split(self.separator_)[-1]}" )