Fix MLModelLoader to search notebooks models/ directory + trained_models/

app.py

@@ -247,8 +247,18 @@ class MLModelLoader:
 
     def initialize(self):
         loaded = 0
+        # All directories where models might exist (notebooks save to ../models)
+        search_dirs = [
+            MODELS_DIR,                          # trained_models/
+            APP_DIR / "models",                  # models/ (where notebooks output)
+            APP_DIR.parent / "models",           # one level up fallback
+        ]
+
         for name in self.MODEL_NAMES:
+            if name in self.models:
+                continue
             try:
+                # 1. Try HuggingFace Hub first
                 model_file = f"{name}/best_model.pkl"
                 scaler_file = f"{name}/scaler.pkl"
                 try:
@@ -264,43 +274,47 @@ class MLModelLoader:
                     self.scalers[name] = joblib.load(scaler_path)
                     loaded += 1
                     logger.info(f"✅ Loaded model from Hub: {name}")
+                    continue
                 except Exception:
-                    # Try flat filename pattern
+                    pass
+
+                # 2. Try all local search directories
+                for sdir in search_dirs:
+                    if name in self.models:
+                        break
+                    for model_fname in [f"{name}/best_model.pkl", f"{name}/model.pkl", f"{name}_model.pkl"]:
+                        candidate = sdir / model_fname
+                        if candidate.exists():
+                            self.models[name] = joblib.load(candidate)
+                            # Try to find matching scaler
+                            for scaler_fname in [f"{name}/scaler.pkl", f"{name}_scaler.pkl"]:
+                                sc = sdir / scaler_fname
+                                if sc.exists():
+                                    self.scalers[name] = joblib.load(sc)
+                                    break
+                            loaded += 1
+                            logger.info(f"✅ Loaded model from {sdir.name}/{model_fname}: {name}")
+                            break
+
+            except Exception as e:
+                logger.warning(f"Error loading model {name}: {e}")
+
+        # Sweep all search dirs for any .pkl files not yet loaded
+        for sdir in search_dirs:
+            if not sdir.exists():
+                continue
+            for pkl in sdir.glob("*.pkl"):
+                stem = pkl.stem.replace("_model", "").replace("_best", "")
+                if stem not in self.models:
                     try:
-                        model_path = hf_hub_download(
-                            repo_id=HF_MODEL_REPO, filename=f"{name}_model.pkl",
-                            token=HF_TOKEN or None, cache_dir=str(MODELS_DIR),
-                        )
-                        self.models[name] = joblib.load(model_path)
+                        self.models[stem] = joblib.load(pkl)
                         loaded += 1
-                        logger.info(f"✅ Loaded model (flat): {name}")
+                        logger.info(f"✅ Loaded model sweep: {stem} from {sdir.name}")
                     except Exception:
                         pass
-                    # Try local
-                    local_model = MODELS_DIR / name / "best_model.pkl"
-                    if local_model.exists() and name not in self.models:
-                        self.models[name] = joblib.load(local_model)
-                        local_scaler = MODELS_DIR / name / "scaler.pkl"
-                        if local_scaler.exists():
-                            self.scalers[name] = joblib.load(local_scaler)
-                        loaded += 1
-                        logger.info(f"✅ Loaded model from local: {name}")
-            except Exception as e:
-                logger.warning(f"Error loading model {name}: {e}")
-
-        # Also check trained_models dir for any .pkl files
-        for pkl in MODELS_DIR.glob("*.pkl"):
-            stem = pkl.stem.replace("_model", "").replace("_best", "")
-            if stem not in self.models:
-                try:
-                    self.models[stem] = joblib.load(pkl)
-                    loaded += 1
-                    logger.info(f"✅ Loaded local model: {stem}")
-                except Exception:
-                    pass
 
         self.ready = loaded > 0
-        logger.info(f"ML Models: {loaded} loaded ({list(ml_loader.models.keys()) if loaded else 'none'})")
+        logger.info(f"ML Models: {loaded} loaded — {list(self.models.keys())}")
 
     def predict(self, model_name: str, features: Dict) -> Dict:
         if model_name not in self.models: