semihyilmaz37/PLNI

1import joblib2import torch3import torch.nn as nn4import numpy as np5from scipy.interpolate import interp1d6from pathlib import Path7from typing import Any, Dict, List, Optional8 9 10BASE_DIR = Path(__file__).resolve().parent11MODELS_DIR = BASE_DIR / "saved_models"12 13FEATURE_NAMES_PATH = MODELS_DIR / "feature_names.joblib"14MODELS_SUMMARY_PATH = MODELS_DIR / "models_summary.joblib"15NN_CALIBRATION_PATH = MODELS_DIR / "neural_network_calibration.joblib"16 17# Map human‑readable model names to their corresponding files18MODEL_PATHS = {19    "Logistic Regression": MODELS_DIR / "logisticregression_model.joblib",20    "Support Vector Machine": MODELS_DIR / "svm_model.joblib",21    "Neural Network": MODELS_DIR / "neural_network_pytorch.pt",  # PyTorch model22    "XGBoost": MODELS_DIR / "xgboost_model.joblib",23}24 25_feature_names_cache: Optional[List[str]] = None26_models_cache: Dict[str, Any] = {}27_models_summary_cache: Optional[Dict[str, Any]] = None28_nn_calibration_cache: Optional[Dict[str, Any]] = None29 30 31class NeuralNetwork(nn.Module):32    """33    PyTorch Neural Network model matching the architecture used during training.34    """35    def __init__(self, input_size: int, hidden_sizes: List[int]):36        super(NeuralNetwork, self).__init__()37        layers = []38        prev_size = input_size39        for hidden_size in hidden_sizes:40            layers.append(nn.Linear(prev_size, hidden_size))41            layers.append(nn.ReLU())42            layers.append(nn.Dropout(0.2))43            prev_size = hidden_size44        layers.append(nn.Linear(prev_size, 1))45        layers.append(nn.Sigmoid())46        self.model = nn.Sequential(*layers)47    48    def forward(self, x):49        return self.model(x)50 51 52def get_feature_names() -> List[str]:53    """54    Load and cache the ordered list of feature names used to train the models.55    """56    global _feature_names_cache57 58    if _feature_names_cache is None:59        try:60            _feature_names_cache = joblib.load(FEATURE_NAMES_PATH)61        except Exception as e:62            raise FileNotFoundError(63                f"Could not load feature names from {FEATURE_NAMES_PATH}: {e}"64            )65 66    return _feature_names_cache67 68 69def get_available_models() -> List[str]:70    """71    Return a list of human‑readable model names that can be used with `predict`.72    """73    return list(MODEL_PATHS.keys())74 75 76def _load_model(model_name: str) -> Any:77    """78    Internal helper that loads (and caches) a model by its human‑readable name.79    Handles PyTorch Neural Network models specially.80    """81    if model_name not in MODEL_PATHS:82        raise ValueError(f"Unknown model '{model_name}'. Available: {get_available_models()}")83 84    if model_name not in _models_cache:85        model_path = MODEL_PATHS[model_name]86        87        if model_name == "Neural Network":88            # Load PyTorch model only (no joblib fallback to avoid NeuralNetworkWrapper issues)89            if not model_path.exists():90                raise FileNotFoundError(91                    f"Neural Network PyTorch model not found at {model_path}. "92                    f"Please ensure the model was saved as a .pt file."93                )94            95            checkpoint = torch.load(model_path, map_location='cpu')96            input_size = checkpoint['input_size']97            hidden_sizes = checkpoint['hidden_sizes']98            99            model = NeuralNetwork(input_size, hidden_sizes)100            model.load_state_dict(checkpoint['model_state_dict'])101            model.eval()102            _models_cache[model_name] = model103 104            # Load calibration data if available105            global _nn_calibration_cache106            if _nn_calibration_cache is None and NN_CALIBRATION_PATH.exists():107                try:108                    _nn_calibration_cache = joblib.load(NN_CALIBRATION_PATH)109                except Exception:110                    _nn_calibration_cache = None111        else:112            # Load sklearn/XGBoost models with joblib113            _models_cache[model_name] = joblib.load(model_path)114 115    return _models_cache[model_name]116 117 118def get_models_summary() -> Dict[str, Any]:119    """120    Return the summary/metadata for all models (e.g. metrics), if available.121    Handles cases where the summary file contains unpicklable classes gracefully.122    """123    global _models_summary_cache124 125    if _models_summary_cache is None:126        try:127            _models_summary_cache = joblib.load(MODELS_SUMMARY_PATH)128        except (AttributeError, ModuleNotFoundError, ImportError) as e:129            # Handle cases where joblib can't unpickle classes (e.g., NeuralNetworkWrapper)130            # Return a basic summary instead131            print(f"Warning: Could not load full models summary: {e}")132            _models_summary_cache = {133                "note": "Full summary unavailable due to missing class definitions",134                "available_models": get_available_models(),135                "feature_names": get_feature_names(),136            }137 138    return _models_summary_cache139 140 141def predict(142    model_name: str,143    features: Dict[str, Any],144    debug: bool = False,145) -> Dict[str, Any]:146    """147    Run a prediction using one of the saved models.148 149    Parameters150    ----------151    model_name:152        Human‑readable model name (see `get_available_models()`).153    features:154        Dict mapping feature name -> value. Any missing features will be filled155        with 0.0; any extra keys are ignored.156    debug:157        If True, print debug information about feature mapping.158 159    Returns160    -------161    Dict with at least:162        - 'prediction': the model's predicted class/value163      and, when available:164        - 'probabilities': list of class probabilities from `predict_proba`165    """166    feature_names = get_feature_names()167 168    # Build a single row in the correct order expected by the models169    ordered_row = [[features.get(name, 0.0) for name in feature_names]]170    171    if debug:172        print(f"\nDEBUG [{model_name}]: Feature vector being sent to model:")173        for i, (name, value) in enumerate(zip(feature_names, ordered_row[0])):174            marker = " <-- SET" if name in features else " <-- DEFAULT (0.0)"175            print(f"  [{i}] {name}: {value}{marker}")176 177    model = _load_model(model_name)178 179    if model_name == "Neural Network":180        # Handle PyTorch model (apply isotonic calibration if available)181        X_tensor = torch.FloatTensor(ordered_row)182        with torch.no_grad():183            y_prob_raw = model(X_tensor).numpy().flatten()184 185        y_prob = _apply_isotonic_calibration(y_prob_raw)186        y_pred = (y_prob > 0.5).astype(int)187 188        result: Dict[str, Any] = {189            "prediction": int(y_pred[0]),190            "probabilities": [float(1 - y_prob[0]), float(y_prob[0])]  # [class 0, class 1]191        }192    else:193        # Handle sklearn/XGBoost models194        y_pred = model.predict(ordered_row)195        result: Dict[str, Any] = {"prediction": int(y_pred[0])}196 197        # Add probabilities if the model supports them198        if hasattr(model, "predict_proba"):199            proba = model.predict_proba(ordered_row)[0]200            # Convert to plain Python list to make it JSON/Gradio‑friendly201            result["probabilities"] = proba.tolist()202 203    return result204 205 206def _apply_isotonic_calibration(y_prob_raw: np.ndarray) -> np.ndarray:207    """208    Apply isotonic calibration (if saved) to raw NN probabilities.209    Falls back to raw probabilities if calibration data is missing or invalid.210    """211    if _nn_calibration_cache is None:212        return y_prob_raw213 214    calibrators_data = _nn_calibration_cache.get("calibrators_data")215    if not calibrators_data:216        return y_prob_raw217 218    calibrated_probs = []219    for cal_data in calibrators_data:220        x_thresh = cal_data.get("X_thresholds")221        y_thresh = cal_data.get("y_thresholds")222        if x_thresh is None or y_thresh is None:223            continue224 225        try:226            f = interp1d(x_thresh, y_thresh, bounds_error=False, fill_value=(0, 1))227            calibrated_probs.append(f(y_prob_raw))228        except Exception:229            continue230 231    if not calibrated_probs:232        return y_prob_raw233 234    return np.mean(calibrated_probs, axis=0)235 236 237def predict_all_models(238    features: Dict[str, Any],239    debug: bool = False,240) -> Dict[str, Dict[str, Any]]:241    """242    Run predictions using all available models and return results together.243 244    Parameters245    ----------246    features:247        Dict mapping feature name -> value. Any missing features will be filled248        with 0.0; any extra keys are ignored.249    debug:250        If True, print debug information about feature mapping.251 252    Returns253    -------254    Dict mapping model name -> prediction result (same format as `predict`)255    """256    results = {}257    for model_name in get_available_models():258        try:259            results[model_name] = predict(model_name, features, debug=debug)260        except Exception as e:261            results[model_name] = {262                "prediction": None,263                "error": str(e)264            }265    return results266 267 268__all__ = [269    "get_feature_names",270    "get_available_models",271    "get_models_summary",272    "predict",273    "predict_all_models",274]275 276 277