fix: check that prediction_proba is not None before checking its type

radicalbit · Jul 1, 2024 · 5508730 · 5508730
1 parent 7b33b62
commit 5508730
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/api/app/models/model_dto.py b/api/app/models/model_dto.py
@@ -99,7 +99,9 @@ def validate_outputs(self) -> Self:
                     raise ValueError(
                         f'prediction must be a number for a ModelType.BINARY, has been provided [{self.outputs.prediction}]'
                     )
-                if not is_optional_float(self.outputs.prediction_proba.type):
+                if not is_none(self.outputs.prediction_proba) and not is_optional_float(
+                    self.outputs.prediction_proba.type
+                ):
                     raise ValueError(
                         f'prediction_proba must be an optional float for a ModelType.BINARY, has been provided [{self.outputs.prediction_proba}]'
                     )
@@ -109,7 +111,9 @@ def validate_outputs(self) -> Self:
                     raise ValueError(
                         f'prediction must be a number or string for a ModelType.MULTI_CLASS, has been provided [{self.outputs.prediction}]'
                     )
-                if not is_optional_float(self.outputs.prediction_proba.type):
+                if not is_none(self.outputs.prediction_proba) and not is_optional_float(
+                    self.outputs.prediction_proba.type
+                ):
                     raise ValueError(
                         f'prediction_proba must be an optional float for a ModelType.MULTI_CLASS, has been provided [{self.outputs.prediction_proba}]'
                     )
@@ -119,7 +123,7 @@ def validate_outputs(self) -> Self:
                     raise ValueError(
                         f'prediction must be a number for a ModelType.REGRESSION, has been provided [{self.outputs.prediction}]'
                     )
-                if not is_none(self.outputs.prediction_proba.type):
+                if not is_none(self.outputs.prediction_proba) and not is_none(self.outputs.prediction_proba.type):
                     raise ValueError(
                         f'prediction_proba must be None for a ModelType.REGRESSION, has been provided [{self.outputs.prediction_proba}]'
                     )