updated fastapi app script

app.py

@@ -10,15 +10,16 @@ from pydantic import BaseModel
 from config import settings
 
 try:
-    path_mlflow_model = "trained_models/knn_random_forest"
+    path_mlflow_model = "./model_for_production/"
     sklearn_pipeline = mlflow.sklearn.load_model(path_mlflow_model)
 except:
-    path_mlflow_model = "/data/models/knn_random_forest"
+    path_mlflow_model = "/data/model_for_production/"
     sklearn_pipeline = mlflow.sklearn.load_model(path_mlflow_model)
 
 app = FastAPI()
 logging.basicConfig(level=logging.INFO)
 
+
 class WaterPotabilityDataItem(BaseModel):
     ph: Union[float, None] = np.nan
     Hardness: Union[float, None] = np.nan
@@ -30,20 +31,53 @@ class WaterPotabilityDataItem(BaseModel):
     Trihalomethanes: Union[float, None] = np.nan
     Turbidity: Union[float, None] = np.nan
 
+
 def predict_pipeline(data_sample):
+    """
+    ---------
+    Arguments
+    ---------
+    data_sample : np.array
+        a numpy array of shape (num_samples, num_feats)
+
+    -------
+    Returns
+    -------
+    pred_sample : np.array
+        a numpy array of shape (num_samples) with predictions
+    """
     pred_sample = sklearn_pipeline.predict(data_sample)
     return pred_sample
 
+
 @app.get("/info")
 def get_app_info():
-    dict_info = {
-        "app_name": settings.app_name,
-        "version": settings.version
-    }
+    """
+    -------
+    Returns
+    -------
+    dict_info : dict
+        a dictionary with info to be sent as a response to get request
+    """
+    dict_info = {"app_name": settings.app_name, "version": settings.version}
     return dict_info
 
+
 @app.post("/predict")
 def predict(wpd_item: WaterPotabilityDataItem):
+    """
+    ---------
+    Arguments
+    ---------
+    wpd_item : object
+        an object of type WaterPotabilityDataItem
+
+    -------
+    Returns
+    -------
+    pred_dict : dict
+        a dictionary of prediction to be sent as a response to post request
+    """
     wpd_arr = np.array(
         [
             wpd_item.ph,
@@ -60,4 +94,5 @@ def predict(wpd_item: WaterPotabilityDataItem):
     logging.info("data sample: %s", wpd_arr)
     pred_sample = predict_pipeline(wpd_arr)
     logging.info("Potability prediction: %s", pred_sample)
-    return {"Potability": int(pred_sample)}
+    pred_dict = {"Potability": int(pred_sample)}
+    return pred_dict