from __future__ import annotations import json import os import threading from typing import Dict, Tuple, Optional, List from pathlib import Path import joblib import numpy as np import pandas as pd from huggingface_hub import hf_hub_download, snapshot_download class _Singleton(type): """Thread-safe Singleton metaclass (una instancia por proceso).""" _instances: Dict[type, object] = {} _lock = threading.Lock() def __call__(cls, *args, **kwargs): # Double-checked locking if cls not in cls._instances: with cls._lock: if cls not in cls._instances: cls._instances[cls] = super().__call__(*args, **kwargs) return cls._instances[cls] class ExoMACModel(metaclass=_Singleton): """ Misión-agnóstico: cargador de modelo (Pipeline sklearn) entrenado con Kepler/K2/TESS. - Descarga artefactos desde Hugging Face SOLO si no existen localmente. - Guarda/lee desde una carpeta local del proyecto (por defecto: ./models/ExoMAC-KKT). - Exposición de helpers de predicción y de features ingenierizadas. """ DEFAULT_REPO = "ZapatoProgramming/ExoMAC-KKT" _FILENAMES = { "model": "exoplanet_best_model.joblib", "feats": "exoplanet_feature_columns.json", "labels": "exoplanet_class_labels.json", "meta": "exoplanet_metadata.json", } def __init__( self, repo_id: Optional[str] = None, token: Optional[str] = None, prefer_snapshot: bool = True, allow_patterns: Optional[List[str]] = None, local_dir: Optional[str | os.PathLike] = None, always_download: bool = False, verbose: bool = True, ): """ Args: repo_id: Hugging Face repo id. Por defecto 'ZapatoProgramming/ExoMAC-KKT'. token: Token HF si el repo es privado. prefer_snapshot: Si True, usa snapshot_download (descarga por patrón). allow_patterns: Patrones a descargar cuando prefer_snapshot=True. local_dir: Carpeta donde se guardan/leen artefactos en tu proyecto. always_download: Si True, fuerza descarga (útil para actualizar). verbose: Imprime mensajes útiles. """ self.repo_id = repo_id or self.DEFAULT_REPO self.token = token self.prefer_snapshot = prefer_snapshot self.allow_patterns = allow_patterns or ["artifacts/*", "*.joblib", "*.json"] self.local_dir = Path(local_dir or (Path("models") / self.repo_id.split("/")[-1])) self.local_dir.mkdir(parents=True, exist_ok=True) self.always_download = always_download self.verbose = verbose self._model = None self._feature_columns: List[str] = [] self._class_labels: List[str] = [] self._metadata: Dict = {} self._load_artifacts() # ------------------------- PUBLIC API ------------------------- @property def model(self): return self._model @property def feature_columns(self) -> List[str]: return list(self._feature_columns) @property def class_labels(self) -> List[str]: return list(self._class_labels) @property def metadata(self) -> Dict: return dict(self._metadata) def predict( self, params: Dict[str, float], return_proba: bool = True, compute_engineered_if_missing: bool = True, ) -> Tuple[str, Optional[Dict[str, float]]]: """ Predice una etiqueta y (opcionalmente) probabilidades para un dict de features. - Rellena features ingenierizadas si el modelo las espera y no están. """ if compute_engineered_if_missing: params = self._ensure_engineered_features(dict(params)) X = pd.DataFrame([params], dtype=float).reindex(columns=self._feature_columns) y_idx = int(self._model.predict(X)[0]) label = self._class_labels[y_idx] if not return_proba: return label, None proba = None try: p = self._model.predict_proba(X)[0] proba = {lbl: float(prob) for lbl, prob in zip(self._class_labels, p)} except Exception: pass return label, proba def predict_with_debug(self, params: Dict[str, float]) -> Tuple[str, Optional[Dict[str, float]]]: """ Igual que predict(), pero imprime features reconocidas/desconocidas y faltantes. """ params2 = self._ensure_engineered_features(dict(params)) X = pd.DataFrame([params2], dtype=float).reindex(columns=self._feature_columns) recognized = [c for c in self._feature_columns if c in params2] unknown = [k for k in params2.keys() if k not in self._feature_columns] missing = X.columns[X.iloc[0].isna()].tolist() print(f"Recognized: {len(recognized)}/{len(self._feature_columns)}") if recognized: print(" •", ", ".join(recognized[:16]) + (" ..." if len(recognized) > 16 else "")) if unknown: print(f"Unknown keys: {len(unknown)}") if unknown: print(" •", ", ".join(unknown[:16]) + (" ..." if len(unknown) > 16 else "")) if missing: print(f"Missing (imputed): {len(missing)}") if missing: print(" •", ", ".join(missing[:16]) + (" ..." if len(missing) > 16 else "")) return self.predict(params2, return_proba=True, compute_engineered_if_missing=False) # ------------------------- INTERNALS ------------------------- def _load_artifacts(self) -> None: """ 1) Si ya existen archivos locales y always_download=False -> NO descarga. 2) Si faltan archivos o always_download=True -> descarga (snapshot o per-file). 3) Carga el modelo + metadata desde disco. """ paths: Optional[Dict[str, str]] = None # (0) Intentar leer desde local sin tocar red if not self.always_download: local_paths = self._try_local_paths() if local_paths is not None: paths = local_paths if self.verbose: print(f"[ExoMAC] Using cached artifacts in {self.local_dir}") else: if self.verbose: print(f"[ExoMAC] Local artifacts not found. Will download to {self.local_dir}.") # (1) Descargar si hace falta if paths is None: if self.prefer_snapshot: # Descarga patrones a la carpeta local (la API ya no usa symlinks) snapshot_download( repo_id=self.repo_id, token=self.token, allow_patterns=self.allow_patterns, local_dir=str(self.local_dir), ) paths = self._resolve_from_dir(self.local_dir) else: paths = {} for key, fname in self._FILENAMES.items(): paths[key] = self._get_artifact_to_local_dir(fname) # (2) Cargar desde disco self._model = joblib.load(paths["model"]) self._feature_columns = json.load(open(paths["feats"], "r", encoding="utf-8")) self._class_labels = json.load(open(paths["labels"], "r", encoding="utf-8")) self._metadata = json.load(open(paths["meta"], "r", encoding="utf-8")) if self.verbose: print(f"[ExoMAC] Loaded model from {paths['model']}") # --- Local path helpers --- def _have_all_files(self, base: Path) -> bool: """¿Están TODOS los artefactos (en artifacts/ o raíz) en 'base'?""" base = Path(base) for _, name in self._FILENAMES.items(): p1 = base / "artifacts" / name p2 = base / name if not (p1.exists() or p2.exists()): return False return True def _try_local_paths(self) -> Optional[Dict[str, str]]: """Devuelve rutas locales si todo existe; si falta algo, None.""" if self._have_all_files(self.local_dir): return self._resolve_from_dir(self.local_dir) return None def _resolve_from_dir(self, base_dir: Path | str) -> Dict[str, str]: """ Selecciona artifacts/ si existe; si no, /. """ base_dir = Path(base_dir) out: Dict[str, str] = {} for key, name in self._FILENAMES.items(): p1 = base_dir / "artifacts" / name p2 = base_dir / name if p1.exists(): out[key] = str(p1) elif p2.exists(): out[key] = str(p2) else: raise FileNotFoundError(f"Could not find {name} under {base_dir}") return out def _get_artifact_to_local_dir(self, fname: str) -> str: """ Descarga a self.local_dir con hf_hub_download (si tu versión soporta local_dir). Si no, descarga a la caché global y copia a self.local_dir. """ self.local_dir.mkdir(parents=True, exist_ok=True) for candidate in (f"artifacts/{fname}", fname): try: # huggingface_hub >= 0.23 soporta local_dir path = hf_hub_download( repo_id=self.repo_id, filename=candidate, token=self.token, local_dir=str(self.local_dir), ) return path except TypeError: # Fallback: versión antigua sin local_dir cache_path = hf_hub_download( repo_id=self.repo_id, filename=candidate, token=self.token, ) dst = self.local_dir / Path(candidate).name os.makedirs(self.local_dir, exist_ok=True) if not os.path.exists(dst): from shutil import copy2 copy2(cache_path, dst) return str(dst) except Exception: # prueba siguiente candidato (raíz en lugar de artifacts/) continue raise FileNotFoundError(f"Could not download {fname} from {self.repo_id}") # --- Engineered features helpers --- def _ensure_engineered_features(self, d: Dict[str, float]) -> Dict[str, float]: """ Rellena features ingenierizadas si el modelo las espera y no están: - duty_cycle, log_koi_period, log_koi_depth, teq_proxy - koi_snr/log_koi_snr o snr_proxy/log_snr_proxy (proxy) """ need = set(self._feature_columns) # Duty cycle if "duty_cycle" in need and "duty_cycle" not in d: if all(k in d for k in ("koi_duration", "koi_period")) and d.get("koi_period"): d["duty_cycle"] = d["koi_duration"] / (d["koi_period"] * 24.0) # Logs if "log_koi_period" in need and "log_koi_period" not in d and d.get("koi_period", 0) > 0: d["log_koi_period"] = np.log10(d["koi_period"]) if "log_koi_depth" in need and "log_koi_depth" not in d and d.get("koi_depth", 0) > 0: d["log_koi_depth"] = np.log10(d["koi_depth"]) # teq_proxy (simple) if "teq_proxy" in need and "teq_proxy" not in d and "koi_steff" in d: d["teq_proxy"] = d["koi_steff"] # SNR real o proxy if "koi_snr" in need and "koi_snr" not in d: d["koi_snr"] = np.nan if "log_koi_snr" in need and "log_koi_snr" not in d and d.get("koi_snr", 0) > 0: d["log_koi_snr"] = np.log10(d["koi_snr"]) if "snr_proxy" in need and "snr_proxy" not in d: if all(k in d for k in ("koi_depth", "koi_duration", "koi_period")) and d.get("koi_period", 0) > 0: d["snr_proxy"] = d["koi_depth"] * np.sqrt(max(d["koi_duration"] / (d["koi_period"] * 24.0), 1e-12)) if "log_snr_proxy" in need and "log_snr_proxy" not in d and d.get("snr_proxy", 0) > 0: d["log_snr_proxy"] = np.log10(d["snr_proxy"]) return d # ------------------------- DEMO ------------------------- if __name__ == "__main__": # Primera ejecución: descargará a ./models/ExoMAC-KKT si no existe. model = ExoMACModel( local_dir="./ExoMACModel/ExoMAC-KKT", prefer_snapshot=True, always_download=False, # <- ejecuciones siguientes NO vuelven a descargar verbose=True, ) # Subsecuentes: misma instancia (singleton) y SIN descarga. same_model = ExoMACModel(local_dir="./ExoMACModel/ExoMAC-KKT") assert model is same_model # Ejemplo mínimo de predicción params = { "koi_period": 12.0, "koi_duration": 3.5, "koi_depth": 600.0, "koi_impact": 0.20, "koi_prad": 2.1, "koi_slogg": 4.4, "koi_sma": 0.10, "koi_smet": 0.0, "koi_srad": 1.0, "koi_steff": 5700.0, "koi_snr": 12.0, } label, proba = model.predict_with_debug(params) print("Predicted:", label) print("Local dir:", model.local_dir.resolve())