Make model not found uniform

Signed-off-by: Frederik Stoel <67371523+wfstoel@users.noreply.github.com>
OpenSTEF · Sep 12, 2023 · 296e1ce · 296e1ce
1 parent deb9eb0
commit 296e1ce
Showing 1 changed file with 3 additions and 9 deletions.
diff --git a/openstef/model/serializer.py b/openstef/model/serializer.py
@@ -147,10 +147,7 @@ def load_model(
             if not models_df.empty:
                 latest_run = models_df.iloc[0]  # Use .iloc[0] to only get latest run
             else:
-                raise LookupError(
-                    f"Model not found for experiment_name {experiment_name}. First"
-                    " train a model!"
-                )
+                raise LookupError("Model not found. First train a model!")
             model_uri = self._get_model_uri(latest_run.artifact_uri)
             loaded_model = mlflow.sklearn.load_model(model_uri)
             loaded_model.age = self._determine_model_age_from_mlflow_run(latest_run)
@@ -162,11 +159,8 @@ def load_model(
             )  # Path without file:///
             self.logger.info("Model successfully loaded with MLflow")
             return loaded_model, model_specs
-        except (AttributeError, MlflowException, OSError) as e:
-            raise AttributeError(
-                f"Model not found for experiment_name {experiment_name}. First train a"
-                " model!"
-            ) from e
+        except (AttributeError, MlflowException, OSError) as exception:
+            raise LookupError("Model not found. First train a model!") from exception
 
     def get_model_age(
         self, experiment_name: str, hyperparameter_optimization_only: bool = False