Spaces:
Sleeping
Sleeping
| import os | |
| import joblib | |
| import threading | |
| import traceback | |
| from huggingface_hub import hf_hub_download | |
| # --- CONFIGURATION --- | |
| HF_REPO_ID = "CodebaseAi/netraids-ml-models" # Replace with your actual public repo ID | |
| # --------------------- | |
| ACTIVE_MODEL = "bcc" | |
| _ACTIVE_LOCK = threading.Lock() | |
| _MODEL_CACHE = {} | |
| # 1. FIXED PATH LOGIC: | |
| # __file__ is /app/utils/model_selector.py | |
| # dirname(__file__) is /app/utils | |
| # dirname(dirname(...)) is /app (the ROOT) | |
| BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| ML_DIR = os.path.join(BASE_DIR, "ml_models") | |
| # Ensure the local ml_models directory exists for caching | |
| if not os.path.exists(ML_DIR): | |
| os.makedirs(ML_DIR, exist_ok=True) | |
| print(f"[model_selector] ROOT BASE_DIR: {BASE_DIR}") | |
| print(f"[model_selector] ML_DIR: {ML_DIR}") | |
| def _get_model_path(filename): | |
| """ | |
| First looks in the local 'ml_models' folder. | |
| If not found, downloads from the public Hugging Face Hub. | |
| """ | |
| local_path = os.path.join(ML_DIR, filename) | |
| # 1. Check if the file is already there | |
| if os.path.exists(local_path): | |
| return local_path | |
| # 2. Download from Hub if missing | |
| try: | |
| print(f"[model_selector] {filename} not found locally. Downloading from Hub...") | |
| # We specify local_dir to force it into our ml_models folder | |
| downloaded_path = hf_hub_download( | |
| repo_id=HF_REPO_ID, | |
| filename=filename, | |
| local_dir=ML_DIR | |
| ) | |
| return downloaded_path | |
| except Exception as e: | |
| print(f"[model_selector] ERROR: Could not find/download {filename}: {e}") | |
| return None | |
| def _try_load(filename): | |
| path = _get_model_path(filename) | |
| if not path or not os.path.exists(path): | |
| print(f"[model_selector] SKIP: {filename} path invalid.") | |
| return None | |
| try: | |
| return joblib.load(path) | |
| except Exception as e: | |
| print(f"[model_selector] FAILED to load {filename}: {e}") | |
| return None | |
| def load_model(model_key): | |
| if model_key in _MODEL_CACHE: | |
| return _MODEL_CACHE[model_key] | |
| if model_key == "bcc": | |
| _MODEL_CACHE["bcc"] = { | |
| "model": _try_load("realtime_model.pkl"), | |
| "scaler": _try_load("realtime_scaler.pkl"), | |
| "encoder": _try_load("realtime_encoder.pkl") | |
| } | |
| return _MODEL_CACHE["bcc"] | |
| if model_key == "cicids": | |
| # It will look for your RF files in the Hub | |
| _MODEL_CACHE["cicids"] = { | |
| "model": _try_load("rf_pipeline.joblib"), | |
| "artifacts": _try_load("training_artifacts.joblib") | |
| } | |
| return _MODEL_CACHE["cicids"] | |
| raise ValueError(f"Unknown model_key: {model_key}") | |
| def set_active_model(key: str): | |
| global ACTIVE_MODEL | |
| with _ACTIVE_LOCK: | |
| ACTIVE_MODEL = key | |
| print(f"[model_selector] ACTIVE_MODEL set to: {ACTIVE_MODEL}") | |
| def get_active_model(): | |
| return ACTIVE_MODEL | |