dragon-ml-toolbox 7.0.0__py3-none-any.whl → 8.1.0__py3-none-any.whl

ml_tools/ML_evaluation_multi.py

@@ -0,0 +1,296 @@
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import seaborn as sns
+import torch
+import shap
+from sklearn.metrics import (
+    classification_report,
+    ConfusionMatrixDisplay,
+    roc_curve,
+    roc_auc_score,
+    precision_recall_curve,
+    average_precision_score,
+    mean_squared_error,
+    mean_absolute_error,
+    r2_score,
+    median_absolute_error,
+    hamming_loss,
+    jaccard_score
+)
+from pathlib import Path
+from typing import Union, List, Optional
+
+from .path_manager import make_fullpath, sanitize_filename
+from ._logger import _LOGGER
+from ._script_info import _script_info
+
+__all__ = [
+    "multi_target_regression_metrics",
+    "multi_label_classification_metrics",
+    "multi_target_shap_summary_plot",
+]
+
+
+def multi_target_regression_metrics(
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    target_names: List[str],
+    save_dir: Union[str, Path]
+):
+    """
+    Calculates and saves regression metrics for each target individually.
+
+    For each target, this function saves a residual plot and a true vs. predicted plot.
+    It also saves a single CSV file containing the key metrics (RMSE, MAE, R², MedAE)
+    for all targets.
+
+    Args:
+        y_true (np.ndarray): Ground truth values, shape (n_samples, n_targets).
+        y_pred (np.ndarray): Predicted values, shape (n_samples, n_targets).
+        target_names (List[str]): A list of names for the target variables.
+        save_dir (str | Path): Directory to save plots and the report.
+    """
+    if y_true.ndim != 2 or y_pred.ndim != 2:
+        raise ValueError("y_true and y_pred must be 2D arrays for multi-target regression.")
+    if y_true.shape != y_pred.shape:
+        raise ValueError("Shapes of y_true and y_pred must match.")
+    if y_true.shape[1] != len(target_names):
+        raise ValueError("Number of target names must match the number of columns in y_true.")
+
+    save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
+    metrics_summary = []
+
+    _LOGGER.info("--- Multi-Target Regression Evaluation ---")
+
+    for i, name in enumerate(target_names):
+        _LOGGER.info(f"  -> Evaluating target: '{name}'")
+        true_i = y_true[:, i]
+        pred_i = y_pred[:, i]
+        sanitized_name = sanitize_filename(name)
+
+        # --- Calculate Metrics ---
+        rmse = np.sqrt(mean_squared_error(true_i, pred_i))
+        mae = mean_absolute_error(true_i, pred_i)
+        r2 = r2_score(true_i, pred_i)
+        medae = median_absolute_error(true_i, pred_i)
+        metrics_summary.append({
+            'target': name,
+            'rmse': rmse,
+            'mae': mae,
+            'r2_score': r2,
+            'median_abs_error': medae
+        })
+
+        # --- Save Residual Plot ---
+        residuals = true_i - pred_i
+        fig_res, ax_res = plt.subplots(figsize=(8, 6), dpi=100)
+        ax_res.scatter(pred_i, residuals, alpha=0.6, edgecolors='k', s=50)
+        ax_res.axhline(0, color='red', linestyle='--')
+        ax_res.set_xlabel("Predicted Values")
+        ax_res.set_ylabel("Residuals (True - Predicted)")
+        ax_res.set_title(f"Residual Plot for '{name}'")
+        ax_res.grid(True, linestyle='--', alpha=0.6)
+        plt.tight_layout()
+        res_path = save_dir_path / f"residual_plot_{sanitized_name}.svg"
+        plt.savefig(res_path)
+        plt.close(fig_res)
+
+        # --- Save True vs. Predicted Plot ---
+        fig_tvp, ax_tvp = plt.subplots(figsize=(8, 6), dpi=100)
+        ax_tvp.scatter(true_i, pred_i, alpha=0.6, edgecolors='k', s=50)
+        ax_tvp.plot([true_i.min(), true_i.max()], [true_i.min(), true_i.max()], 'k--', lw=2)
+        ax_tvp.set_xlabel('True Values')
+        ax_tvp.set_ylabel('Predicted Values')
+        ax_tvp.set_title(f'True vs. Predicted Values for "{name}"')
+        ax_tvp.grid(True, linestyle='--', alpha=0.6)
+        plt.tight_layout()
+        tvp_path = save_dir_path / f"true_vs_predicted_plot_{sanitized_name}.svg"
+        plt.savefig(tvp_path)
+        plt.close(fig_tvp)
+
+    # --- Save Summary Report ---
+    summary_df = pd.DataFrame(metrics_summary)
+    report_path = save_dir_path / "regression_report_multi.csv"
+    summary_df.to_csv(report_path, index=False)
+    _LOGGER.info(f"✅ Full regression report saved to '{report_path.name}'")
+
+
+def multi_label_classification_metrics(
+    y_true: np.ndarray,
+    y_prob: np.ndarray,
+    target_names: List[str],
+    save_dir: Union[str, Path],
+    threshold: float = 0.5
+):
+    """
+    Calculates and saves classification metrics for each label individually.
+
+    This function first computes overall multi-label metrics (Hamming Loss, Jaccard Score)
+    and then iterates through each label to generate and save individual reports,
+    confusion matrices, ROC curves, and Precision-Recall curves.
+
+    Args:
+        y_true (np.ndarray): Ground truth binary labels, shape (n_samples, n_labels).
+        y_prob (np.ndarray): Predicted probabilities, shape (n_samples, n_labels).
+        target_names (List[str]): A list of names for the labels.
+        save_dir (str | Path): Directory to save plots and reports.
+        threshold (float): The probability threshold to convert probabilities into
+                           binary predictions for metrics like the confusion matrix.
+    """
+    if y_true.ndim != 2 or y_prob.ndim != 2:
+        raise ValueError("y_true and y_prob must be 2D arrays for multi-label classification.")
+    if y_true.shape != y_prob.shape:
+        raise ValueError("Shapes of y_true and y_prob must match.")
+    if y_true.shape[1] != len(target_names):
+        raise ValueError("Number of target names must match the number of columns in y_true.")
+
+    save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
+    
+    # Generate binary predictions from probabilities
+    y_pred = (y_prob >= threshold).astype(int)
+
+    _LOGGER.info("--- Multi-Label Classification Evaluation ---")
+
+    # --- Calculate and Save Overall Metrics ---
+    h_loss = hamming_loss(y_true, y_pred)
+    j_score_micro = jaccard_score(y_true, y_pred, average='micro')
+    j_score_macro = jaccard_score(y_true, y_pred, average='macro')
+
+    overall_report = (
+        f"Overall Multi-Label Metrics (Threshold = {threshold}):\n"
+        f"--------------------------------------------------\n"
+        f"Hamming Loss: {h_loss:.4f}\n"
+        f"Jaccard Score (micro): {j_score_micro:.4f}\n"
+        f"Jaccard Score (macro): {j_score_macro:.4f}\n"
+        f"--------------------------------------------------\n"
+    )
+    _LOGGER.info(overall_report)
+    overall_report_path = save_dir_path / "classification_report_overall.txt"
+    overall_report_path.write_text(overall_report)
+
+    # --- Per-Label Metrics and Plots ---
+    for i, name in enumerate(target_names):
+        _LOGGER.info(f"  -> Evaluating label: '{name}'")
+        true_i = y_true[:, i]
+        pred_i = y_pred[:, i]
+        prob_i = y_prob[:, i]
+        sanitized_name = sanitize_filename(name)
+
+        # --- Save Classification Report for the label ---
+        report_text = classification_report(true_i, pred_i)
+        report_path = save_dir_path / f"classification_report_{sanitized_name}.txt"
+        report_path.write_text(report_text) # type: ignore
+
+        # --- Save Confusion Matrix ---
+        fig_cm, ax_cm = plt.subplots(figsize=(6, 6), dpi=100)
+        ConfusionMatrixDisplay.from_predictions(true_i, pred_i, cmap="Blues", ax=ax_cm)
+        ax_cm.set_title(f"Confusion Matrix for '{name}'")
+        cm_path = save_dir_path / f"confusion_matrix_{sanitized_name}.svg"
+        plt.savefig(cm_path)
+        plt.close(fig_cm)
+
+        # --- Save ROC Curve ---
+        fpr, tpr, _ = roc_curve(true_i, prob_i)
+        auc = roc_auc_score(true_i, prob_i)
+        fig_roc, ax_roc = plt.subplots(figsize=(6, 6), dpi=100)
+        ax_roc.plot(fpr, tpr, label=f'AUC = {auc:.2f}')
+        ax_roc.plot([0, 1], [0, 1], 'k--')
+        ax_roc.set_title(f'ROC Curve for "{name}"')
+        ax_roc.set_xlabel('False Positive Rate'); ax_roc.set_ylabel('True Positive Rate')
+        ax_roc.legend(loc='lower right'); ax_roc.grid(True, linestyle='--', alpha=0.6)
+        roc_path = save_dir_path / f"roc_curve_{sanitized_name}.svg"
+        plt.savefig(roc_path)
+        plt.close(fig_roc)
+
+        # --- Save Precision-Recall Curve ---
+        precision, recall, _ = precision_recall_curve(true_i, prob_i)
+        ap_score = average_precision_score(true_i, prob_i)
+        fig_pr, ax_pr = plt.subplots(figsize=(6, 6), dpi=100)
+        ax_pr.plot(recall, precision, label=f'AP = {ap_score:.2f}')
+        ax_pr.set_title(f'Precision-Recall Curve for "{name}"')
+        ax_pr.set_xlabel('Recall'); ax_pr.set_ylabel('Precision')
+        ax_pr.legend(loc='lower left'); ax_pr.grid(True, linestyle='--', alpha=0.6)
+        pr_path = save_dir_path / f"pr_curve_{sanitized_name}.svg"
+        plt.savefig(pr_path)
+        plt.close(fig_pr)
+
+    _LOGGER.info(f"✅ All individual label reports and plots saved to '{save_dir_path.name}'")
+
+
+def multi_target_shap_summary_plot(
+    model: torch.nn.Module,
+    background_data: Union[torch.Tensor, np.ndarray],
+    instances_to_explain: Union[torch.Tensor, np.ndarray],
+    feature_names: List[str],
+    target_names: List[str],
+    save_dir: Union[str, Path]
+):
+    """
+    Calculates SHAP values for a multi-target model and saves summary plots for each target.
+
+    Args:
+        model (torch.nn.Module): The trained PyTorch model.
+        background_data (torch.Tensor | np.ndarray): A sample of data for the explainer background.
+        instances_to_explain (torch.Tensor | np.ndarray): The specific data instances to explain.
+        feature_names (List[str]): Names of the features for plot labeling.
+        target_names (List[str]): Names of the output targets.
+        save_dir (str | Path): Directory to save SHAP artifacts.
+    """
+    # Convert all data to numpy
+    background_data_np = background_data.numpy() if isinstance(background_data, torch.Tensor) else background_data
+    instances_to_explain_np = instances_to_explain.numpy() if isinstance(instances_to_explain, torch.Tensor) else instances_to_explain
+
+    if np.isnan(background_data_np).any() or np.isnan(instances_to_explain_np).any():
+        _LOGGER.error("❌ Input data for SHAP contains NaN values. Aborting explanation.")
+        return
+
+    _LOGGER.info("\n--- Multi-Target SHAP Value Explanation ---")
+    model.eval()
+    model.cpu()
+
+    # 1. Summarize the background data.
+    background_summary = shap.kmeans(background_data_np, 30)
+
+    # 2. Define a prediction function wrapper for the multi-target model.
+    def prediction_wrapper(x_np: np.ndarray) -> np.ndarray:
+        x_torch = torch.from_numpy(x_np).float()
+        with torch.no_grad():
+            output = model(x_torch)
+        return output.cpu().numpy()
+
+    # 3. Create the KernelExplainer.
+    explainer = shap.KernelExplainer(prediction_wrapper, background_summary)
+
+    _LOGGER.info("Calculating SHAP values with KernelExplainer...")
+    # For multi-output models, shap_values is a list of arrays.
+    shap_values_list = explainer.shap_values(instances_to_explain_np, l1_reg="aic")
+
+    save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
+    plt.ioff()
+
+    # 4. Iterate through each target's SHAP values and generate plots.
+    for i, target_name in enumerate(target_names):
+        _LOGGER.info(f"  -> Generating SHAP plots for target: '{target_name}'")
+        shap_values_for_target = shap_values_list[i]
+        sanitized_target_name = sanitize_filename(target_name)
+
+        # Save Bar Plot for the target
+        shap.summary_plot(shap_values_for_target, instances_to_explain_np, feature_names=feature_names, plot_type="bar", show=False)
+        plt.title(f"SHAP Feature Importance for '{target_name}'")
+        plt.tight_layout()
+        bar_path = save_dir_path / f"shap_bar_plot_{sanitized_target_name}.svg"
+        plt.savefig(bar_path)
+        plt.close()
+
+        # Save Dot Plot for the target
+        shap.summary_plot(shap_values_for_target, instances_to_explain_np, feature_names=feature_names, plot_type="dot", show=False)
+        plt.title(f"SHAP Feature Importance for '{target_name}'")
+        plt.tight_layout()
+        dot_path = save_dir_path / f"shap_dot_plot_{sanitized_target_name}.svg"
+        plt.savefig(dot_path)
+        plt.close()
+
+    plt.ion()
+    _LOGGER.info(f"✅ All SHAP plots saved to '{save_dir_path.name}'")
+

