dragon-ml-toolbox 14.8.0__py3-none-any.whl → 16.0.0__py3-none-any.whl

ml_tools/ML_callbacks.py

@@ -4,23 +4,22 @@ from tqdm.auto import tqdm
 from typing import Union, Literal, Optional
 from pathlib import Path
 
-from .path_manager import make_fullpath, sanitize_filename
-from .keys import PyTorchLogKeys, PyTorchCheckpointKeys
+from .path_manager import make_fullpath
+from ._keys import PyTorchLogKeys, PyTorchCheckpointKeys
 from ._logger import _LOGGER
 from ._script_info import _script_info
 
 
 __all__ = [
-    "Callback", 
     "History", 
     "TqdmProgressBar",
-    "EarlyStopping",  
-    "ModelCheckpoint",
-    "LRScheduler"
+    "DragonEarlyStopping",  
+    "DragonModelCheckpoint",
+    "DragonLRScheduler"
 ]
 
 
-class Callback:
+class _Callback:
     """
     Abstract base class used to build new callbacks.
     
@@ -60,7 +59,7 @@ class Callback:
         pass
 
 
-class History(Callback):
+class History(_Callback):
     """
     Callback that records events into a `history` dictionary.
     
@@ -79,7 +78,7 @@ class History(Callback):
             self.trainer.history.setdefault(k, []).append(v) # type: ignore
 
 
-class TqdmProgressBar(Callback):
+class TqdmProgressBar(_Callback):
     """Callback that provides a tqdm progress bar for training."""
     def __init__(self):
         self.epoch_bar = None
@@ -110,7 +109,7 @@ class TqdmProgressBar(Callback):
         self.epoch_bar.close() # type: ignore
 
 
-class EarlyStopping(Callback):
+class DragonEarlyStopping(_Callback):
     """
     Stop training when a monitored metric has stopped improving.
     """
@@ -187,11 +186,11 @@ class EarlyStopping(Callback):
                     _LOGGER.info(f"Epoch {epoch+1}: early stopping after {self.wait} epochs with no improvement.")
 
 
-class ModelCheckpoint(Callback):
+class DragonModelCheckpoint(_Callback):
     """
     Saves the model weights, optimizer state, LR scheduler state (if any), and epoch number to a directory with automated filename generation and rotation. 
     """
