dragon-ml-toolbox 6.3.0__tar.gz → 6.4.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of dragon-ml-toolbox might be problematic. Click here for more details.

Files changed (39) hide show
  1. {dragon_ml_toolbox-6.3.0/dragon_ml_toolbox.egg-info → dragon_ml_toolbox-6.4.0}/PKG-INFO +2 -1
  2. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/README.md +1 -0
  3. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0/dragon_ml_toolbox.egg-info}/PKG-INFO +2 -1
  4. dragon_ml_toolbox-6.4.0/ml_tools/ML_inference.py +287 -0
  5. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/ML_models.py +1 -1
  6. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/pyproject.toml +1 -1
  7. dragon_ml_toolbox-6.3.0/ml_tools/ML_inference.py +0 -140
  8. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/LICENSE +0 -0
  9. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/LICENSE-THIRD-PARTY.md +0 -0
  10. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/dragon_ml_toolbox.egg-info/SOURCES.txt +0 -0
  11. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/dragon_ml_toolbox.egg-info/dependency_links.txt +0 -0
  12. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/dragon_ml_toolbox.egg-info/requires.txt +0 -0
  13. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/dragon_ml_toolbox.egg-info/top_level.txt +0 -0
  14. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/ETL_engineering.py +0 -0
  15. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/GUI_tools.py +0 -0
  16. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/MICE_imputation.py +0 -0
  17. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/ML_callbacks.py +0 -0
  18. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/ML_datasetmaster.py +0 -0
  19. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/ML_evaluation.py +0 -0
  20. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/ML_optimization.py +0 -0
  21. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/ML_trainer.py +0 -0
  22. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/PSO_optimization.py +0 -0
  23. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/RNN_forecast.py +0 -0
  24. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/SQL.py +0 -0
  25. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/VIF_factor.py +0 -0
  26. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/__init__.py +0 -0
  27. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/_logger.py +0 -0
  28. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/_script_info.py +0 -0
  29. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/custom_logger.py +0 -0
  30. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/data_exploration.py +0 -0
  31. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/ensemble_evaluation.py +0 -0
  32. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/ensemble_inference.py +0 -0
  33. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/ensemble_learning.py +0 -0
  34. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/handle_excel.py +0 -0
  35. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/keys.py +0 -0
  36. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/optimization_tools.py +0 -0
  37. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/path_manager.py +0 -0
  38. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/ml_tools/utilities.py +0 -0
  39. {dragon_ml_toolbox-6.3.0 → dragon_ml_toolbox-6.4.0}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dragon-ml-toolbox
3
- Version: 6.3.0
3
+ Version: 6.4.0
4
4
  Summary: A collection of tools for data science and machine learning projects.
5
5
  Author-email: Karl Loza <luigiloza@gmail.com>
6
6
  License-Expression: MIT
@@ -240,6 +240,7 @@ pip install "dragon-ml-toolbox[gui-torch,plot]"
240
240
  ```Bash
241
241
  custom_logger
242
242
  GUI_tools
243
+ ML_models
243
244
  ML_inference
244
245
  path_manager
245
246
  ```
@@ -159,6 +159,7 @@ pip install "dragon-ml-toolbox[gui-torch,plot]"
159
159
  ```Bash
160
160
  custom_logger
161
161
  GUI_tools
162
+ ML_models
162
163
  ML_inference
163
164
  path_manager
164
165
  ```
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dragon-ml-toolbox
3
- Version: 6.3.0
3
+ Version: 6.4.0
4
4
  Summary: A collection of tools for data science and machine learning projects.
5
5
  Author-email: Karl Loza <luigiloza@gmail.com>
6
6
  License-Expression: MIT
@@ -240,6 +240,7 @@ pip install "dragon-ml-toolbox[gui-torch,plot]"
240
240
  ```Bash
241
241
  custom_logger
242
242
  GUI_tools
243
+ ML_models
243
244
  ML_inference
244
245
  path_manager
245
246
  ```
