dragon-ml-toolbox 19.13.0__py3-none-any.whl → 20.0.0__py3-none-any.whl

ml_tools/_core/_ML_inference.py

@@ -1,646 +0,0 @@
-import torch
-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_finalize_handler import FinalizedFileHandler
-from ._ML_scaler import DragonScaler
-from ._script_info import _script_info
-from ._logger import get_logger
-from ._path_manager import make_fullpath
-from ._keys import PyTorchInferenceKeys, PyTorchCheckpointKeys, MLTaskKeys, ScalerKeys, MagicWords
-
-
-_LOGGER = get_logger("Inference Handler")
-
-
-__all__ = [
-    "DragonInferenceHandler",
-    "multi_inference_regression",
-    "multi_inference_classification"
-]
-
-
-class _BaseInferenceHandler(ABC):
-    """
-    Abstract base class for PyTorch inference handlers.
-
-    Manages common tasks like loading a model's state dictionary via FinalizedFileHandler,
-    validating the target device, and preprocessing input features.
-    """
-    def __init__(self,
-                 model: nn.Module,
-                 state_dict: Union[str, Path],
-                 device: str = 'cpu',
-                 scaler: Optional[Union[str, Path]] = None,
-                 task: Optional[str] = None):
-        """
-        Initializes the handler.
-
-        Args:
-            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').
-            scaler (str | Path | None): An optional scaler or path to a saved scaler state.
-            task (str | None): The specific machine learning task. If None, it attempts to read it from the finalized-file.
-        """
-        self.model = model
-        self.device = self._validate_device(device)
-        self._classification_threshold = 0.5
-        self._loaded_threshold: bool = False
-        self._loaded_class_map: bool = False
-        self._class_map: Optional[dict[str,int]] = None
-        self._idx_to_class: Optional[Dict[int, str]] = None
-        
-        # --- 1. Load File Handler ---
-        # This loads the content on CPU and validates structure
-        self._file_handler = FinalizedFileHandler(state_dict)
-        
-        # Silence warnings of the filehandler internally
-        self._file_handler._verbose = False
-        
-        # --- 2. Task Resolution ---
-        file_task = self._file_handler.task
-        
-        if task is None:
-            # User didn't provide task, must be in file
-            if file_task == MagicWords.UNKNOWN:
-                _LOGGER.error(f"Task not specified in arguments and not found in file '{make_fullpath(state_dict).name}'.")
-                raise ValueError()
-            self.task = file_task
-            _LOGGER.info(f"Task '{self.task}' detected from file.")
-        else:
-            # User provided task
-            if file_task != MagicWords.UNKNOWN and file_task != task:
-                _LOGGER.warning(f"Provided task '{task}' differs from file metadata task '{file_task}'. Using provided task '{task}'.")
-            self.task = task
-
-        # --- 3. Load Model Weights ---
-        # Weights are already loaded in file_handler (on CPU)
-        try:
-            self.model.load_state_dict(self._file_handler.model_state_dict)
-        except RuntimeError as e:
-            _LOGGER.error(f"State dict mismatch: {e}")
-            raise
-
-        # --- 4. Load Metadata (Thresholds, Class Maps) ---
-        if self._file_handler.classification_threshold is not None:
-            self._classification_threshold = self._file_handler.classification_threshold
-            self._loaded_threshold = True
-            
-        if self._file_handler.class_map is not None:
-            self.set_class_map(self._file_handler.class_map)
-            # set_class_map sets _loaded_class_map to True
-        
-        # --- 5. Move to Device ---
-        self.model.to(self.device)
-        self.model.eval()
-        _LOGGER.info(f"Model loaded and moved to {self.device} in evaluation mode.")
-
-        # --- 6. Load Scalers ---
-        self.feature_scaler: Optional[DragonScaler] = None
-        self.target_scaler: Optional[DragonScaler] = None
-
-        if scaler is not None:
-            if isinstance(scaler, (str, Path)):
-                path_obj = make_fullpath(scaler, enforce="file")
-                loaded_scaler_data = torch.load(path_obj)
-                
-                if isinstance(loaded_scaler_data, dict) and (ScalerKeys.FEATURE_SCALER in loaded_scaler_data or ScalerKeys.TARGET_SCALER in loaded_scaler_data):
-                    if ScalerKeys.FEATURE_SCALER in loaded_scaler_data:
-                        self.feature_scaler = DragonScaler.load(loaded_scaler_data[ScalerKeys.FEATURE_SCALER], verbose=False)
-                        _LOGGER.info("Loaded DragonScaler state for feature scaling.")
-                    if ScalerKeys.TARGET_SCALER in loaded_scaler_data:
-                        self.target_scaler = DragonScaler.load(loaded_scaler_data[ScalerKeys.TARGET_SCALER], verbose=False)
-                        _LOGGER.info("Loaded DragonScaler state for target scaling.")
-                else:
-                    _LOGGER.warning("Loaded scaler file does not contain separate feature/target scalers. Assuming it is a feature scaler (legacy format).")
-                    self.feature_scaler = DragonScaler.load(loaded_scaler_data)
-            else:
-                _LOGGER.error("Scaler must be a file path (str or Path) to a saved DragonScaler state file.")
-                raise ValueError()
-
-    def _validate_device(self, device: str) -> torch.device:
-        """Validates the selected device and returns a torch.device object."""
-        device_lower = device.lower()
-        if "cuda" in device_lower and not torch.cuda.is_available():
-            _LOGGER.warning("CUDA not available, switching to CPU.")
-            device_lower = "cpu"
-        elif device_lower == "mps" and not torch.backends.mps.is_available():
-            _LOGGER.warning("Apple Metal Performance Shaders (MPS) not available, switching to CPU.")
-            device_lower = "cpu"
-        return torch.device(device_lower)
-    
-    def set_class_map(self, class_map: Dict[str, int], force_overwrite: bool = False):
-        """
-        Sets the class name mapping to translate predicted integer labels back into string names.
-        
-        Args:
-            class_map (Dict[str, int]): The class_to_idx dictionary.
-            force_overwrite (bool): If True, allows overwriting a map that was loaded from a configuration file.
-        """
-        if self._loaded_class_map:
-            warning_message = f"A '{PyTorchCheckpointKeys.CLASS_MAP}' was loaded from the model configuration file."
-            if not force_overwrite:
-                warning_message += " Use 'force_overwrite=True' if you are sure you want to modify it. This will not affect the value from the file."
-                _LOGGER.warning(warning_message)
-                return
-            else:
-                warning_message += " Overwriting it for this inference instance."
-                _LOGGER.warning(warning_message)
-        
-        self._class_map = class_map
-        self._idx_to_class = {v: k for k, v in class_map.items()}
-        self._loaded_class_map = True
-        _LOGGER.info("InferenceHandler: Class map set for label-to-name translation.")
-
-    @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 DragonInferenceHandler(_BaseInferenceHandler):
-    """
-    Handles loading a PyTorch model's state dictionary and performing inference for tabular data.
-    """
-    def __init__(self,
-                 model: nn.Module,
-                 state_dict: Union[str, Path],
-                 task: Optional[Literal["regression", 
-                               "binary classification", 
-                               "multiclass classification", 
-                               "multitarget regression", 
-                               "multilabel binary classification"]] = None,
-                 device: str = 'cpu',
-                 scaler: Optional[Union[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, optional): The type of task. If None, it will be detected from file.
-            device (str): The device to run inference on ('cpu', 'cuda', 'mps').
-            scaler (str | Path | None): A path to a saved DragonScaler state.
-            
-        Note: class_map (Dict[int, str]) will be loaded from the model file, to set or override it use `.set_class_map()`.
-        """
-        # Call the parent constructor to handle model loading, device, and scaler
-        # The parent constructor resolves 'task'
-        super().__init__(model=model, 
-                         state_dict=state_dict, 
-                         device=device, 
-                         scaler=scaler, 
-                         task=task)
-        
-        # --- Validation of resolved task ---
-        valid_tasks = [
-            MLTaskKeys.REGRESSION,
-            MLTaskKeys.BINARY_CLASSIFICATION, 
-            MLTaskKeys.MULTICLASS_CLASSIFICATION,
-            MLTaskKeys.MULTITARGET_REGRESSION,
-            MLTaskKeys.MULTILABEL_BINARY_CLASSIFICATION
-        ]
-        
-        if self.task not in valid_tasks:
-            _LOGGER.error(f"'task' recognized as '{self.task}', but this inference handler only supports: {valid_tasks}.")
-            raise ValueError()
-
-        self.target_ids: Optional[list[str]] = None
-        self._target_ids_set: bool = False
-        
-        # --- Attempt to load target names from FinalizedFileHandler ---
-        if self._file_handler.target_names is not None:
-            self.set_target_ids(self._file_handler.target_names)
-        elif self._file_handler.target_name is not None:
-            self.set_target_ids([self._file_handler.target_name])
-        else:
-            _LOGGER.warning("No target names found in file metadata.")
-
-    def _preprocess_input(self, features: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
-        """
-        Converts input to a torch.Tensor, applies FEATURE scaling if a scaler is
-        present, and moves it to the correct device.
-        """
-        if isinstance(features, np.ndarray):
-            features_tensor = torch.from_numpy(features).float()
-        else:
-            features_tensor = features.float()
-
-        if self.feature_scaler:
-            features_tensor = self.feature_scaler.transform(features_tensor)
-
-        return features_tensor.to(self.device)
-    
-    def set_target_ids(self, target_names: list[str], force_overwrite: bool=False):
-        """
-        Assigns the provided list of strings as the target variable names.
-        
-        If target IDs have already been set, this method will log a warning.
-
-        Args:
-            target_names (list[str]): A list of target names.
-            force_overwrite (bool): If True, allows the method to overwrite previously set target IDs.
-        """
-        if self._target_ids_set:
-            warning_message = "Target IDs was previously set."
-            if not force_overwrite:
-                warning_message += " Use `force_overwrite=True` to overwrite."
-                _LOGGER.warning(warning_message)
-                return
-            else:
-                warning_message += " Overwriting..."
-                _LOGGER.warning(warning_message)
-
-        self.target_ids = target_names
-        self._target_ids_set = True
-        _LOGGER.info("Target IDs set.")
-
-    def predict_batch(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
-        """
-        Core batch prediction method.
-
-        Args:
-            features (np.ndarray | torch.Tensor): A 2D array/tensor of input features.
-
-        Returns:
-            Dict: A dictionary containing the raw output tensors from the model.
-        """
-        if features.ndim != 2:
-            _LOGGER.error("Input for batch prediction must be a 2D array or tensor.")
-            raise ValueError()
-
-        input_tensor = self._preprocess_input(features)
-
-        with torch.no_grad():
-            output = self.model(input_tensor)
-
-            # --- Target Scaling Logic (Inverse Transform) ---
-            # Only for regression tasks and if a target scaler exists
-            if self.target_scaler:
-                if self.task not in [MLTaskKeys.REGRESSION, MLTaskKeys.MULTITARGET_REGRESSION]:
-                    # raise error
-                    _LOGGER.error("Target scaler is only applicable for regression tasks. A target scaler was provided for a non-regression task.")
-                    raise ValueError()
-                
-                # Ensure output is 2D (N, Targets) for the scaler
-                original_shape = output.shape
-                if output.ndim == 1:
-                    output = output.reshape(-1, 1)
-                
-                # Apply inverse transform (de-scale)
-                output = self.target_scaler.inverse_transform(output)
-                
-                # Restore original shape if necessary (though usually we want 2D or 1D flat)
-                if len(original_shape) == 1:
-                    output = output.flatten()
-
-            # --- Task Specific Formatting ---
-            if self.task == MLTaskKeys.MULTICLASS_CLASSIFICATION:
-                probs = torch.softmax(output, dim=1)
-                labels = torch.argmax(probs, dim=1)
-                return {
-                    PyTorchInferenceKeys.LABELS: labels,
-                    PyTorchInferenceKeys.PROBABILITIES: probs
-                }
-                
-            elif self.task == MLTaskKeys.BINARY_CLASSIFICATION:
-                if output.ndim == 2 and output.shape[1] == 1:
-                    output = output.squeeze(1)
-                    
-                probs = torch.sigmoid(output) 
-                labels = (probs >= self._classification_threshold).int()
-                return {
-                    PyTorchInferenceKeys.LABELS: labels,
-                    PyTorchInferenceKeys.PROBABILITIES: probs
-                }
-                
-            elif self.task == MLTaskKeys.REGRESSION:
-                # For single-target regression, ensure output is flattened
-                return {PyTorchInferenceKeys.PREDICTIONS: output.flatten()}
-            
-            elif self.task == MLTaskKeys.MULTILABEL_BINARY_CLASSIFICATION:
-                probs = torch.sigmoid(output)
-                labels = (probs >= self._classification_threshold).int()
-                return {
-                    PyTorchInferenceKeys.LABELS: labels,
-                    PyTorchInferenceKeys.PROBABILITIES: probs
-                }
-            
-            elif self.task == MLTaskKeys.MULTITARGET_REGRESSION:
-                return {PyTorchInferenceKeys.PREDICTIONS: output}
-            
-            else:
-                _LOGGER.error(f"Unrecognized task '{self.task}'.")
-                raise ValueError()
-
-    def predict(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
-        """
-        Core single-sample prediction method for single-target models.
-
-        Args:
-            features (np.ndarray | torch.Tensor): A 1D array/tensor of input features.
-
-        Returns:
-            Dict: A dictionary containing the raw output tensors for a single sample.
-        """
-        if features.ndim == 1:
-            features = features.reshape(1, -1) # Reshape to a batch of one
-
-        if features.shape[0] != 1:
-            _LOGGER.error("The 'predict()' method is for a single sample. Use 'predict_batch()' for multiple samples.")
-            raise ValueError()
-
-        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
-    
-    # --- NumPy Convenience Wrappers (on CPU) ---
-
-    def predict_batch_numpy(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, np.ndarray]:
-        """
-        Convenience wrapper for predict_batch that returns NumPy arrays
-        and adds string labels for classification tasks if a class_map is set.
-        """
-        tensor_results = self.predict_batch(features)
-        numpy_results = {key: value.cpu().numpy() for key, value in tensor_results.items()}
-        
-        # Add string names for classification if map exists
-        is_classification = self.task in [
-            MLTaskKeys.BINARY_CLASSIFICATION, 
-            MLTaskKeys.MULTICLASS_CLASSIFICATION
-        ]
-        
-        if is_classification and self._idx_to_class and PyTorchInferenceKeys.LABELS in numpy_results:
-            int_labels = numpy_results[PyTorchInferenceKeys.LABELS] # This is a (B,) array
-            numpy_results[PyTorchInferenceKeys.LABEL_NAMES] = [ # type: ignore
-                self._idx_to_class.get(label_id, "Unknown")
-                for label_id in int_labels
-            ]
-        
-        return numpy_results
-
-    def predict_numpy(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, Any]:
-        """
-        Convenience wrapper for predict that returns NumPy arrays or scalars
-        and adds string labels for classification tasks if a class_map is set.
-        """
-        tensor_results = self.predict(features)
-
-        if self.task == MLTaskKeys.REGRESSION:
-            # .item() implicitly moves to CPU and returns a Python scalar
-            return {PyTorchInferenceKeys.PREDICTIONS: tensor_results[PyTorchInferenceKeys.PREDICTIONS].item()}
-        
-        elif self.task in [MLTaskKeys.BINARY_CLASSIFICATION, MLTaskKeys.MULTICLASS_CLASSIFICATION]:
-            int_label = tensor_results[PyTorchInferenceKeys.LABELS].item()
-            label_name = "Unknown"
-            if self._idx_to_class:
-                label_name = self._idx_to_class.get(int_label, "Unknown") # type: ignore
-
-            return {
-                PyTorchInferenceKeys.LABELS: int_label,
-                PyTorchInferenceKeys.LABEL_NAMES: label_name,
-                PyTorchInferenceKeys.PROBABILITIES: tensor_results[PyTorchInferenceKeys.PROBABILITIES].cpu().numpy()
-            }
-            
-        elif self.task in [MLTaskKeys.MULTILABEL_BINARY_CLASSIFICATION, MLTaskKeys.MULTITARGET_REGRESSION]:
-            # For multi-target models, the output is always an array.
-            numpy_results = {key: value.cpu().numpy() for key, value in tensor_results.items()}
-            return numpy_results
-        else:
-            # should never happen
-            _LOGGER.error(f"Unrecognized task '{self.task}'.")
-            raise ValueError()
-    
-    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:
-            _LOGGER.error(f"'target_ids' has not been implemented.")
-            raise AttributeError()
-        
-        if self.task == MLTaskKeys.REGRESSION:
-            result = self.predict_numpy(features)[PyTorchInferenceKeys.PREDICTIONS]
-            return {self.target_ids[0]: result}
-        
-        elif self.task in [MLTaskKeys.BINARY_CLASSIFICATION, MLTaskKeys.MULTICLASS_CLASSIFICATION]:
-            result = self.predict_numpy(features)[PyTorchInferenceKeys.LABELS]
-            return {self.target_ids[0]: result}
-        
-        elif self.task == MLTaskKeys.MULTITARGET_REGRESSION:
-            result = self.predict_numpy(features)[PyTorchInferenceKeys.PREDICTIONS].flatten().tolist()
-            return {key: value for key, value in zip(self.target_ids, result)}
-        
-        elif self.task == MLTaskKeys.MULTILABEL_BINARY_CLASSIFICATION:
-            result = self.predict_numpy(features)[PyTorchInferenceKeys.LABELS].flatten().tolist()
-            return {key: value for key, value in zip(self.target_ids, result)}
-        
-        else:
-            # should never happen
-            _LOGGER.error(f"Unrecognized task '{self.task}'.")
-            raise ValueError()
-        
-    def set_classification_threshold(self, threshold: float, force_overwrite: bool=False):
-        """
-        Sets the classification threshold for the current inference instance.
-        
-        If a threshold was previously loaded from a model configuration, this
-        method will log a warning and refuse to update the value. This
-        prevents accidentally overriding a setting from a loaded checkpoint.
-        
-        To bypass this safety check set `force_overwrite` to `True`.
-
-        Args:
-            threshold (float): The new classification threshold value to set.
-            force_overwrite (bool): If True, allows overwriting a threshold that was loaded from a configuration file. 
-        """
-        if self._loaded_threshold:
-            warning_message = f"The current '{PyTorchCheckpointKeys.CLASSIFICATION_THRESHOLD}={self._classification_threshold}' was loaded and set from a model configuration file."
-            if not force_overwrite:
-                warning_message += " Use 'force_overwrite' if you are sure you want to modify it. This will not affect the value from the file."
-                _LOGGER.warning(warning_message)
-                return
-            else:
-                warning_message += f" Overwriting it to {threshold}."
-                _LOGGER.warning(warning_message)
- 
-        self._classification_threshold = threshold
-
-
-def multi_inference_regression(handlers: list[DragonInferenceHandler], 
-                               feature_vector: Union[np.ndarray, torch.Tensor], 
-                               output: Literal["numpy","torch"]="numpy") -> dict[str,Any]:
-    """
-    Performs regression inference using multiple models on a single feature vector.
-
-    This function iterates through a list of DragonInferenceHandler objects,
-    each configured for a different regression target. It runs a prediction for
-    each handler using the same input feature vector and returns the results
-    in a dictionary.
-    
-    The function adapts its behavior based on the input dimensions:
-    - 1D input: Returns a dictionary mapping target ID to a single value.
-    - 2D input: Returns a dictionary mapping target ID to a list of values.
-
-    Args:
-        handlers (list[DragonInferenceHandler]): A list of initialized inference
-            handlers. Each handler must have a unique `target_id` and be configured with `task="regression"`.
-        feature_vector (Union[np.ndarray, torch.Tensor]): An input sample (1D) or a batch of samples (2D) to be fed into each regression model.
-        output (Literal["numpy", "torch"], optional): The desired format for the output predictions.
-            - "numpy": Returns predictions as Python scalars or NumPy arrays.
-            - "torch": Returns predictions as PyTorch tensors.
-
-    Returns:
-        (dict[str, Any]): A dictionary mapping each handler's `target_id` to its
-        predicted regression values. 
-
-    Raises:
-        AttributeError: If any handler in the list is missing a `target_id`.
-        ValueError: If any handler's `task` is not 'regression' or if the input `feature_vector` is not 1D or 2D.
-    """
-    # check batch dimension
-    is_single_sample = feature_vector.ndim == 1
-    
-    # Reshape a 1D vector to a 2D batch of one for uniform processing.
-    if is_single_sample:
-        feature_vector = feature_vector.reshape(1, -1)
-    
-    # Validate that the input is a 2D tensor.
-    if feature_vector.ndim != 2:
-        _LOGGER.error("Input feature_vector must be a 1D or 2D array/tensor.")
-        raise ValueError()
-    
-    results: dict[str,Any] = dict()
-    for handler in handlers:
-        # validation
-        if handler.target_ids is None:
-            _LOGGER.error("All inference handlers must have a 'target_ids' attribute.")
-            raise AttributeError()
-        if handler.task != MLTaskKeys.REGRESSION:
-            _LOGGER.error(f"Invalid task type: The handler for target_id '{handler.target_ids[0]}' is for '{handler.task}', only single target regression tasks are supported.")
-            raise ValueError()
-            
-        # inference
-        if output == "numpy":
-            # This path returns NumPy arrays or standard Python scalars
-            numpy_result = handler.predict_batch_numpy(feature_vector)[PyTorchInferenceKeys.PREDICTIONS]
-            if is_single_sample:
-                # For a single sample, convert the 1-element array to a Python scalar
-                results[handler.target_ids[0]] = numpy_result.item()
-            else:
-                # For a batch, return the full NumPy array of predictions
-                results[handler.target_ids[0]] = numpy_result
-
-        else:  # output == "torch"
-            # This path returns PyTorch tensors on the model's device
-            torch_result = handler.predict_batch(feature_vector)[PyTorchInferenceKeys.PREDICTIONS]
-            if is_single_sample:
-                # For a single sample, return the 0-dim tensor
-                results[handler.target_ids[0]] = torch_result[0]
-            else:
-                # For a batch, return the full tensor of predictions
-                results[handler.target_ids[0]] = torch_result
-
-    return results
-
-
-def multi_inference_classification(
-    handlers: list[DragonInferenceHandler], 
-    feature_vector: Union[np.ndarray, torch.Tensor], 
-    output: Literal["numpy","torch"]="numpy"
-    ) -> tuple[dict[str, Any], dict[str, Any]]:
-    """
-    Performs classification inference on a single sample or a batch.
-
-    This function iterates through a list of DragonInferenceHandler objects,
-    each configured for a different classification target. It returns two
-    dictionaries: one for the predicted labels and one for the probabilities.
-
-    The function adapts its behavior based on the input dimensions:
-    - 1D input: The dictionaries map target ID to a single label and a single probability array.
-    - 2D input: The dictionaries map target ID to an array of labels and an array of probability arrays.
-
-    Args:
-        handlers (list[DragonInferenceHandler]): A list of initialized inference handlers. Each must have a unique `target_id` and be configured
-            with `task="classification"`.
-        feature_vector (Union[np.ndarray, torch.Tensor]): An input sample (1D)
-            or a batch of samples (2D) for prediction.
-        output (Literal["numpy", "torch"], optional): The desired format for the
-            output predictions.
-
-    Returns:
-        (tuple[dict[str, Any], dict[str, Any]]): A tuple containing two dictionaries:
-        1.  A dictionary mapping `target_id` to the predicted label(s).
-        2.  A dictionary mapping `target_id` to the prediction probabilities.
-
-    Raises:
-        AttributeError: If any handler in the list is missing a `target_id`.
-        ValueError: If any handler's `task` is not 'classification' or if the input `feature_vector` is not 1D or 2D.
-    """
-    # Store if the original input was a single sample
-    is_single_sample = feature_vector.ndim == 1
-    
-    # Reshape a 1D vector to a 2D batch of one for uniform processing
-    if is_single_sample:
-        feature_vector = feature_vector.reshape(1, -1)
-    
-    if feature_vector.ndim != 2:
-        _LOGGER.error("Input feature_vector must be a 1D or 2D array/tensor.")
-        raise ValueError()
-
-    # Initialize two dictionaries for results
-    labels_results: dict[str, Any] = dict()
-    probs_results: dict[str, Any] = dict()
-
-    for handler in handlers:
-        # Validation
-        if handler.target_ids is None:
-            _LOGGER.error("All inference handlers must have a 'target_id' attribute.")
-            raise AttributeError()
-        if handler.task not in [MLTaskKeys.BINARY_CLASSIFICATION, MLTaskKeys.MULTICLASS_CLASSIFICATION]:
-            _LOGGER.error(f"Invalid task type: The handler for target_id '{handler.target_ids[0]}' is for '{handler.task}', but this function only supports binary and multiclass classification.")
-            raise ValueError()
-            
-        # Inference
-        if output == "numpy":
-            # predict_batch_numpy returns a dict of NumPy arrays
-            result = handler.predict_batch_numpy(feature_vector)
-        else: # torch
-            # predict_batch returns a dict of Torch tensors
-            result = handler.predict_batch(feature_vector)
-        
-        labels = result[PyTorchInferenceKeys.LABELS]
-        probabilities = result[PyTorchInferenceKeys.PROBABILITIES]
-        
-        if is_single_sample:
-            # For "numpy", convert the single label to a Python int scalar.
-            # For "torch", get the 0-dim tensor label.
-            if output == "numpy":
-                labels_results[handler.target_ids[0]] = labels.item()
-            else: # torch
-                labels_results[handler.target_ids[0]] = labels[0]
-            
-            # The probabilities are an array/tensor of values
-            probs_results[handler.target_ids[0]] = probabilities[0]
-        else:
-            labels_results[handler.target_ids[0]] = labels
-            probs_results[handler.target_ids[0]] = probabilities
-            
-    return labels_results, probs_results
-
-
-def info():
-    _script_info(__all__)