-    def __init__(self, save_dir: Union[str,Path], checkpoint_name: Optional[str]=None, monitor: str = PyTorchLogKeys.VAL_LOSS,
+    def __init__(self, save_dir: Union[str,Path], monitor: str = PyTorchLogKeys.VAL_LOSS,
                  save_best_only: bool = True, mode: Literal['auto', 'min', 'max']= 'auto', verbose: int = 0):
         """
         - If `save_best_only` is True, it saves the single best model, deleting the previous best. 
@@ -199,7 +198,6 @@ class ModelCheckpoint(Callback):
 
         Args:
             save_dir (str): Directory where checkpoint files will be saved.
-            checkpoint_name (str| None): If None, the filename will include the epoch and score.
             monitor (str): Metric to monitor.
             save_best_only (bool): If true, save only the best model.
             mode (str): One of {'auto', 'min', 'max'}.
@@ -215,9 +213,8 @@ class ModelCheckpoint(Callback):
         self.monitor = monitor
         self.save_best_only = save_best_only
         self.verbose = verbose
-        if checkpoint_name:
-            checkpoint_name = sanitize_filename(checkpoint_name)
-        self.checkpoint_name = checkpoint_name
+        self._latest_checkpoint_path = None
+        self._checkpoint_name = PyTorchCheckpointKeys.CHECKPOINT_NAME
 
         # State variables to be managed during training
         self.saved_checkpoints = []
@@ -261,10 +258,7 @@ class ModelCheckpoint(Callback):
             old_best_str = f"{self.best:.4f}" if self.best not in [np.inf, -np.inf] else "inf"
             
             # Create a descriptive filename
-            if self.checkpoint_name is None:
-                filename = f"epoch_{epoch}-{self.monitor}_{current:.4f}.pth"
-            else:
-                filename = f"epoch{epoch}_{self.checkpoint_name}.pth"
+            filename = f"epoch{epoch}_{self._checkpoint_name}_{current:.4f}.pth"
             new_filepath = self.save_dir / filename
             
             if self.verbose > 0:
@@ -279,6 +273,7 @@ class ModelCheckpoint(Callback):
                 PyTorchCheckpointKeys.MODEL_STATE: self.trainer.model.state_dict(), # type: ignore
                 PyTorchCheckpointKeys.OPTIMIZER_STATE: self.trainer.optimizer.state_dict(), # type: ignore
                 PyTorchCheckpointKeys.BEST_SCORE: self.best, 
+                PyTorchCheckpointKeys.HISTORY: self.trainer.history, # type: ignore
             }
             
             # Check for scheduler
@@ -287,6 +282,7 @@ class ModelCheckpoint(Callback):
             
             # Save the new best model
             torch.save(checkpoint_data, new_filepath)
+            self._latest_checkpoint_path = new_filepath
 
             # Delete the old best model file
             if self.last_best_filepath and self.last_best_filepath.exists():
@@ -298,10 +294,8 @@ class ModelCheckpoint(Callback):
     def _save_rolling_checkpoints(self, epoch, logs):
         """Saves the latest model and keeps only the most recent ones."""
         current = logs.get(self.monitor)
-        if self.checkpoint_name is None:
-            filename = f"epoch_{epoch}-{self.monitor}_{current:.4f}.pth"
-        else:
-            filename = f"epoch{epoch}_{self.checkpoint_name}.pth"
+
+        filename = f"epoch{epoch}_{self._checkpoint_name}_{current:.4f}.pth"
         filepath = self.save_dir / filename
         
         if self.verbose > 0:
@@ -313,12 +307,15 @@ class ModelCheckpoint(Callback):
             PyTorchCheckpointKeys.MODEL_STATE: self.trainer.model.state_dict(), # type: ignore
             PyTorchCheckpointKeys.OPTIMIZER_STATE: self.trainer.optimizer.state_dict(), # type: ignore
             PyTorchCheckpointKeys.BEST_SCORE: self.best, # Save the current best score
+            PyTorchCheckpointKeys.HISTORY: self.trainer.history, # type: ignore
         }
         
         if hasattr(self.trainer, 'scheduler') and self.trainer.scheduler is not None: # type: ignore
             checkpoint_data[PyTorchCheckpointKeys.SCHEDULER_STATE] = self.trainer.scheduler.state_dict() # type: ignore
         
         torch.save(checkpoint_data, filepath)
+        
+        self._latest_checkpoint_path = filepath
 
         self.saved_checkpoints.append(filepath)
 
@@ -330,8 +327,16 @@ class ModelCheckpoint(Callback):
                     _LOGGER.info(f"  -> Deleting old checkpoint: {file_to_delete.name}")
                 file_to_delete.unlink()
 
+    @property
+    def best_checkpoint_path(self):
+        if self._latest_checkpoint_path:
+            return self._latest_checkpoint_path
+        else:
+            _LOGGER.error("No checkpoint paths saved.")
+            raise ValueError()
+
 
-class LRScheduler(Callback):
+class DragonLRScheduler(_Callback):
     """
     Callback to manage a PyTorch learning rate scheduler.
     """
@@ -361,6 +366,8 @@ class LRScheduler(Callback):
 
     def on_epoch_end(self, epoch, logs=None):
         """Step the scheduler and log any change in learning rate."""
+        logs = logs or {}
+        
         # For schedulers that need a metric (e.g., val_loss)
         if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
             if self.monitor is None:
@@ -376,12 +383,22 @@ class LRScheduler(Callback):
         # For all other schedulers
         else:
             self.scheduler.step()
+            
+        # Get the current learning rate
+        current_lr = self.trainer.optimizer.param_groups[0]['lr'] # type: ignore
 
         # Log the change if the LR was updated
-        current_lr = self.trainer.optimizer.param_groups[0]['lr'] # type: ignore
         if current_lr != self.previous_lr:
             _LOGGER.info(f"Epoch {epoch}: Learning rate changed to {current_lr:.6f}")
             self.previous_lr = current_lr
+        
+        # --- Add LR to logs and history ---
+        # Add to the logs dict for any subsequent callbacks
+        logs[PyTorchLogKeys.LEARNING_RATE] = current_lr
+        
+        # Also add directly to the trainer's history dict
+        if hasattr(self.trainer, 'history'):
+            self.trainer.history.setdefault(PyTorchLogKeys.LEARNING_RATE, []).append(current_lr) # type: ignore
 
 
 def info():

ml_tools/ML_configuration.py

@@ -1,16 +1,20 @@
-from typing import Optional
+from typing import Optional, Union
 from ._script_info import _script_info
 
 
 __all__ = [
     "ClassificationMetricsFormat",
-    "MultiClassificationMetricsFormat"
+    "MultiClassificationMetricsFormat",
+    "RegressionMetricsFormat",
+    "SegmentationMetricsFormat",
+    "SequenceValueMetricsFormat",
+    "SequenceSequenceMetricsFormat"
 ]
 
 
 class ClassificationMetricsFormat:
     """
-    Optional configuration for classification tasks, use in the '.evaluate()' method of the MLTrainer.
+    Optional configuration for classification tasks.
     """
     def __init__(self, 
                  cmap: str="Blues",
@@ -65,10 +69,9 @@ class ClassificationMetricsFormat:
 
 class MultiClassificationMetricsFormat:
     """
-    Optional configuration for multi-label classification tasks, use in the '.evaluate()' method of the MLTrainer.
+    Optional configuration for multi-label classification tasks.
     """
     def __init__(self,
-                 threshold: float=0.5,
                  ROC_PR_line: str='darkorange',
                  cmap: str = "Blues",
                  font_size: int = 16) -> None:
@@ -76,10 +79,6 @@ class MultiClassificationMetricsFormat:
         Initializes the formatting configuration for multi-label classification metrics.
 
         Args:
-            threshold (float): The probability threshold (0.0 to 1.0) used
-                to convert sigmoid outputs into binary (0 or 1) predictions for
-                calculating the confusion matrix and overall metrics. Defaults to 0.5.
-            
             ROC_PR_line (str): The color name or hex code for the line plotted
                 on the ROC and Precision-Recall curves (one for each label). 
                 Defaults to 'darkorange'.
@@ -97,14 +96,12 @@ class MultiClassificationMetricsFormat:
         
         ## [Matplotlib Colormaps](https://matplotlib.org/stable/users/explain/colors/colormaps.html)    
         """
-        self.threshold = threshold
         self.cmap = cmap
         self.ROC_PR_line = ROC_PR_line
         self.font_size = font_size
         
     def __repr__(self) -> str:
         parts = [
-            f"threshold={self.threshold}",
             f"ROC_PR_line='{self.ROC_PR_line}'",
             f"cmap='{self.cmap}'",
             f"font_size={self.font_size}"
@@ -112,5 +109,201 @@ class MultiClassificationMetricsFormat:
         return f"MultiClassificationMetricsFormat({', '.join(parts)})"
 
 
+class RegressionMetricsFormat:
+    """
+    Optional configuration for single-target regression and multi-target regression tasks.
+    """
+    def __init__(self, 
+                 font_size: int=16,
+                 scatter_color: str='tab:blue',
+                 scatter_alpha: float=0.6,
+                 ideal_line_color: str='k',
+                 residual_line_color: str='red',
+                 hist_bins: Union[int, str] = 'auto') -> None:
+        """
+        Initializes the formatting configuration for regression metrics.
+
+        Args:
+            font_size (int): The base font size to apply to the plots. Defaults to 16.
+            scatter_color (str): Matplotlib color for the scatter plot points. Defaults to 'tab:blue'.
+                - Common color names: 'tab:blue', 'crimson', 'forestgreen', '#4682B4'
+            scatter_alpha (float): Alpha transparency for scatter plot points. Defaults to 0.6.
+            ideal_line_color (str): Matplotlib color for the 'ideal' y=x line in the 
+                True vs. Predicted plot. Defaults to 'k' (black).
+                - Common color names: 'k', 'red', 'darkgrey', '#FF6347'
+            residual_line_color (str): Matplotlib color for the y=0 line in the 
+                Residual plot. Defaults to 'red'.
+                - Common color names: 'red', 'blue', 'k', '#4682B4'
+            hist_bins (int | str): The number of bins for the residuals histogram. 
+                Defaults to 'auto' to use seaborn's automatic bin selection.
+                - Options: 'auto', 'sqrt', 10, 20
+        
+        <br>
+        
+        ## [Matplotlib Colors](https://matplotlib.org/stable/users/explain/colors/colors.html)
+        """
+        self.font_size = font_size
+        self.scatter_color = scatter_color
+        self.scatter_alpha = scatter_alpha
+        self.ideal_line_color = ideal_line_color
+        self.residual_line_color = residual_line_color
+        self.hist_bins = hist_bins
+        
+    def __repr__(self) -> str:
+        parts = [
+            f"font_size={self.font_size}",
+            f"scatter_color='{self.scatter_color}'",
+            f"scatter_alpha={self.scatter_alpha}",
+            f"ideal_line_color='{self.ideal_line_color}'",
+            f"residual_line_color='{self.residual_line_color}'",
+            f"hist_bins='{self.hist_bins}'"
+        ]
+        return f"RegressionMetricsFormat({', '.join(parts)})"
+
+
+class SegmentationMetricsFormat:
+    """
+    Optional configuration for segmentation tasks.
+    """
+    def __init__(self,
+                 heatmap_cmap: str = 'viridis',
+                 cm_cmap: str = "Blues",
+                 font_size: int = 16) -> None:
+        """
+        Initializes the formatting configuration for segmentation metrics.
+
+        Args:
+            heatmap_cmap (str): The matplotlib colormap name for the per-class
+                metrics heatmap. Defaults to "viridis".
+                - Sequential options: 'viridis', 'plasma', 'inferno', 'cividis'
+                - Diverging options: 'coolwarm', 'bwr', 'seismic'
+            cm_cmap (str): The matplotlib colormap name for the pixel-level
+                confusion matrix. Defaults to "Blues".
+                - Sequential options: 'Blues', 'Greens', 'Reds', 'Oranges'
+            font_size (int): The base font size to apply to the plots. Defaults to 16.
+        
+        <br>
+        
+        ## [Matplotlib Colormaps](https://matplotlib.org/stable/users/explain/colors/colormaps.html)
+        """
+        self.heatmap_cmap = heatmap_cmap
+        self.cm_cmap = cm_cmap
+        self.font_size = font_size
+        
+    def __repr__(self) -> str:
+        parts = [
+            f"heatmap_cmap='{self.heatmap_cmap}'",
+            f"cm_cmap='{self.cm_cmap}'",
+            f"font_size={self.font_size}"
+        ]
+        return f"SegmentationMetricsFormat({', '.join(parts)})"
+
+
+# Similar to regression configuration
+class SequenceValueMetricsFormat:
+    """
+    Optional configuration for sequence to value prediction tasks.
+    """
+    def __init__(self, 
+                 font_size: int=16,
+                 scatter_color: str='tab:blue',
+                 scatter_alpha: float=0.6,
+                 ideal_line_color: str='k',
+                 residual_line_color: str='red',
+                 hist_bins: Union[int, str] = 'auto') -> None:
+        """
+        Initializes the formatting configuration for sequence to value metrics.
+
+        Args:
+            font_size (int): The base font size to apply to the plots. Defaults to 16.
+            scatter_color (str): Matplotlib color for the scatter plot points. Defaults to 'tab:blue'.
+                - Common color names: 'tab:blue', 'crimson', 'forestgreen', '#4682B4'
+            scatter_alpha (float): Alpha transparency for scatter plot points. Defaults to 0.6.
+            ideal_line_color (str): Matplotlib color for the 'ideal' y=x line in the 
+                True vs. Predicted plot. Defaults to 'k' (black).
+                - Common color names: 'k', 'red', 'darkgrey', '#FF6347'
+            residual_line_color (str): Matplotlib color for the y=0 line in the 
+                Residual plot. Defaults to 'red'.
+                - Common color names: 'red', 'blue', 'k', '#4682B4'
+            hist_bins (int | str): The number of bins for the residuals histogram. 
+                Defaults to 'auto' to use seaborn's automatic bin selection.
+                - Options: 'auto', 'sqrt', 10, 20
+        
+        <br>
+        
+        ## [Matplotlib Colors](https://matplotlib.org/stable/users/explain/colors/colors.html)
+        """
+        self.font_size = font_size
+        self.scatter_color = scatter_color
+        self.scatter_alpha = scatter_alpha
+        self.ideal_line_color = ideal_line_color
+        self.residual_line_color = residual_line_color
+        self.hist_bins = hist_bins
+        
+    def __repr__(self) -> str:
+        parts = [
+            f"font_size={self.font_size}",
+            f"scatter_color='{self.scatter_color}'",
+            f"scatter_alpha={self.scatter_alpha}",
+            f"ideal_line_color='{self.ideal_line_color}'",
+            f"residual_line_color='{self.residual_line_color}'",
+            f"hist_bins='{self.hist_bins}'"
+        ]
+        return f"SequenceValueMetricsFormat({', '.join(parts)})"
+
+
+class SequenceSequenceMetricsFormat:
+    """
+    Optional configuration for sequence-to-sequence evaluation plots.
+    """
+    def __init__(self,
+                 font_size: int = 16,
+                 plot_figsize: tuple[int, int] = (10, 6),
+                 grid_style: str = '--',
+                 rmse_color: str = 'tab:blue',
+                 rmse_marker: str = 'o-',
+                 mae_color: str = 'tab:orange',
+                 mae_marker: str = 's--'):
+        """
+        Initializes the formatting configuration for seq-to-seq metrics.
+
+        Args:
+            font_size (int): The base font size to apply to the plots. Defaults to 16.
+            plot_figsize (Tuple[int, int]): Figure size for the plot. Defaults to (10, 6).
+            grid_style (str): Matplotlib linestyle for the plot grid. Defaults to '--'.
+                - Options: '--' (dashed), ':' (dotted), '-.' (dash-dot), '-' (solid)
+            rmse_color (str): Matplotlib color for the RMSE line. Defaults to 'tab:blue'.
+                - Common color names: 'tab:blue', 'crimson', 'forestgreen', '#4682B4'
+            rmse_marker (str): Matplotlib marker style for the RMSE line. Defaults to 'o-'.
+                - Options: 'o-' (circle), 's--' (square), '^:' (triangle), 'x' (x marker)
+            mae_color (str): Matplotlib color for the MAE line. Defaults to 'tab:orange'.
+                - Common color names: 'tab:orange', 'purple', 'black', '#FF6347'
+            mae_marker (str): Matplotlib marker style for the MAE line. Defaults to 's--'.
+                - Options: 's--', 'o-', 'v:', '+' (plus marker)
+        
+        <br>
+        
+        ## [Matplotlib Colors](https://matplotlib.org/stable/users/explain/colors/colors.html)
+        ## [Matplotlib Linestyles](https://matplotlib.org/stable/gallery/lines_bars_and_markers/linestyles.html)
+        ## [Matplotlib Markers](https://matplotlib.org/stable/api/markers_api.html)
+        """
+        self.font_size = font_size
+        self.plot_figsize = plot_figsize
+        self.grid_style = grid_style
+        self.rmse_color = rmse_color
+        self.rmse_marker = rmse_marker
+        self.mae_color = mae_color
+        self.mae_marker = mae_marker
+
+    def __repr__(self) -> str:
+        parts = [
+            f"font_size={self.font_size}",
+            f"plot_figsize={self.plot_figsize}",
+            f"grid_style='{self.grid_style}'",
+            f"rmse_color='{self.rmse_color}'",
+            f"mae_color='{self.mae_color}'"
+        ]
+        return f"SequenceMetricsFormat({', '.join(parts)})"
+
 def info():
     _script_info(__all__)