ml_tools/ML_inference.py

@@ -3,6 +3,7 @@ from torch import nn
 import numpy as np
 from pathlib import Path
 from typing import Union, Literal, Dict, Any, Optional
+from abc import ABC, abstractmethod
 
 from .ML_scaler import PytorchScaler
 from ._script_info import _script_info
@@ -12,38 +13,36 @@ from .keys import PyTorchInferenceKeys
 
 __all__ = [
     "PyTorchInferenceHandler",
+    "PyTorchInferenceHandlerMulti",
     "multi_inference_regression",
     "multi_inference_classification"
 ]
 
-class PyTorchInferenceHandler:
+
+class _BaseInferenceHandler(ABC):
     """
-    Handles loading a PyTorch model's state dictionary and performing inference
-    for either regression or classification tasks.
+    Abstract base class for PyTorch inference handlers.
+
+    Manages common tasks like loading a model's state dictionary, validating
+    the target device, and preprocessing input features.
     """
     def __init__(self,
                  model: nn.Module,
                  state_dict: Union[str, Path],
-                 task: Literal["classification", "regression"],
                  device: str = 'cpu',
-                 target_id: Optional[str]=None,
                  scaler: Optional[Union[PytorchScaler, str, Path]] = None):
         """
-        Initializes the handler by loading a model's state_dict.
+        Initializes the handler.
 
         Args:
-            model (nn.Module): An instantiated PyTorch model with the correct architecture.
-            state_dict (str | Path): The path to the saved .pth model state_dict file.
-            task (str): The type of task, 'regression' or 'classification'.
+            model (nn.Module): An instantiated PyTorch model.
+            state_dict (str | Path): Path to the saved .pth model state_dict file.
             device (str): The device to run inference on ('cpu', 'cuda', 'mps').
-            target_id (str | None): Target name as used in the training set.
-            scaler (PytorchScaler | str | Path | None): A PytorchScaler instance or the file path to a saved PytorchScaler state.
+            scaler (PytorchScaler | str | Path | None): An optional scaler or path to a saved scaler state.
         """
         self.model = model
