FranciscoLozDataScience commited on
Commit
7f67122
·
1 Parent(s): 11c425a

return the str label not the array

Browse files
Files changed (1) hide show
  1. model.py +7 -5
model.py CHANGED
@@ -12,6 +12,7 @@ class SmokerModel:
12
  def __init__(self, model_path, scaler_path):
13
  self.model = load(model_path)
14
  self.scaler = load(scaler_path)
 
15
 
16
  def scale(self, X):
17
  """
@@ -30,7 +31,7 @@ class SmokerModel:
30
 
31
  return new_data_scaled
32
 
33
- def predict(self, X: np.ndarray):
34
  """
35
  Make a prediction on one sample using the loaded model.
36
 
@@ -55,8 +56,9 @@ class SmokerModel:
55
  X_scaled = self.scale(X)
56
 
57
  # Now, use the scaled data to make predictions using the loaded model
58
- predicted_label = self.model.predict(X_scaled)
59
 
60
- print(predicted_label)
61
-
62
- return predicted_label
 
 
12
  def __init__(self, model_path, scaler_path):
13
  self.model = load(model_path)
14
  self.scaler = load(scaler_path)
15
+ self.labels = ["non-smoker", "smoker"]
16
 
17
  def scale(self, X):
18
  """
 
31
 
32
  return new_data_scaled
33
 
34
+ def predict(self, X: np.ndarray) -> str:
35
  """
36
  Make a prediction on one sample using the loaded model.
37
 
 
56
  X_scaled = self.scale(X)
57
 
58
  # Now, use the scaled data to make predictions using the loaded model
59
+ array = self.model.predict(X_scaled)
60
 
61
+ #predict
62
+ predicted_label = array[0]
63
+ str_label = self.labels[predicted_label]
64
+ return str_label