Update predict_utils.py

predict_utils.py

@@ -1,21 +1,19 @@
+# predict_utils.py
 import os
-import joblib
 import logging
+import joblib
 from huggingface_hub import hf_hub_download
 
-# -----------------------------------------------------------
-# Logging Setup
-# -----------------------------------------------------------
+# Standard logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# -----------------------------------------------------------
-# Environment Variables
-# -----------------------------------------------------------
+# 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),
@@ -23,62 +21,281 @@ LOCAL_CANDIDATES = [
     HF_MODEL_FILENAME
 ]
 
-# -----------------------------------------------------------
-# Model Loader
-# -----------------------------------------------------------
+# -------------------------
+# Helpers: inspect, try loaders
+# -------------------------
+def inspect_file(path):
+    info = {"path": path, "exists": False}
+    try:
+        info["exists"] = os.path.exists(path)
+        if not info["exists"]:
+            return info
+        info["size"] = os.path.getsize(path)
+        with open(path, "rb") as f:
+            head = f.read(1024)
+        info["head_bytes"] = head
+        try:
+            info["head_text"] = head.decode("utf-8", errors="replace")
+        except:
+            info["head_text"] = None
+    except Exception as e:
+        info["inspect_error"] = str(e)
+    return info
+
+def try_joblib_load(path):
+    try:
+        logger.info(f"Trying joblib.load on {path}")
+        m = joblib.load(path)
+        logger.info("joblib.load succeeded")
+        return ("joblib", m)
+    except Exception as e:
+        logger.exception(f"joblib.load failed: {e}")
+        return ("joblib", e)
+
+def try_xgboost_booster(path):
+    try:
+        import xgboost as xgb
+    except Exception as e:
+        logger.exception(f"xgboost import failed: {e}")
+        return ("xgboost_import", e)
+
+    try:
+        logger.info(f"Trying xgboost.Booster().load_model on {path}")
+        booster = xgb.Booster()
+        booster.load_model(path)
+        logger.info("xgboost.Booster.load_model succeeded")
+
+        class BoosterWrapper:
+            def __init__(self, booster):
+                self.booster = booster
+                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:
+                    return (_np.where(pred >= 0.5, 1, 0)).tolist()
+                return pred.tolist()
+
+            def predict_proba(self, X):
+                import numpy as _np, xgboost as _xgb
+                arr = _np.array(X, dtype=float)
+                dmat = _xgb.DMatrix(arr)
+                pred = self.booster.predict(dmat)
+                if pred.ndim == 1:
+                    return (_np.vstack([1 - pred, pred]).T).tolist()
+                return pred.tolist()
+
+        return ("xgboost_booster", BoosterWrapper(booster))
+    except Exception as e:
+        logger.exception(f"xgboost.Booster.load_model failed: {e}")
+        return ("xgboost_booster", e)
+
+# -------------------------
+# Core loader
+# -------------------------
 def load_model():
     logger.info("==== MODEL LOAD START ====")
     logger.info(f"Repo: {HF_MODEL_REPO}")
     logger.info(f"Filename: {HF_MODEL_FILENAME}")
     logger.info(f"HF_TOKEN present? {bool(HF_TOKEN)}")
 
-    # 1) Try local paths
+    # Try local candidates
     for path in LOCAL_CANDIDATES:
         try:
-            if os.path.exists(path):
-                logger.info(f"Attempting to load local model: {path}")
-                model = joblib.load(path)
-                logger.info("Model loaded successfully from local file.")
-                return model
+            info = inspect_file(path)
+            logger.info(f"Inspecting local candidate: {info}")
+            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
+
         except Exception as e:
-            logger.exception(f"Failed reading local model at {path}: {e}")
+            logger.exception(f"Error while trying local candidate {path}: {e}")
 
-    # 2) Try downloading from Hugging Face Hub
+    # 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
-        )
+        model_path = hf_hub_download(repo_id=HF_MODEL_REPO, filename=HF_MODEL_FILENAME, token=HF_TOKEN)
         logger.info(f"Downloaded model to: {model_path}")
-        model = joblib.load(model_path)
-        logger.info("Model loaded successfully from HF Hub.")
-        return model
-    except Exception as e:
-        logger.exception(f"hf_hub_download failed: {e}")
+        info = inspect_file(model_path)
+        logger.info(f"Inspecting downloaded file: {info}")
+
+        t, res = try_joblib_load(model_path)
+        if t == "joblib" and not isinstance(res, Exception):
+            return res
 