-        self.task = task
         self.device = self._validate_device(device)
-        self.target_id = target_id
-        
+
         # Load the scaler if a path is provided
         if scaler is not None:
             if isinstance(scaler, (str, Path)):
@@ -52,7 +51,7 @@ class PyTorchInferenceHandler:
                 self.scaler = scaler
         else:
             self.scaler = None
-        
+
         model_p = make_fullpath(state_dict, enforce="file")
 
         try:
@@ -72,6 +71,7 @@ class PyTorchInferenceHandler:
             _LOGGER.warning("⚠️ CUDA not available, switching to CPU.")
             device_lower = "cpu"
         elif device_lower == "mps" and not torch.backends.mps.is_available():
+            # Your M-series Mac will appreciate this check!
             _LOGGER.warning("⚠️ Apple Metal Performance Shaders (MPS) not available, switching to CPU.")
             device_lower = "cpu"
         return torch.device(device_lower)
@@ -84,51 +84,103 @@ class PyTorchInferenceHandler:
         if isinstance(features, np.ndarray):
             features_tensor = torch.from_numpy(features).float()
         else:
-            # Ensure it's a float tensor for the model
             features_tensor = features.float()
-        
-        # Apply the scaler transformation if the scaler is available
+
         if self.scaler:
             features_tensor = self.scaler.transform(features_tensor)
