Update app/configs.py

app/configs.py

@@ -1,5 +1,5 @@
 import mlflow
-import mlflow.tensorflow
+import mlflow.pyfunc
 
 from pathlib import Path
 from typing import Dict, Any, List
@@ -20,7 +20,7 @@ CLASSIFICATION_MODEL_URI = "runs:/7766854bcd074c8e856aea7c47680bf4/model"
 SEGMENTATION_MODEL_URI = "runs:/35f0bea9bf774ed59f2298a897c8cdf5/model"
 
 # ======================
-# Lazy Loaded Models
+# Cached models
 # ======================
 
 _classification_model = None
@@ -28,7 +28,7 @@ _segmentation_model = None
 
 
 # ======================
-# Classification Model
+# Classification
 # ======================
 
 def get_classification_model():
@@ -37,17 +37,17 @@ def get_classification_model():
     if _classification_model is None:
         print("Loading Classification Model...")
 
-        _classification_model = mlflow.tensorflow.load_model(
+        _classification_model = mlflow.pyfunc.load_model(
             CLASSIFICATION_MODEL_URI
         )
 
-        print("Classification Model Loaded Successfully")
+        print("Classification Model Loaded")
 
     return _classification_model
 
 
 # ======================
-# Segmentation Model
+# Segmentation
 # ======================
 
 def get_segmentation_model():
@@ -56,11 +56,11 @@ def get_segmentation_model():
     if _segmentation_model is None:
         print("Loading Segmentation Model...")
 
-        _segmentation_model = mlflow.tensorflow.load_model(
+        _segmentation_model = mlflow.pyfunc.load_model(
             SEGMENTATION_MODEL_URI
         )
 
-        print("Segmentation Model Loaded Successfully")
+        print("Segmentation Model Loaded")
 
     return _segmentation_model
 
@@ -74,15 +74,11 @@ model_classes = {
     1: "malignant"
 }
 
-# ======================
-# Request History
-# ======================
-
 request_history: List[Dict[str, Any]] = []
 
 
 # ======================
-# Storage Setup
+# Storage
 # ======================
 
 STORAGE_DIR = Path("storage")