Update predict_utils.py

predict_utils.py

@@ -1,19 +1,17 @@
-# predict_utils.py
+# predict_utils.py (patched to handle XGBClassifier use_label_encoder issue + robust loader)
 import os
 import logging
 import joblib
 from huggingface_hub import hf_hub_download
 
-# Standard logging
+# Logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Env vars
 HF_MODEL_REPO = os.getenv("HF_MODEL_REPO", "sathishaiuse/wellness-classifier-model")
 HF_MODEL_FILENAME = os.getenv("HF_MODEL_FILENAME", "best_overall_XGBoost.joblib")
 HF_TOKEN = os.getenv("HF_TOKEN") or None
 
-# Local candidate paths to look for the model file
 LOCAL_CANDIDATES = [
     os.path.join("/app", HF_MODEL_FILENAME),
     os.path.join("/tmp", HF_MODEL_FILENAME),
@@ -22,7 +20,32 @@ LOCAL_CANDIDATES = [
 ]
 
 # -------------------------
-# Helpers: inspect, try loaders
+# Monkey-patch xgboost sklearn wrappers to add missing attributes before unpickling
+# This prevents errors like: "'XGBClassifier' object has no attribute 'use_label_encoder'"
+# -------------------------
+def ensure_xgb_sklearn_compat():
+    try:
+        import xgboost as xgb
+        # XGBClassifier
+        clf = getattr(xgb, "XGBClassifier", None)
+        if clf is not None:
+            if not hasattr(clf, "use_label_encoder"):
+                setattr(clf, "use_label_encoder", False)
+                logger.info("Patched XGBClassifier.use_label_encoder = False")
+        # XGBRegressor
+        reg = getattr(xgb, "XGBRegressor", None)
+        if reg is not None:
+            if not hasattr(reg, "use_label_encoder"):
+                setattr(reg, "use_label_encoder", False)
+                logger.info("Patched XGBRegressor.use_label_encoder = False")
+    except Exception as e:
+        logger.debug(f"xgboost not available to patch: {e}")
+
+# Call the patch early so joblib.load can succeed
+ensure_xgb_sklearn_compat()
+
+# -------------------------
+# Helpers
 # -------------------------
 def inspect_file(path):
     info = {"path": path, "exists": False}
@@ -44,6 +67,8 @@ def inspect_file(path):
 
 def try_joblib_load(path):
     try:
+        # Ensure patch just before load (in case xgboost gets imported lazily)
+        ensure_xgb_sklearn_compat()
         logger.info(f"Trying joblib.load on {path}")
         m = joblib.load(path)
         logger.info("joblib.load succeeded")
@@ -71,13 +96,11 @@ def try_xgboost_booster(path):
                 self._is_xgb_booster = True
 
             def predict(self, X):
-                # X -> 2D list/array
                 import numpy as _np, xgboost as _xgb
                 arr = _np.array(X, dtype=float)
                 dmat = _xgb.DMatrix(arr)
                 pred = self.booster.predict(dmat)
-                # binary prob -> class decision
-                if pred.ndim == 1:
+                if hasattr(pred, "ndim") and pred.ndim == 1:
                     return (_np.where(pred >= 0.5, 1, 0)).tolist()
                 return pred.tolist()
 
@@ -86,7 +109,7 @@ def try_xgboost_booster(path):
                 arr = _np.array(X, dtype=float)
                 dmat = _xgb.DMatrix(arr)
                 pred = self.booster.predict(dmat)
-                if pred.ndim == 1:
+                if hasattr(pred, "ndim") and pred.ndim == 1:
                     return (_np.vstack([1 - pred, pred]).T).tolist()
                 return pred.tolist()
 
@@ -96,7 +119,7 @@ def try_xgboost_booster(path):
         return ("xgboost_booster", e)
 
 # -------------------------
-# Core loader
+# Loader
 # -------------------------
 def load_model():
     logger.info("==== MODEL LOAD START ====")
@@ -104,7 +127,6 @@ def load_model():
     logger.info(f"Filename: {HF_MODEL_FILENAME}")
     logger.info(f"HF_TOKEN present? {bool(HF_TOKEN)}")
 
-    # Try local candidates
     for path in LOCAL_CANDIDATES:
         try:
             info = inspect_file(path)
@@ -112,12 +134,10 @@ def load_model():
             if not info.get("exists"):
                 continue
 
-            # try joblib
             t, res = try_joblib_load(path)
             if t == "joblib" and not isinstance(res, Exception):
                 return res
 
-            # try xgboost booster
             t, res = try_xgboost_booster(path)
             if t == "xgboost_booster" and not isinstance(res, Exception):
                 return res
@@ -125,7 +145,6 @@ def load_model():
         except Exception as e:
             logger.exception(f"Error while trying local candidate {path}: {e}")
 
-    # Try HF hub download
     try:
         logger.info(f"Trying hf_hub_download from {HF_MODEL_REPO}/{HF_MODEL_FILENAME}")
         model_path = hf_hub_download(repo_id=HF_MODEL_REPO, filename=HF_MODEL_FILENAME, token=HF_TOKEN)
@@ -148,40 +167,30 @@ def load_model():
         return None
 
 # -------------------------
-# Prediction helper (robust)
-# Accepts: features as dict, list, or list-of-lists
-# Ensures sklearn pipelines that need DataFrame get a pandas.DataFrame
+# Robust predict (creates DataFrame when model expects column names)
 # -------------------------
 def predict(model, features):
     """
-    model: object returned by load_model()
-    features: dict (preferred) mapping column_name -> value (order preserved),
-              OR list/tuple representing feature vector in correct order,
-              OR list-of-lists for batch.
-    Returns: {"prediction": ..., "probability": ...} or {"error": "..."}
+    Accepts dict (col->val), list, or list-of-lists.
+    Returns dict with prediction and probability, or error.
     """
     if model is None:
         return {"error": "Model not loaded"}
 
     try:
-        # detect xgboost booster wrapper (we set attribute _is_xgb_booster)
-        is_booster = hasattr(model, "_is_xgb_booster")
-
-        # prepare input for sklearn-pipeline style models: DataFrame with column names
         import pandas as _pd
         import numpy as _np
 
-        # Case A: features is a dict -> preserve key order and create single-row DataFrame
+        is_booster = hasattr(model, "_is_xgb_booster")
+
+        # dict -> DataFrame with columns in order of keys
         if isinstance(features, dict):
-            # ensure keys are strings (column names the pipeline expects)
             col_names = [str(k) for k in features.keys()]
             row_values = [features[k] for k in features.keys()]
-            # Create DataFrame preserving column order
             df = _pd.DataFrame([row_values], columns=col_names)
             logger.info(f"Prepared DataFrame for prediction with columns: {col_names}")
 
             if is_booster:
-                # booster expects numeric array
                 arr = df.values.astype(float)
                 preds = model.predict(arr)
                 prob = None
@@ -194,7 +203,6 @@ def predict(model, features):
                 pred_val = int(preds[0]) if isinstance(preds, (list, tuple)) else int(preds)
                 return {"prediction": pred_val, "probability": prob}
 
-            # sklearn-like pipeline
             if hasattr(model, "predict"):
                 pred = model.predict(df)[0]
                 prob = None
@@ -204,7 +212,6 @@ def predict(model, features):
                         prob = float(max(p))
                     except:
                         prob = None
-                # convert numpy types to native
                 try:
                     pred = int(pred)
                 except:
@@ -213,9 +220,9 @@ def predict(model, features):
 
             return {"error": "Loaded model object not recognized (no predict method)"}
 
-        # Case B: features is list or tuple -> single row without column names
+        # list -> numpy array single row
+        import numpy as _np
         if isinstance(features, (list, tuple)):
-            # single-row list
             arr2d = _np.array([features], dtype=float)
             if is_booster:
                 preds = model.predict(arr2d)
@@ -229,9 +236,6 @@ def predict(model, features):
                 pred_val = int(preds[0]) if isinstance(preds, (list, tuple)) else int(preds)
                 return {"prediction": pred_val, "probability": prob}
 
-            # sklearn pipeline without column names -> create DataFrame with numeric column names
-            # but many scikit-learn ColumnTransformer setups expect string column names; this is risky.
-            # Try passing numpy array directly to predict() if pipeline accepts it.
             if hasattr(model, "predict"):
                 try:
                     pred = model.predict(arr2d)[0]
@@ -243,8 +247,7 @@ def predict(model, features):
                         except:
                             prob = None
                     return {"prediction": pred, "probability": prob}
-                except Exception as e:
-                    # as last resort, build DataFrame with string column names "0","1",... and hope pipeline uses positional selection
+                except Exception:
                     cols = [str(i) for i in range(arr2d.shape[1])]
                     df = _pd.DataFrame(arr2d, columns=cols)
                     pred = model.predict(df)[0]
@@ -257,7 +260,7 @@ def predict(model, features):
                             prob = None
                     return {"prediction": pred, "probability": prob}
 
-        # Case C: features is list-of-lists (batch)
+        # batch
         if isinstance(features, list) and len(features) > 0 and isinstance(features[0], (list, tuple)):
             arr = _np.array(features, dtype=float)
             if is_booster:
@@ -281,8 +284,7 @@ def predict(model, features):
                         except:
                             prob = None
                     return {"prediction": pred.tolist(), "probability": prob}
-                except Exception as e:
-                    # try DataFrame fallback
+                except Exception:
                     cols = [str(i) for i in range(arr.shape[1])]
                     df = _pd.DataFrame(arr, columns=cols)
                     pred = model.predict(df)
@@ -296,6 +298,7 @@ def predict(model, features):
                     return {"prediction": pred.tolist(), "probability": prob}
 
         return {"error": "Unsupported features format. Provide dict (col->val) or list of values."}
+
     except Exception as e:
         logger.exception(f"Prediction error: {e}")
         return {"error": str(e)}