-        
-        # Ensure tensor is on the correct device
+
         return features_tensor.to(self.device)
-    
+
+    @abstractmethod
+    def predict_batch(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """Core batch prediction method. Must be implemented by subclasses."""
+        pass
+
+    @abstractmethod
+    def predict(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """Core single-sample prediction method. Must be implemented by subclasses."""
+        pass
+
+
+class PyTorchInferenceHandler(_BaseInferenceHandler):
+    """
+    Handles loading a PyTorch model's state dictionary and performing inference
+    for single-target regression or classification tasks.
+    """
+    def __init__(self,
+                 model: nn.Module,
+                 state_dict: Union[str, Path],
+                 task: Literal["classification", "regression"],
+                 device: str = 'cpu',
+                 target_id: Optional[str] = None,
+                 scaler: Optional[Union[PytorchScaler, str, Path]] = None):
+        """
+        Initializes the handler for single-target tasks.
+
+        Args:
+            model (nn.Module): An instantiated PyTorch model architecture.
+            state_dict (str | Path): Path to the saved .pth model state_dict file.
+            task (str): The type of task, 'regression' or 'classification'.
+            device (str): The device to run inference on ('cpu', 'cuda', 'mps').
+            target_id (str | None): An optional identifier for the target.
+            scaler (PytorchScaler | str | Path | None): A PytorchScaler instance or the file path to a saved PytorchScaler state.
+        """
+        # Call the parent constructor to handle model loading, device, and scaler
+        super().__init__(model, state_dict, device, scaler)
+
+        if task not in ["classification", "regression"]:
+            raise ValueError("`task` must be 'classification' or 'regression'.")
+        self.task = task
+        self.target_id = target_id
+
     def predict_batch(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
         """
-        Core batch prediction method. Returns results as PyTorch tensors on the model's device.
+        Core batch prediction method for single-target models.
+
+        Args:
+            features (np.ndarray | torch.Tensor): A 2D array/tensor of input features.
+
+        Returns:
+            A dictionary containing the raw output tensors from the model.
         """
         if features.ndim != 2:
             raise ValueError("Input for batch prediction must be a 2D array or tensor.")
 
         input_tensor = self._preprocess_input(features)
-        
+
         with torch.no_grad():
-            # Output tensor remains on the model's device (e.g., 'mps' or 'cuda')
             output = self.model(input_tensor)
 
             if self.task == "classification":
-                probs = nn.functional.softmax(output, dim=1)
+                probs = torch.softmax(output, dim=1)
                 labels = torch.argmax(probs, dim=1)
                 return {
                     PyTorchInferenceKeys.LABELS: labels,
                     PyTorchInferenceKeys.PROBABILITIES: probs
                 }
             else:  # regression
-                return {PyTorchInferenceKeys.PREDICTIONS: output}
+                # For single-target regression, ensure output is flattened
+                return {PyTorchInferenceKeys.PREDICTIONS: output.flatten()}
 
     def predict(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
         """
-        Core single-sample prediction. Returns results as PyTorch tensors on the model's device.
+        Core single-sample prediction method for single-target models.
+
+        Args:
+            features (np.ndarray | torch.Tensor): A 1D array/tensor of input features.
+
+        Returns:
+            A dictionary containing the raw output tensors for a single sample.
         """
         if features.ndim == 1:
-            features = features.reshape(1, -1)
-        
+            features = features.reshape(1, -1) # Reshape to a batch of one
+
         if features.shape[0] != 1:
             raise ValueError("The predict() method is for a single sample. Use predict_batch() for multiple samples.")
 
         batch_results = self.predict_batch(features)
-        
+
+        # Extract the first (and only) result from the batch output
         single_results = {key: value[0] for key, value in batch_results.items()}
         return single_results
 
@@ -139,7 +191,6 @@ class PyTorchInferenceHandler:
         Convenience wrapper for predict_batch that returns NumPy arrays.
         """
         tensor_results = self.predict_batch(features)
-        # Move tensor to CPU before converting to NumPy
         numpy_results = {key: value.cpu().numpy() for key, value in tensor_results.items()}
         return numpy_results
 
@@ -148,16 +199,163 @@ class PyTorchInferenceHandler:
         Convenience wrapper for predict that returns NumPy arrays or scalars.
         """
         tensor_results = self.predict(features)
-        
+
         if self.task == "regression":
-            # .item() implicitly moves to CPU
+            # .item() implicitly moves to CPU and returns a Python scalar
             return {PyTorchInferenceKeys.PREDICTIONS: tensor_results[PyTorchInferenceKeys.PREDICTIONS].item()}
         else: # classification
             return {
                 PyTorchInferenceKeys.LABELS: tensor_results[PyTorchInferenceKeys.LABELS].item(),
-                # Move tensor to CPU before converting to NumPy
                 PyTorchInferenceKeys.PROBABILITIES: tensor_results[PyTorchInferenceKeys.PROBABILITIES].cpu().numpy()
             }
+    
+    def quick_predict(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, Any]:
+        """
+        Convenience wrapper to get the mapping {target_name: prediction} or {target_name: label}
+        
+        `target_id` must be implemented.
+        """
+        if self.target_id is None:
+            raise AttributeError(f"'target_id' has not been implemented.")
+        
+        if self.task == "regression":
+            result = self.predict_numpy(features)[PyTorchInferenceKeys.PREDICTIONS]
+        else:
+            result = self.predict_numpy(features)[PyTorchInferenceKeys.LABELS]
+        
+        return {self.target_id: result}
+
+
+class PyTorchInferenceHandlerMulti(_BaseInferenceHandler):
+    """
+    Handles loading a PyTorch model's state dictionary and performing inference
+    for multi-target regression or multi-label classification tasks.
+    """
+    def __init__(self,
+                 model: nn.Module,
+                 state_dict: Union[str, Path],
+                 task: Literal["multi_target_regression", "multi_label_classification"],
+                 device: str = 'cpu',
+                 target_ids: Optional[list[str]] = None,
+                 scaler: Optional[Union[PytorchScaler, str, Path]] = None):
+        """
+        Initializes the handler for multi-target tasks.
+
+        Args:
+            model (nn.Module): An instantiated PyTorch model.
+            state_dict (str | Path): Path to the saved .pth model state_dict file.
+            task (str): The type of task, 'multi_target_regression' or 'multi_label_classification'.
+            device (str): The device to run inference on ('cpu', 'cuda', 'mps').
+            target_ids (list[str] | None): An optional identifier for the targets.
+            scaler (PytorchScaler | str | Path | None): A PytorchScaler instance or the file path to a saved PytorchScaler state.
+        """
+        super().__init__(model, state_dict, device, scaler)
+
+        if task not in ["multi_target_regression", "multi_label_classification"]:
+            raise ValueError("`task` must be 'multi_target_regression' or 'multi_label_classification'.")
+        self.task = task
+        self.target_ids = target_ids
+
+    def predict_batch(self,
+                      features: Union[np.ndarray, torch.Tensor],
+                      classification_threshold: float = 0.5
+                      ) -> Dict[str, torch.Tensor]:
+        """
+        Core batch prediction method for multi-target models.
+
+        Args:
+            features (np.ndarray | torch.Tensor): A 2D array/tensor of input features.
+            classification_threshold (float): The threshold to convert probabilities
+                into binary predictions for multi-label classification.
+
+        Returns:
+            A dictionary containing the raw output tensors from the model.
+        """
+        if features.ndim != 2:
+            raise ValueError("Input for batch prediction must be a 2D array or tensor.")
+
+        input_tensor = self._preprocess_input(features)
+
+        with torch.no_grad():
+            output = self.model(input_tensor)
+
+            if self.task == "multi_label_classification":
+                probs = torch.sigmoid(output)
+                # Get binary predictions based on the threshold
+                labels = (probs >= classification_threshold).int()
+                return {
+                    PyTorchInferenceKeys.LABELS: labels,
+                    PyTorchInferenceKeys.PROBABILITIES: probs
+                }
+            else:  # multi_target_regression
+                # The output is already in the correct [batch_size, n_targets] shape
+                return {PyTorchInferenceKeys.PREDICTIONS: output}
+
+    def predict(self,
+                features: Union[np.ndarray, torch.Tensor],
+                classification_threshold: float = 0.5
+                ) -> Dict[str, torch.Tensor]:
+        """
+        Core single-sample prediction method for multi-target models.
+
+        Args:
+            features (np.ndarray | torch.Tensor): A 1D array/tensor of input features.
+            classification_threshold (float): The threshold for multi-label tasks.
+
+        Returns:
+            A dictionary containing the raw output tensors for a single sample.
+        """
+        if features.ndim == 1:
+            features = features.reshape(1, -1)
+
+        if features.shape[0] != 1:
+            raise ValueError("The predict() method is for a single sample. Use predict_batch() for multiple samples.")
+
+        batch_results = self.predict_batch(features, classification_threshold)
+
+        single_results = {key: value[0] for key, value in batch_results.items()}
+        return single_results
+
+    # --- NumPy Convenience Wrappers (on CPU) ---
+
+    def predict_batch_numpy(self,
+                            features: Union[np.ndarray, torch.Tensor],
+                            classification_threshold: float = 0.5
+                            ) -> Dict[str, np.ndarray]:
+        """
+        Convenience wrapper for predict_batch that returns NumPy arrays.
+        """
+        tensor_results = self.predict_batch(features, classification_threshold)
+        numpy_results = {key: value.cpu().numpy() for key, value in tensor_results.items()}
+        return numpy_results
+
+    def predict_numpy(self,
+                      features: Union[np.ndarray, torch.Tensor],
+                      classification_threshold: float = 0.5
+                      ) -> Dict[str, np.ndarray]:
+        """
+        Convenience wrapper for predict that returns NumPy arrays for a single sample.
+        Note: For multi-target models, the output is always an array.
+        """
+        tensor_results = self.predict(features, classification_threshold)
+        numpy_results = {key: value.cpu().numpy() for key, value in tensor_results.items()}
+        return numpy_results
+    
+    def quick_predict(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, Any]:
+        """
+        Convenience wrapper to get the mapping {target_name: prediction} or {target_name: label}
+        
+        `target_ids` must be implemented.
+        """
+        if self.target_ids is None:
+            raise AttributeError(f"'target_id' has not been implemented.")
+        
+        if self.task == "multi_target_regression":
+            result = self.predict_numpy(features)[PyTorchInferenceKeys.PREDICTIONS].flatten().tolist()
+        else:
+            result = self.predict_numpy(features)[PyTorchInferenceKeys.LABELS].flatten().tolist()
+        
+        return {key: value for key, value in zip(self.target_ids, result)}
 
 
 def multi_inference_regression(handlers: list[PyTorchInferenceHandler], 

ml_tools/ML_models.py

@@ -89,10 +89,6 @@ class _BaseMLP(nn.Module):
 class MultilayerPerceptron(_BaseMLP):
     """
     Creates a versatile Multilayer Perceptron (MLP) for regression or classification tasks.
-
-    This model generates raw output values (logits) suitable for use with loss
-    functions like `nn.CrossEntropyLoss` (for classification) or `nn.MSELoss`
-    (for regression).
     """
     def __init__(self, in_features: int, out_targets: int,
                  hidden_layers: List[int] = [256, 128], drop_out: float = 0.2) -> None: