dragon-ml-toolbox 3.12.5__py3-none-any.whl → 4.0.0__py3-none-any.whl

dragon_ml_toolbox-4.0.0.dist-info/METADATA

@@ -0,0 +1,230 @@
+Metadata-Version: 2.4
+Name: dragon-ml-toolbox
+Version: 4.0.0
+Summary: A collection of tools for data science and machine learning projects.
+Author-email: Karl Loza <luigiloza@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/DrAg0n-BoRn/ML_tools
+Project-URL: Changelog, https://github.com/DrAg0n-BoRn/ML_tools/blob/master/CHANGELOG.md
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+License-File: LICENSE-THIRD-PARTY.md
+Provides-Extra: ml
+Requires-Dist: numpy; extra == "ml"
+Requires-Dist: pandas; extra == "ml"
+Requires-Dist: polars; extra == "ml"
+Requires-Dist: joblib; extra == "ml"
+Requires-Dist: scikit-learn; extra == "ml"
+Requires-Dist: matplotlib; extra == "ml"
+Requires-Dist: seaborn; extra == "ml"
+Requires-Dist: imbalanced-learn; extra == "ml"
+Requires-Dist: ipython; extra == "ml"
+Requires-Dist: ipykernel; extra == "ml"
+Requires-Dist: notebook; extra == "ml"
+Requires-Dist: jupyterlab; extra == "ml"
+Requires-Dist: ipywidgets; extra == "ml"
+Requires-Dist: xgboost; extra == "ml"
+Requires-Dist: lightgbm; extra == "ml"
+Requires-Dist: shap; extra == "ml"
+Requires-Dist: tqdm; extra == "ml"
+Requires-Dist: Pillow; extra == "ml"
+Provides-Extra: mice
+Requires-Dist: numpy<2.0; extra == "mice"
+Requires-Dist: pandas; extra == "mice"
+Requires-Dist: polars; extra == "mice"
+Requires-Dist: joblib; extra == "mice"
+Requires-Dist: miceforest>=6.0.0; extra == "mice"
+Requires-Dist: plotnine>=0.12; extra == "mice"
+Requires-Dist: matplotlib; extra == "mice"
+Requires-Dist: statsmodels; extra == "mice"
+Requires-Dist: lightgbm<=4.5.0; extra == "mice"
+Requires-Dist: shap; extra == "mice"
+Provides-Extra: pytorch
+Requires-Dist: torch; extra == "pytorch"
+Requires-Dist: torchvision; extra == "pytorch"
+Provides-Extra: excel
+Requires-Dist: pandas; extra == "excel"
+Requires-Dist: openpyxl; extra == "excel"
+Requires-Dist: ipython; extra == "excel"
+Requires-Dist: ipykernel; extra == "excel"
+Requires-Dist: notebook; extra == "excel"
+Requires-Dist: jupyterlab; extra == "excel"
+Requires-Dist: ipywidgets; extra == "excel"
+Provides-Extra: gui-boost
+Requires-Dist: numpy; extra == "gui-boost"
+Requires-Dist: joblib; extra == "gui-boost"
+Requires-Dist: FreeSimpleGUI>=5.2; extra == "gui-boost"
+Requires-Dist: pyinstaller; extra == "gui-boost"
+Requires-Dist: xgboost; extra == "gui-boost"
+Requires-Dist: lightgbm; extra == "gui-boost"
+Provides-Extra: gui-torch
+Requires-Dist: numpy; extra == "gui-torch"
+Requires-Dist: FreeSimpleGUI>=5.2; extra == "gui-torch"
+Requires-Dist: pyinstaller; extra == "gui-torch"
+Provides-Extra: plot
+Requires-Dist: matplotlib; extra == "plot"
+Requires-Dist: seaborn; extra == "plot"
+Dynamic: license-file
+
+# dragon-ml-toolbox
+
+A collection of Python utilities for data science and machine learning, structured as a modular package for easy reuse and installation. This package has no base dependencies, allowing for lightweight and customized virtual environments.
+
+### Features:
+
+- Modular scripts for data exploration, logging, machine learning, and more.
+- Designed for seamless integration as a Git submodule or installable Python package.
+
+## Installation
+
+**Python 3.10+**
+
+### Via PyPI
+
+Install the latest stable release from PyPI:
+
+```bash
+pip install dragon-ml-toolbox
+```
+
+### Via GitHub (Editable)
+
+Clone the repository and install in editable mode with optional dependencies:
+
+```bash
+git clone https://github.com/DrAg0n-BoRn/ML_tools.git
+cd ML_tools
+pip install -e .
+```
+
+### Via conda-forge
+
+Install from the conda-forge channel:
+
+```bash
+conda install -c conda-forge dragon-ml-toolbox
+```
+**Note:** This version is outdated or broken due to dependency incompatibilities. Use PyPi instead.
+
+## Modular Installation
+
+### 📦 Core Machine Learning Toolbox [ML]
+
+Installs a comprehensive set of tools for typical data science workflows, including data manipulation, modeling, and evaluation. PyTorch is required.
+
+```Bash
+pip install "dragon-ml-toolbox[ML]"
+```
+
+To install the standard CPU-only versions of Torch and Torchvision:
+
+```Bash
+pip install "dragon-ml-toolbox[pytorch]"
+```
+
+⚠️ To make use of GPU acceleration (highly recommended), follow the official instructions: [PyTorch website](https://pytorch.org/get-started/locally/)
+
+#### Modules:
+
+```bash
+custom_logger
+data_exploration
+datasetmaster
+ensemble_learning
+ensemble_inference
+ETL_engineering
+ML_callbacks
+ML_evaluation
+ML_trainer
+ML_inference
+path_manager
+PSO_optimization
+RNN_forecast
+utilities
+```
+
+### 🔬 MICE Imputation and Variance Inflation Factor [mice]
+
+⚠️ Important: This group has strict version requirements. It is highly recommended to install this group in a separate virtual environment.
+
+```Bash
+pip install "dragon-ml-toolbox[mice]"
+```
+
+#### Modules:
+
+```bash
+custom_logger
+MICE_imputation
+VIF_factor
+path_manager
+utilities
+```
+
+### 📋 Excel File Handling [excel]
+
+Installs dependencies required to process and handle .xlsx or .xls files.
+
+```Bash
+pip install "dragon-ml-toolbox[excel]"
+```
+
+#### Modules:
+
+```bash
+custom_logger
+handle_excel
+path_manager
+```
+
+### 🎰 GUI for Boosting Algorithms (XGBoost, LightGBM) [gui-boost]
+
+For GUIs that include plotting functionality, you must also install the [plot] extra.
+
+```Bash
+pip install "dragon-ml-toolbox[gui-boost]"
+```
+
+```Bash
+pip install "dragon-ml-toolbox[gui-boost,plot]"
+```
+
+#### Modules:
+
+```bash
+GUI_tools
+ensemble_inference
+path_manager
+```
+
+### 🤖 GUI for PyTorch Models [gui-torch]
+
+For GUIs that include plotting functionality, you must also install the [plot] extra.
+
+```Bash
+pip install "dragon-ml-toolbox[gui-torch]"
+```
+
+```Bash
+pip install "dragon-ml-toolbox[gui-torch,plot]"
+```
+
+#### Modules:
+
+```bash
+GUI_tools
+ML_inference
+path_manager
+```
+
+## Usage
+
+After installation, import modules like this:
+
+```python
+from ml_tools.utilities import serialize_object, deserialize_object
+from ml_tools.custom_logger import custom_logger
+```

dragon_ml_toolbox-4.0.0.dist-info/RECORD

@@ -0,0 +1,29 @@
+dragon_ml_toolbox-4.0.0.dist-info/licenses/LICENSE,sha256=2uUFNy7D0TLgHim1K5s3DIJ4q_KvxEXVilnU20cWliY,1066
+dragon_ml_toolbox-4.0.0.dist-info/licenses/LICENSE-THIRD-PARTY.md,sha256=lY4_rJPnLnMu7YBQaY-_iz1JRDcLdQzNCyeLAF1glJY,1837
+ml_tools/ETL_engineering.py,sha256=m_IY-4hSp5X5TfJbWQ-MJNRxkxl4fcsxOnsivMs8tiM,39506
+ml_tools/GUI_tools.py,sha256=n4ZZ5kEjwK5rkOCFJE41HeLFfjhpJVLUSzk9Kd9Kr_0,45410
+ml_tools/MICE_imputation.py,sha256=b6ZTs8RedXFifOpuMCzr68xM16mCBVh1Ua6kcGfiVtg,11462
+ml_tools/ML_callbacks.py,sha256=0a-Rbr0Xp_B1FNopOKBBmuJ4MqazS5JgDiT7wx1dHvE,13161
+ml_tools/ML_evaluation.py,sha256=4dVqe6JF1Ukmk1sAcY8E5EG1oB1_oy2HXE5OT-pZwCs,10273
+ml_tools/ML_inference.py,sha256=Fh-X2UQn3AznWBjf-7iPSxwE-EzkGQm1VEIRUAkURmE,5336
+ml_tools/ML_trainer.py,sha256=dJjMfCEEM07Txy9KEH-2srZ3CZUa4lFWTJhpNWQ4Ndk,14974
+ml_tools/PSO_optimization.py,sha256=z8zPyoMtE-Vl5LcB24ZtNNTcEw9kVHftkV1VkV5pLD8,24662
+ml_tools/RNN_forecast.py,sha256=2CyjBLSYYc3xLHxwLXUmP5Qv8AmV1OB_EndETNX1IBk,1956
+ml_tools/VIF_factor.py,sha256=2nUMupfUoogf8o6ghoFZk_OwWhFXU0R3C9Gj0HOlI14,10415
+ml_tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ml_tools/_logger.py,sha256=TpgYguxO-CWYqqgLW0tqFjtwZ58PE_W2OCfWNGZr0n0,1175
+ml_tools/_pytorch_models.py,sha256=ewPPsTHgmRPzMMWwObZOdH1vxm2Ij2VWZP38NC6zSH4,10135
+ml_tools/_script_info.py,sha256=21r83LV3RubsNZ_RTEUON6RbDf7Mh4_udweNcvdF_Fk,212
+ml_tools/custom_logger.py,sha256=a3ywSCQT7j5ypR-usnKh2l861d_aVJ93ZRVqxrHsBBw,4112
+ml_tools/data_exploration.py,sha256=rJhvxUqVbEuB_7HG-PfLH3vaA7hrZEtbVHg9QO9VS4A,22837
+ml_tools/datasetmaster.py,sha256=_tNC2v98eCQGr3nMW_EFs83TRgRme8Uc7ttg1vosmQU,30106
+ml_tools/ensemble_inference.py,sha256=0SNX3YAz5bpvtwYmqEwqyWeIJP2Pb-v-bemENRSO7qg,9426
+ml_tools/ensemble_learning.py,sha256=Zi1oy6G2FWnTI5hBwjlexwF3JKALFS2FN6F8HAlVi_s,35391
+ml_tools/handle_excel.py,sha256=J9iwIqMZemoxK49J5osSwp9Ge0h9YTKyYGbOm53hcno,13007
+ml_tools/keys.py,sha256=kK9UF-hek2VcPGFILCKl5geoN6flmMOu7IzhdEA6z5Y,1068
+ml_tools/path_manager.py,sha256=ElDa25bntANujTjY7xN4ZfCDiZp-9Ud3x0aJSJptZBY,13419
+ml_tools/utilities.py,sha256=mz-M351DzxWxnYVcLX-7ZQ6c-RGoCV9g4VTS9Qif2Es,18348
+dragon_ml_toolbox-4.0.0.dist-info/METADATA,sha256=tGUq_S7xEHszYM0vUifYUzV4dmV_TQ9W8Ja4sJOEPTs,5994
+dragon_ml_toolbox-4.0.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dragon_ml_toolbox-4.0.0.dist-info/top_level.txt,sha256=wm-oxax3ciyez6VoO4zsFd-gSok2VipYXnbg3TH9PtU,9
+dragon_ml_toolbox-4.0.0.dist-info/RECORD,,

ml_tools/ETL_engineering.py

@@ -1,8 +1,8 @@
 import polars as pl
 import re
 from typing import Literal, Union, Optional, Any, Callable, List, Dict, Tuple
-from .utilities import _script_info
-from .logger import _LOGGER
+from ._script_info import _script_info
+from ._logger import _LOGGER
 
 
 __all__ = [

ml_tools/GUI_tools.py

@@ -4,9 +4,9 @@ import traceback
 import FreeSimpleGUI as sg
 from functools import wraps
 from typing import Any, Dict, Tuple, List, Literal, Union, Optional, Callable
-from .utilities import _script_info
+from ._script_info import _script_info
 import numpy as np
-from .logger import _LOGGER
+from ._logger import _LOGGER
 from .keys import _OneHotOtherPlaceholder
 
 
@@ -191,13 +191,14 @@ class GUIFactory:
         }
         return sg.Button(text.title(), key=key, **style_args)
 
-    def make_frame(self, title: str, layout: List[List[Union[sg.Element, sg.Column]]], **kwargs) -> sg.Frame:
+    def make_frame(self, title: str, layout: List[List[Union[sg.Element, sg.Column]]], center_layout: bool = False, **kwargs) -> sg.Frame:
         """
         Creates a styled frame around a given layout.
 
         Args:
             title (str): The title displayed on the frame's border.
             layout (list): The layout to enclose within the frame.
+            center_layout (bool): If True, the content within the frame will be horizontally centered.
             **kwargs: Override default styles or add other sg.Frame parameters
                       (e.g., `title_color='red'`, `relief=sg.RELIEF_SUNKEN`).
         """
@@ -210,6 +211,10 @@ class GUIFactory:
             "background_color": sg.theme_background_color(),
             **kwargs
         }
+        
+        if center_layout:
+            style_args["element_justification"] = 'center'
+        
         return sg.Frame(title, layout, **style_args)
 
     # --- General-Purpose Layout Generators ---
@@ -218,7 +223,8 @@ class GUIFactory:
         data_dict: Dict[str, Union[Tuple[Union[int,float,None], Union[int,float,None]],List[Union[int,float,None]]]],
         is_target: bool = False,
         layout_mode: Literal["grid", "row"] = 'grid',
-        number_columns: int = 5
+        number_columns: int = 5,
+        center_layout: bool = True
     ) -> List[List[sg.Column]]:
         """
         Generates a layout for continuous features or targets.
@@ -228,6 +234,7 @@ class GUIFactory:
             is_target (bool): If True, creates disabled inputs for displaying results.
             layout_mode (str): 'grid' for a multi-row grid layout, or 'row' for a single horizontal row.
             number_columns (int): Number of columns when `layout_mode` is 'grid'.
+            center_layout (bool): If True, the entire grid will be horizontally centered.
 
         Returns:
             A list of lists of sg.Column elements, ready to be used in a window layout.
@@ -264,7 +271,7 @@ class GUIFactory:
                 layout = [[label], [element]]
             else:
                 range_font = (cfg.fonts.font_family, cfg.fonts.range_size) # type: ignore
-                range_text = sg.Text(f"Range: {val_min}-{val_max}", font=range_font, background_color=bg_color) # type: ignore
+                range_text = sg.Text(f"Range: {val_min} - {val_max}", font=range_font, background_color=bg_color) # type: ignore
                 layout = [[label], [element], [range_text]]
             
             # each feature is wrapped as a column element
@@ -275,13 +282,14 @@ class GUIFactory:
             return [all_feature_layouts] # A single row containing all features
         
         # Default to 'grid' layout: delegate to the helper method
-        return self._build_grid_layout(all_feature_layouts, number_columns, bg_color) # type: ignore
+        return self._build_grid_layout(all_feature_layouts, number_columns, bg_color, center_layout) # type: ignore
 
     def generate_combo_layout(
         self,
         data_dict: Dict[str, Union[List[Any],Tuple[Any,...]]],
         layout_mode: Literal["grid", "row"] = 'grid',
-        number_columns: int = 5
+        number_columns: int = 5,
+        center_layout: bool = True
     ) -> List[List[sg.Column]]:
         """
         Generates a layout for categorical or binary features using Combo boxes.
@@ -290,6 +298,7 @@ class GUIFactory:
             data_dict (dict): Keys are feature names, values are lists of options.
             layout_mode (str): 'grid' for a multi-row grid layout, or 'row' for a single horizontal row.
             number_columns (int): Number of columns when `layout_mode` is 'grid'.
+            center_layout (bool): If True, the entire grid will be horizontally centered.
 
         Returns:
             A list of lists of sg.Column elements, ready to be used in a window layout.
@@ -315,13 +324,14 @@ class GUIFactory:
             return [all_feature_layouts] # A single row containing all features
             
         # Default to 'grid' layout: delegate to the helper method
-        return self._build_grid_layout(all_feature_layouts, number_columns, bg_color) # type: ignore
+        return self._build_grid_layout(all_feature_layouts, number_columns, bg_color, center_layout) # type: ignore
     
     def generate_multiselect_layout(
         self,
         data_dict: Dict[str, Union[List[Any], Tuple[Any, ...]]],
         layout_mode: Literal["grid", "row"] = 'grid',
-        number_columns: int = 5
+        number_columns: int = 5,
+        center_layout: bool = True
     ) -> List[List[sg.Column]]:
         """
         Generates a layout for features using Listbox elements for multiple selections.
@@ -333,6 +343,7 @@ class GUIFactory:
             data_dict (dict): Keys are feature names, values are lists of options.
             layout_mode (str): 'grid' for a multi-row grid layout, or 'row' for a single horizontal row.
             number_columns (int): Number of columns when `layout_mode` is 'grid'.
+            center_layout (bool): If True, the entire grid will be horizontally centered.
 
         Returns:
             A list of lists of sg.Column elements, ready to be used in a window layout.
@@ -366,7 +377,7 @@ class GUIFactory:
             return [all_feature_layouts]  # A single row containing all features
 
         # Default to 'grid' layout: delegate to the helper method
-        return self._build_grid_layout(all_feature_layouts, number_columns, bg_color) # type: ignore
+        return self._build_grid_layout(all_feature_layouts, number_columns, bg_color, center_layout) # type: ignore
 
     # --- Window Creation ---
     def create_window(self, title: str, layout: List[List[sg.Element]], **kwargs) -> sg.Window:
@@ -396,7 +407,7 @@ class GUIFactory:
         
         return window
     
-    def _build_grid_layout(self, all_feature_layouts: List[sg.Column], num_columns: int, bg_color: str) -> List[List[sg.Column]]:
+    def _build_grid_layout(self, all_feature_layouts: List[sg.Column], num_columns: int, bg_color: str, center_layout: bool = True) -> List[List[sg.Column]]:
         """
         Private helper to distribute feature layouts vertically into a grid of columns.
         """
@@ -412,7 +423,12 @@ class GUIFactory:
         gui_columns = [sg.Column([[c] for c in col], background_color=bg_color) for col in final_columns]
 
         # Return a single row containing all the generated vertical columns
-        return [gui_columns]
+        if center_layout:
+            # Return a single row containing the columns, centered with Push elements.
+            return [[sg.Push()] + gui_columns + [sg.Push()]] # type: ignore
+        else:
+            # Return a single row containing just the columns.
+            return [gui_columns]
 
 
 # --- Exception Handling Decorator ---

ml_tools/MICE_imputation.py

@@ -3,11 +3,12 @@ import miceforest as mf
 from pathlib import Path
 import matplotlib.pyplot as plt
 import numpy as np
-from .utilities import load_dataframe, list_csv_paths, sanitize_filename, _script_info, merge_dataframes, save_dataframe, threshold_binary_values, make_fullpath
+from .utilities import load_dataframe, merge_dataframes, save_dataframe, threshold_binary_values
+from .path_manager import sanitize_filename, make_fullpath, list_csv_paths
 from plotnine import ggplot, labs, theme, element_blank # type: ignore
 from typing import Optional, Union
-from .logger import _LOGGER
-
+from ._logger import _LOGGER
+from ._script_info import _script_info
 
 __all__ = [
     "apply_mice",

ml_tools/ML_callbacks.py

@@ -1,10 +1,11 @@
 import numpy as np
 import torch
 from tqdm.auto import tqdm
-from .utilities import make_fullpath
+from .path_manager import make_fullpath
 from .keys import LogKeys
-from .logger import _LOGGER
+from ._logger import _LOGGER
 from typing import Optional
+from ._script_info import _script_info
 
 
 __all__ = [
@@ -270,7 +271,7 @@ class ModelCheckpoint(Callback):
             self.last_best_filepath = new_filepath
 
     def _save_rolling_checkpoints(self, epoch, logs):
-        """Saves the latest model and keeps only the last 5."""
+        """Saves the latest model and keeps only the most recent ones."""
         filename = f"epoch_{epoch}.pth"
         filepath = self.save_dir / filename
         
@@ -334,4 +335,7 @@ class LRScheduler(Callback):
         if current_lr != self.previous_lr:
             _LOGGER.info(f"Epoch {epoch}: Learning rate changed to {current_lr:.6f}")
             self.previous_lr = current_lr
-            
+
+
+def info():
+    _script_info(__all__)

ml_tools/ML_evaluation.py

@@ -14,9 +14,10 @@ from sklearn.metrics import (
 import torch
 import shap
 from pathlib import Path
-from .utilities import make_fullpath
-from .logger import _LOGGER
+from .path_manager import make_fullpath
+from ._logger import _LOGGER
 from typing import Union, Optional
+from ._script_info import _script_info
 
 
 __all__ = [
@@ -62,7 +63,7 @@ def plot_losses(history: dict, save_dir: Optional[Union[str, Path]] = None):
     plt.tight_layout()
     
     if save_dir:
-        save_dir_path = make_fullpath(save_dir, make=True)
+        save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
         save_path = save_dir_path / "loss_plot.svg"
         plt.savefig(save_path)
         _LOGGER.info(f"📉 Loss plot saved as '{save_path.name}'")
@@ -88,7 +89,7 @@ def classification_metrics(y_true: np.ndarray, y_pred: np.ndarray, y_prob: Optio
     print(report)
     
     if save_dir:
-        save_dir_path = make_fullpath(save_dir, make=True)
+        save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
         # Save text report
         report_path = save_dir_path / "classification_report.txt"
         report_path.write_text(report, encoding="utf-8")
@@ -158,7 +159,7 @@ def regression_metrics(y_true: np.ndarray, y_pred: np.ndarray, save_dir: Optiona
     print(report_string)
 
     if save_dir:
-        save_dir_path = make_fullpath(save_dir, make=True)
+        save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
         # Save text report
         report_path = save_dir_path / "regression_report.txt"
         report_path.write_text(report_string)
@@ -220,7 +221,7 @@ def shap_summary_plot(model, background_data: torch.Tensor, instances_to_explain
         _LOGGER.info("Using SHAP values for the positive class (class 1) for plots.")
 
     if save_dir:
-        save_dir_path = make_fullpath(save_dir, make=True)
+        save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
         # Save Bar Plot
         bar_path = save_dir_path / "shap_bar_plot.svg"
         shap.summary_plot(shap_values_for_plot, instances_to_explain, feature_names=feature_names, plot_type="bar", show=False)
@@ -253,3 +254,7 @@ def shap_summary_plot(model, background_data: torch.Tensor, instances_to_explain
     else:
         _LOGGER.info("No save directory provided. Displaying SHAP dot plot.")
         shap.summary_plot(shap_values_for_plot, instances_to_explain, feature_names=feature_names, plot_type="dot")
+
+
+def info():
+    _script_info(__all__)

ml_tools/ML_inference.py

@@ -0,0 +1,131 @@
+import torch
+from torch import nn
+import numpy as np
+from pathlib import Path
+from typing import Union, Literal, Dict, Any
+
+from ._script_info import _script_info
+from ._logger import _LOGGER
+from .path_manager import make_fullpath
+from .keys import PyTorchInferenceKeys
+
+__all__ = [
+    "PyTorchInferenceHandler"
+]
+
+class PyTorchInferenceHandler:
+    """
+    Handles loading a PyTorch model's state dictionary and performing inference
+    for either regression or classification tasks.
+    """
+    def __init__(self,
+                 model: nn.Module,
+                 state_dict: Union[str, Path],
+                 task: Literal["classification", "regression"],
+                 device: str = 'cpu'):
+        """
+        Initializes the handler by loading a model's state_dict.
+
+        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'.
+            device (str): The device to run inference on ('cpu', 'cuda', 'mps').
+        """
+        self.model = model
+        self.task = task
+        self.device = self._validate_device(device)
+
+        model_p = make_fullpath(state_dict, enforce="file")
+
+        try:
+            # Load the state dictionary and apply it to the model structure
+            self.model.load_state_dict(torch.load(model_p, map_location=self.device))
+            self.model.to(self.device)
+            self.model.eval()  # Set the model to evaluation mode
+            _LOGGER.info(f"✅ Model state loaded from '{model_p.name}' and set to evaluation mode.")
+        except Exception as e:
+            _LOGGER.error(f"❌ Failed to load model state from '{model_p}': {e}")
+            raise
+
+    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 _preprocess_input(self, features: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
+        """Converts input to a torch.Tensor and moves it to the correct device."""
+        if isinstance(features, np.ndarray):
+            features = torch.from_numpy(features).float()
+        
+        # Ensure tensor is on the correct device
+        return features.to(self.device)
+
+    def predict(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, Any]:
+        """
+        Predicts on a single feature vector.
+
+        Args:
+            features (np.ndarray | torch.Tensor): A 1D or 2D array/tensor for a single sample.
+
+        Returns:
+            Dict[str, Any]: A dictionary containing the prediction.
+                - For regression: {'predictions': float}
+                - For classification: {'labels': int, 'probabilities': np.ndarray}
+        """
+        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.")
+
+        results_batch = self.predict_batch(features)
+
+        # Extract the single result from the batch
+        if self.task == "regression":
+            return {PyTorchInferenceKeys.PREDICTIONS: results_batch[PyTorchInferenceKeys.PREDICTIONS].item()}
+        else: # classification
+            return {
+                PyTorchInferenceKeys.LABELS: results_batch[PyTorchInferenceKeys.LABELS].item(),
+                PyTorchInferenceKeys.PROBABILITIES: results_batch[PyTorchInferenceKeys.PROBABILITIES][0]
+            }
+
+    def predict_batch(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, Any]:
+        """
+        Predicts on a batch of feature vectors.
+
+        Args:
+            features (np.ndarray | torch.Tensor): A 2D array/tensor where each row is a sample.
+
+        Returns:
+            Dict[str, Any]: A dictionary containing the predictions.
+                - For regression: {'predictions': np.ndarray}
+                - For classification: {'labels': np.ndarray, 'probabilities': np.ndarray}
+        """
+        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).cpu()
+
+            if self.task == "classification":
+                probs = nn.functional.softmax(output, dim=1)
+                labels = torch.argmax(probs, dim=1)
+                return {
+                    PyTorchInferenceKeys.LABELS: labels.numpy(),
+                    PyTorchInferenceKeys.PROBABILITIES: probs.numpy()
+                }
+            else:  # regression
+                return {PyTorchInferenceKeys.PREDICTIONS: output.numpy()}
+
+
+def info():
+    _script_info(__all__)