@@ -0,0 +1,287 @@
1
+ import torch
2
+ from torch import nn
3
+ import numpy as np
4
+ from pathlib import Path
5
+ from typing import Union, Literal, Dict, Any, Optional
6
+
7
+ from ._script_info import _script_info
8
+ from ._logger import _LOGGER
9
+ from .path_manager import make_fullpath
10
+ from .keys import PyTorchInferenceKeys
11
+
12
+ __all__ = [
13
+ "PyTorchInferenceHandler",
14
+ "multi_inference_regression",
15
+ "multi_inference_classification"
16
+ ]
17
+
18
+ class PyTorchInferenceHandler:
19
+ """
20
+ Handles loading a PyTorch model's state dictionary and performing inference
21
+ for either regression or classification tasks.
22
+ """
23
+ def __init__(self,
24
+ model: nn.Module,
25
+ state_dict: Union[str, Path],
26
+ task: Literal["classification", "regression"],
27
+ device: str = 'cpu',
28
+ target_id: Optional[str]=None):
29
+ """
30
+ Initializes the handler by loading a model's state_dict.
31
+
32
+ Args:
33
+ model (nn.Module): An instantiated PyTorch model with the correct architecture.
34
+ state_dict (str | Path): The path to the saved .pth model state_dict file.
35
+ task (str): The type of task, 'regression' or 'classification'.
36
+ device (str): The device to run inference on ('cpu', 'cuda', 'mps').
37
+ target_id (str | None): Target name as used in the training set.
38
+ """
39
+ self.model = model
40
+ self.task = task
41
+ self.device = self._validate_device(device)
42
+ self.target_id = target_id
43
+
44
+ model_p = make_fullpath(state_dict, enforce="file")
45
+
46
+ try:
47
+ # Load the state dictionary and apply it to the model structure
48
+ self.model.load_state_dict(torch.load(model_p, map_location=self.device))
49
+ self.model.to(self.device)
50
+ self.model.eval() # Set the model to evaluation mode
51
+ _LOGGER.info(f"✅ Model state loaded from '{model_p.name}' and set to evaluation mode.")
52
+ except Exception as e:
53
+ _LOGGER.error(f"❌ Failed to load model state from '{model_p}': {e}")
54
+ raise
55
+
56
+ def _validate_device(self, device: str) -> torch.device:
57
+ """Validates the selected device and returns a torch.device object."""
58
+ device_lower = device.lower()
59
+ if "cuda" in device_lower and not torch.cuda.is_available():
60
+ _LOGGER.warning("⚠️ CUDA not available, switching to CPU.")
61
+ device_lower = "cpu"
62
+ elif device_lower == "mps" and not torch.backends.mps.is_available():
63
+ _LOGGER.warning("⚠️ Apple Metal Performance Shaders (MPS) not available, switching to CPU.")
64
+ device_lower = "cpu"
65
+ return torch.device(device_lower)
66
+
67
+ def _preprocess_input(self, features: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
68
+ """Converts input to a torch.Tensor and moves it to the correct device."""
69
+ if isinstance(features, np.ndarray):
70
+ features = torch.from_numpy(features).float()
71
+
72
+ # Ensure tensor is on the correct device
73
+ return features.to(self.device)
74
+
75
+ def predict_batch(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
76
+ """
77
+ Core batch prediction method. Returns results as PyTorch tensors on the model's device.
78
+ """
79
+ if features.ndim != 2:
80
+ raise ValueError("Input for batch prediction must be a 2D array or tensor.")
81
+
82
+ input_tensor = self._preprocess_input(features)
83
+
84
+ with torch.no_grad():
85
+ # Output tensor remains on the model's device (e.g., 'mps' or 'cuda')
86
+ output = self.model(input_tensor)
87
+
88
+ if self.task == "classification":
89
+ probs = nn.functional.softmax(output, dim=1)
90
+ labels = torch.argmax(probs, dim=1)
91
+ return {
92
+ PyTorchInferenceKeys.LABELS: labels,
93
+ PyTorchInferenceKeys.PROBABILITIES: probs
94
+ }
95
+ else: # regression
96
+ return {PyTorchInferenceKeys.PREDICTIONS: output}
97
+
98
+ def predict(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
99
+ """
100
+ Core single-sample prediction. Returns results as PyTorch tensors on the model's device.
101
+ """
102
+ if features.ndim == 1:
103
+ features = features.reshape(1, -1)
104
+
105
+ if features.shape[0] != 1:
106
+ raise ValueError("The predict() method is for a single sample. Use predict_batch() for multiple samples.")
107
+
108
+ batch_results = self.predict_batch(features)
109
+
110
+ single_results = {key: value[0] for key, value in batch_results.items()}
111
+ return single_results
112
+
113
+ # --- NumPy Convenience Wrappers (on CPU) ---
114
+
115
+ def predict_batch_numpy(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, np.ndarray]:
116
+ """
117
+ Convenience wrapper for predict_batch that returns NumPy arrays.
118
+ """
119
+ tensor_results = self.predict_batch(features)
120
+ # Move tensor to CPU before converting to NumPy
121
+ numpy_results = {key: value.cpu().numpy() for key, value in tensor_results.items()}
122
+ return numpy_results
123
+
124
+ def predict_numpy(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, Any]:
125
+ """
126
+ Convenience wrapper for predict that returns NumPy arrays or scalars.
127
+ """
128
+ tensor_results = self.predict(features)
129
+
130
+ if self.task == "regression":
131
+ # .item() implicitly moves to CPU
132
+ return {PyTorchInferenceKeys.PREDICTIONS: tensor_results[PyTorchInferenceKeys.PREDICTIONS].item()}
133
+ else: # classification
134
+ return {
135
+ PyTorchInferenceKeys.LABELS: tensor_results[PyTorchInferenceKeys.LABELS].item(),
136
+ # Move tensor to CPU before converting to NumPy
137
+ PyTorchInferenceKeys.PROBABILITIES: tensor_results[PyTorchInferenceKeys.PROBABILITIES].cpu().numpy()
138
+ }
139
+
140
+
141
+ def multi_inference_regression(handlers: list[PyTorchInferenceHandler],
142
+ feature_vector: Union[np.ndarray, torch.Tensor],
143
+ output: Literal["numpy","torch"]="numpy") -> dict[str,Any]:
144
+ """
145
+ Performs regression inference using multiple models on a single feature vector.
146
+
147
+ This function iterates through a list of PyTorchInferenceHandler objects,
148
+ each configured for a different regression target. It runs a prediction for
149
+ each handler using the same input feature vector and returns the results
150
+ in a dictionary.
151
+
152
+ The function adapts its behavior based on the input dimensions:
153
+ - 1D input: Returns a dictionary mapping target ID to a single value.
154
+ - 2D input: Returns a dictionary mapping target ID to a list of values.
155
+
156
+ Args:
157
+ handlers (list[PyTorchInferenceHandler]): A list of initialized inference
158
+ handlers. Each handler must have a unique `target_id` and be configured with `task="regression"`.
159
+ feature_vector (Union[np.ndarray, torch.Tensor]): An input sample (1D) or a batch of samples (2D) to be fed into each regression model.
160
+ output (Literal["numpy", "torch"], optional): The desired format for the output predictions.
161
+ - "numpy": Returns predictions as Python scalars or NumPy arrays.
162
+ - "torch": Returns predictions as PyTorch tensors.
163
+
164
+ Returns:
165
+ (dict[str, Any]): A dictionary mapping each handler's `target_id` to its
166
+ predicted regression values.
167
+
168
+ Raises:
169
+ AttributeError: If any handler in the list is missing a `target_id`.
170
+ ValueError: If any handler's `task` is not 'regression' or if the input `feature_vector` is not 1D or 2D.
171
+ """
172
+ # check batch dimension
173
+ is_single_sample = feature_vector.ndim == 1
174
+
175
+ # Reshape a 1D vector to a 2D batch of one for uniform processing.
176
+ if is_single_sample:
177
+ feature_vector = feature_vector.reshape(1, -1)
178
+
179
+ # Validate that the input is a 2D tensor.
180
+ if feature_vector.ndim != 2:
181
+ raise ValueError("Input feature_vector must be a 1D or 2D array/tensor.")
182
+
183
+ results: dict[str,Any] = dict()
184
+ for handler in handlers:
185
+ # validation
186
+ if handler.target_id is None:
187
+ raise AttributeError("All inference handlers must have a 'target_id' attribute.")
188
+ if handler.task != "regression":
189
+ raise ValueError(
190
+ f"Invalid task type: The handler for target_id '{handler.target_id}' "
191
+ f"is for '{handler.task}', but only 'regression' tasks are supported."
192
+ )
193
+ # inference
194
+ if output == "numpy":
195
+ result = handler.predict_batch_numpy(feature_vector)[PyTorchInferenceKeys.PREDICTIONS]
196
+ else: # torch
197
+ result = handler.predict_batch(feature_vector)[PyTorchInferenceKeys.PREDICTIONS]
198
+
199
+ # Unpack single results and update result dictionary
200
+ # If the original input was 1D, extract the single prediction from the array.
201
+ if is_single_sample:
202
+ results[handler.target_id] = result[0]
203
+ else:
204
+ results[handler.target_id] = result
205
+
206
+ return results
207
+
208
+
209
+ def multi_inference_classification(
210
+ handlers: list[PyTorchInferenceHandler],
211
+ feature_vector: Union[np.ndarray, torch.Tensor],
212
+ output: Literal["numpy","torch"]="numpy"
213
+ ) -> tuple[dict[str, Any], dict[str, Any]]:
214
+ """
215
+ Performs classification inference on a single sample or a batch.
216
+
217
+ This function iterates through a list of PyTorchInferenceHandler objects,
218
+ each configured for a different classification target. It returns two
219
+ dictionaries: one for the predicted labels and one for the probabilities.
220
+
221
+ The function adapts its behavior based on the input dimensions:
222
+ - 1D input: The dictionaries map target ID to a single label and a single probability array.
223
+ - 2D input: The dictionaries map target ID to an array of labels and an array of probability arrays.
224
+
225
+ Args:
226
+ handlers (list[PyTorchInferenceHandler]): A list of initialized inference handlers. Each must have a unique `target_id` and be configured
227
+ with `task="classification"`.
228
+ feature_vector (Union[np.ndarray, torch.Tensor]): An input sample (1D)
229
+ or a batch of samples (2D) for prediction.
230
+ output (Literal["numpy", "torch"], optional): The desired format for the
231
+ output predictions.
232
+
233
+ Returns:
234
+ (tuple[dict[str, Any], dict[str, Any]]): A tuple containing two dictionaries:
235
+ 1. A dictionary mapping `target_id` to the predicted label(s).
236
+ 2. A dictionary mapping `target_id` to the prediction probabilities.
237
+
238
+ Raises:
239
+ AttributeError: If any handler in the list is missing a `target_id`.
240
+ ValueError: If any handler's `task` is not 'classification' or if the input `feature_vector` is not 1D or 2D.
241
+ """
242
+ # Store if the original input was a single sample
243
+ is_single_sample = feature_vector.ndim == 1
244
+
245
+ # Reshape a 1D vector to a 2D batch of one for uniform processing
246
+ if is_single_sample:
247
+ feature_vector = feature_vector.reshape(1, -1)
248
+
249
+ if feature_vector.ndim != 2:
250
+ raise ValueError("Input feature_vector must be a 1D or 2D array/tensor.")
251
+
252
+ # Initialize two dictionaries for results
253
+ labels_results: dict[str, Any] = dict()
254
+ probs_results: dict[str, Any] = dict()
255
+
256
+ for handler in handlers:
257
+ # Validation
258
+ if handler.target_id is None:
259
+ raise AttributeError("All inference handlers must have a 'target_id' attribute.")
260
+ if handler.task != "classification":
261
+ raise ValueError(
262
+ f"Invalid task type: The handler for target_id '{handler.target_id}' "
263
+ f"is for '{handler.task}', but this function only supports 'classification'."
264
+ )
265
+
266
+ # Always use the batch method to get both labels and probabilities
267
+ if output == "numpy":
268
+ result = handler.predict_batch_numpy(feature_vector)
269
+ else: # torch
270
+ result = handler.predict_batch(feature_vector)
271
+
272
+ labels = result[PyTorchInferenceKeys.LABELS]
273
+ probabilities = result[PyTorchInferenceKeys.PROBABILITIES]
274
+
275
+ # If the original input was 1D, unpack the single result from the batch array
276
+ if is_single_sample:
277
+ labels_results[handler.target_id] = labels[0]
278
+ probs_results[handler.target_id] = probabilities[0]
279
+ else:
280
+ labels_results[handler.target_id] = labels
281
+ probs_results[handler.target_id] = probabilities
282
+
283
+ return labels_results, probs_results
284
+
285
+
286
+ def info():
287
+ _script_info(__all__)
@@ -223,7 +223,7 @@ def save_architecture(model: nn.Module, directory: Union[str, Path], verbose: bo
223
223
  json.dump(config, f, indent=4)
224
224
 
225
225
  if verbose:
226
- _LOGGER.info(f"✅ Architecture for '{model.__class__.__name__}' saved to '{path_dir}'")
226
+ _LOGGER.info(f"✅ Architecture for '{model.__class__.__name__}' saved to '{path_dir.name}'")
227
227
 
228
228
 
229
229
  def load_architecture(filepath: Union[str, Path], expected_model_class: type, verbose: bool=True) -> nn.Module:
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "dragon-ml-toolbox"
3
- version = "6.3.0"
3
+ version = "6.4.0"
4
4
  description = "A collection of tools for data science and machine learning projects."
5
5
  authors = [
6
6
  { name = "Karl Loza", email = "luigiloza@gmail.com" }
@@ -1,140 +0,0 @@
1
- import torch
2
- from torch import nn
3
- import numpy as np
4
- from pathlib import Path
5
- from typing import Union, Literal, Dict, Any, Optional
6
-
7
- from ._script_info import _script_info
8
- from ._logger import _LOGGER
9
- from .path_manager import make_fullpath
10
- from .keys import PyTorchInferenceKeys
11
-
12
- __all__ = [
13
- "PyTorchInferenceHandler"
14
- ]
15
-
16
- class PyTorchInferenceHandler:
17
- """
18
- Handles loading a PyTorch model's state dictionary and performing inference
19
- for either regression or classification tasks.
20
- """
21
- def __init__(self,
22
- model: nn.Module,
23
- state_dict: Union[str, Path],
24
- task: Literal["classification", "regression"],
25
- device: str = 'cpu',
26
- target_id: Optional[str]=None):
27
- """
28
- Initializes the handler by loading a model's state_dict.
29
-
30
- Args:
31
- model (nn.Module): An instantiated PyTorch model with the correct architecture.
32
- state_dict (str | Path): The path to the saved .pth model state_dict file.
33
- task (str): The type of task, 'regression' or 'classification'.
34
- device (str): The device to run inference on ('cpu', 'cuda', 'mps').
35
- target_id (str | None): Target name as used in the training set.
36
- """
37
- self.model = model
38
- self.task = task
39
- self.device = self._validate_device(device)
40
- self.target_id = target_id
41
-
42
- model_p = make_fullpath(state_dict, enforce="file")
43
-
44
- try:
45
- # Load the state dictionary and apply it to the model structure
46
- self.model.load_state_dict(torch.load(model_p, map_location=self.device))
47
- self.model.to(self.device)
48
- self.model.eval() # Set the model to evaluation mode
49
- _LOGGER.info(f"✅ Model state loaded from '{model_p.name}' and set to evaluation mode.")
50
- except Exception as e:
51
- _LOGGER.error(f"❌ Failed to load model state from '{model_p}': {e}")
52
- raise
53
-
54
- def _validate_device(self, device: str) -> torch.device:
55
- """Validates the selected device and returns a torch.device object."""
56
- device_lower = device.lower()
57
- if "cuda" in device_lower and not torch.cuda.is_available():
58
- _LOGGER.warning("⚠️ CUDA not available, switching to CPU.")
59
- device_lower = "cpu"
60
- elif device_lower == "mps" and not torch.backends.mps.is_available():
61
- _LOGGER.warning("⚠️ Apple Metal Performance Shaders (MPS) not available, switching to CPU.")
62
- device_lower = "cpu"
63
- return torch.device(device_lower)
64
-
65
- def _preprocess_input(self, features: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
66
- """Converts input to a torch.Tensor and moves it to the correct device."""
67
- if isinstance(features, np.ndarray):
68
- features = torch.from_numpy(features).float()
69
-
70
- # Ensure tensor is on the correct device
71
- return features.to(self.device)
72
-
73
- def predict_batch(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
74
- """
75
- Core batch prediction method. Returns results as PyTorch tensors on the model's device.
76
- """
77
- if features.ndim != 2:
78
- raise ValueError("Input for batch prediction must be a 2D array or tensor.")
79
-
80
- input_tensor = self._preprocess_input(features)
81
-
82
- with torch.no_grad():
83
- # Output tensor remains on the model's device (e.g., 'mps' or 'cuda')
84
- output = self.model(input_tensor)
85
-
86
- if self.task == "classification":
87
- probs = nn.functional.softmax(output, dim=1)
88
- labels = torch.argmax(probs, dim=1)
89
- return {
90
- PyTorchInferenceKeys.LABELS: labels,
91
- PyTorchInferenceKeys.PROBABILITIES: probs
92
- }
93
- else: # regression
94
- return {PyTorchInferenceKeys.PREDICTIONS: output}
95
-
96
- def predict(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, torch.Tensor]:
97
- """
98
- Core single-sample prediction. Returns results as PyTorch tensors on the model's device.
99
- """
100
- if features.ndim == 1:
101
- features = features.reshape(1, -1)
102
-
103
- if features.shape[0] != 1:
104
- raise ValueError("The predict() method is for a single sample. Use predict_batch() for multiple samples.")
105
-
106
- batch_results = self.predict_batch(features)
107
-
108
- single_results = {key: value[0] for key, value in batch_results.items()}
109
- return single_results
110
-
111
- # --- NumPy Convenience Wrappers (on CPU) ---
112
-
113
- def predict_batch_numpy(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, np.ndarray]:
114
- """
115
- Convenience wrapper for predict_batch that returns NumPy arrays.
116
- """
117
- tensor_results = self.predict_batch(features)
118
- # Move tensor to CPU before converting to NumPy
119
- numpy_results = {key: value.cpu().numpy() for key, value in tensor_results.items()}
120
- return numpy_results
121
-
122
- def predict_numpy(self, features: Union[np.ndarray, torch.Tensor]) -> Dict[str, Any]:
123
- """
124
- Convenience wrapper for predict that returns NumPy arrays or scalars.
125
- """
126
- tensor_results = self.predict(features)
127
-
128
- if self.task == "regression":
129
- # .item() implicitly moves to CPU
130
- return {PyTorchInferenceKeys.PREDICTIONS: tensor_results[PyTorchInferenceKeys.PREDICTIONS].item()}
131
- else: # classification
132
- return {
133
- PyTorchInferenceKeys.LABELS: tensor_results[PyTorchInferenceKeys.LABELS].item(),
134
- # ✅ Move tensor to CPU before converting to NumPy
135
- PyTorchInferenceKeys.PROBABILITIES: tensor_results[PyTorchInferenceKeys.PROBABILITIES].cpu().numpy()
136
- }
137
-
138
-
139
- def info():
140
- _script_info(__all__)