-    # 3) Fast fail — model not available
-    logger.error("❌ Model could NOT be loaded from local or HF Hub.")
-    return None
+        t, res = try_xgboost_booster(model_path)
+        if t == "xgboost_booster" and not isinstance(res, Exception):
+            return res
 
+        logger.error("Tried joblib and xgboost loader on downloaded file but both failed.")
+        return None
+    except Exception as e:
+        logger.exception(f"hf_hub_download failed: {e}")
+        return None
 
-# -----------------------------------------------------------
-# Prediction Function
-# -----------------------------------------------------------
-def predict(model, features: dict):
+# -------------------------
+# Prediction helper (robust)
+# Accepts: features as dict, list, or list-of-lists
+# Ensures sklearn pipelines that need DataFrame get a pandas.DataFrame
+# -------------------------
+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": "..."}
+    """
     if model is None:
         return {"error": "Model not loaded"}
 
     try:
-        df = features  # Streamlit form already creates proper dict
-        probabilities = model.predict_proba([list(df.values())])[0]
-        prediction = model.predict([list(df.values())])[0]
-
-        return {
-            "prediction": prediction,
-            "probability": probabilities.max()
-        }
+        # 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
+        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
+                if hasattr(model, "predict_proba"):
+                    p = model.predict_proba(arr)
+                    try:
+                        prob = float(p[0][1])
+                    except:
+                        prob = None
+                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
+                if hasattr(model, "predict_proba"):
+                    p = model.predict_proba(df)[0]
+                    try:
+                        prob = float(max(p))
+                    except:
+                        prob = None
+                # convert numpy types to native
+                try:
+                    pred = int(pred)
+                except:
+                    pass
+                return {"prediction": pred, "probability": prob}
+
+            return {"error": "Loaded model object not recognized (no predict method)"}
+
+        # Case B: features is list or tuple -> single row without column names
+        if isinstance(features, (list, tuple)):
+            # single-row list
+            arr2d = _np.array([features], dtype=float)
+            if is_booster:
+                preds = model.predict(arr2d)
+                prob = None
+                if hasattr(model, "predict_proba"):
+                    p = model.predict_proba(arr2d)
+                    try:
+                        prob = float(p[0][1])
+                    except:
+                        prob = None
+                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]
+                    prob = None
+                    if hasattr(model, "predict_proba"):
+                        p = model.predict_proba(arr2d)[0]
+                        try:
+                            prob = float(max(p))
+                        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
+                    cols = [str(i) for i in range(arr2d.shape[1])]
+                    df = _pd.DataFrame(arr2d, columns=cols)
+                    pred = model.predict(df)[0]
+                    prob = None
+                    if hasattr(model, "predict_proba"):
+                        p = model.predict_proba(df)[0]
+                        try:
+                            prob = float(max(p))
+                        except:
+                            prob = None
+                    return {"prediction": pred, "probability": prob}
+
+        # Case C: features is list-of-lists (batch)
+        if isinstance(features, list) and len(features) > 0 and isinstance(features[0], (list, tuple)):
+            arr = _np.array(features, dtype=float)
+            if is_booster:
+                preds = model.predict(arr)
+                prob = None
+                if hasattr(model, "predict_proba"):
+                    p = model.predict_proba(arr)
+                    try:
+                        prob = float(p[0][1])
+                    except:
+                        prob = None
+                return {"prediction": preds.tolist(), "probability": prob}
+            if hasattr(model, "predict"):
+                try:
+                    pred = model.predict(arr)
+                    prob = None
+                    if hasattr(model, "predict_proba"):
+                        p = model.predict_proba(arr)
+                        try:
+                            prob = float(max(p[0]))
+                        except:
+                            prob = None
+                    return {"prediction": pred.tolist(), "probability": prob}
+                except Exception as e:
+                    # try DataFrame fallback
+                    cols = [str(i) for i in range(arr.shape[1])]
+                    df = _pd.DataFrame(arr, columns=cols)
+                    pred = model.predict(df)
+                    prob = None
+                    if hasattr(model, "predict_proba"):
+                        p = model.predict_proba(df)
+                        try:
+                            prob = float(max(p[0]))
+                        except:
+                            prob = None
+                    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)}