bearking58 commited on
Commit
cacd064
·
1 Parent(s): 279839c

fix: 2nd model returning array (change to float)

Browse files
core-model-prediction/secondary_model.py CHANGED
@@ -20,5 +20,5 @@ class SecondaryModel:
20
  features_df[self.secondary_model_features])
21
  return features_df
22
 
23
- def predict(self, secondary_model_features: List[float]):
24
- return int(self.model.predict(self.preprocess_input(secondary_model_features))[0])
 
20
  features_df[self.secondary_model_features])
21
  return features_df
22
 
23
+ def predict(self, secondary_model_features: List[float]) -> float:
24
+ return self.model.predict_proba(self.preprocess_input(secondary_model_features))[:, -1][0]