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

ml_tools/_core/_ML_trainer.py

@@ -1,2323 +0,0 @@
-from typing import List, Literal, Union, Optional, Callable, Dict, Any
-from pathlib import Path
-from torch.utils.data import DataLoader, Dataset
-import torch
-from torch import nn
-import numpy as np
-from abc import ABC, abstractmethod
-
-from ._path_manager import make_fullpath
-from ._ML_callbacks import _Callback, History, TqdmProgressBar, DragonModelCheckpoint, _DragonEarlyStopping, _DragonLRScheduler
-from ._ML_evaluation import classification_metrics, regression_metrics, plot_losses, shap_summary_plot, plot_attention_importance
-from ._ML_evaluation_multi import multi_target_regression_metrics, multi_label_classification_metrics, multi_target_shap_summary_plot
-from ._ML_vision_evaluation import segmentation_metrics, object_detection_metrics
-from ._ML_sequence_evaluation import sequence_to_sequence_metrics, sequence_to_value_metrics
-from ._ML_evaluation_captum import captum_feature_importance, _is_captum_available, captum_segmentation_heatmap, captum_image_heatmap
-from ._ML_configuration import (RegressionMetricsFormat, 
-                            MultiTargetRegressionMetricsFormat,
-                            BinaryClassificationMetricsFormat,
-                            MultiClassClassificationMetricsFormat,
-                            BinaryImageClassificationMetricsFormat,
-                            MultiClassImageClassificationMetricsFormat,
-                            MultiLabelBinaryClassificationMetricsFormat,
-                            BinarySegmentationMetricsFormat,
-                            MultiClassSegmentationMetricsFormat,
-                            SequenceValueMetricsFormat,
-                            SequenceSequenceMetricsFormat,
-
-                            FinalizeBinaryClassification,
-                            FinalizeBinarySegmentation,
-                            FinalizeBinaryImageClassification,
-                            FinalizeMultiClassClassification,
-                            FinalizeMultiClassImageClassification,
-                            FinalizeMultiClassSegmentation,
-                            FinalizeMultiLabelBinaryClassification,
-                            FinalizeMultiTargetRegression,
-                            FinalizeRegression,
-                            FinalizeObjectDetection,
-                            FinalizeSequenceSequencePrediction,
-                            FinalizeSequenceValuePrediction)
-
-from ._script_info import _script_info
-from ._keys import PyTorchLogKeys, PyTorchCheckpointKeys, DatasetKeys, MLTaskKeys, MagicWords, DragonTrainerKeys, SequenceDatasetKeys, ScalerKeys
-from ._logger import get_logger
-
-
-_LOGGER = get_logger("DragonTrainer")
-
-
-__all__ = [
-    "DragonTrainer",
-    "DragonDetectionTrainer",
-    "DragonSequenceTrainer"
-]
-
-class _BaseDragonTrainer(ABC):
-    """
-    Abstract base class for Dragon Trainers.
-    
-    Handles the common training loop orchestration, checkpointing, callback
-    management, and device handling. Subclasses must implement the
-    task-specific logic (dataloaders, train/val steps, evaluation).
-    """
-    def __init__(self, 
-                 model: nn.Module, 
-                 optimizer: torch.optim.Optimizer, 
-                 device: Union[Literal['cuda', 'mps', 'cpu'],str], 
-                 dataloader_workers: int = 2,
-                 checkpoint_callback: Optional[DragonModelCheckpoint] = None,
-                 early_stopping_callback: Optional[_DragonEarlyStopping] = None,
-                 lr_scheduler_callback: Optional[_DragonLRScheduler] = None,
-                 extra_callbacks: Optional[List[_Callback]] = None):
-
-        self.model = model
-        self.optimizer = optimizer
-        self.scheduler = None
-        self.device = self._validate_device(device)
-        self.dataloader_workers = dataloader_workers
-        
-        # Callback handler
-        default_callbacks = [History(), TqdmProgressBar()]
-        
-        self._checkpoint_callback = None
-        if checkpoint_callback:
-            default_callbacks.append(checkpoint_callback)
-            self._checkpoint_callback = checkpoint_callback
-        if early_stopping_callback:
-            default_callbacks.append(early_stopping_callback)
-        if lr_scheduler_callback:
-            default_callbacks.append(lr_scheduler_callback)
-        
-        user_callbacks = extra_callbacks if extra_callbacks is not None else []
-        self.callbacks = default_callbacks + user_callbacks
-        self._set_trainer_on_callbacks()
-
-        # Internal state
-        self.train_loader: Optional[DataLoader] = None
-        self.validation_loader: Optional[DataLoader] = None 
-        self.history: Dict[str, List[Any]] = {}
-        self.epoch = 0
-        self.epochs = 0 # Total epochs for the fit run
-        self.start_epoch = 1
-        self.stop_training = False
-        self._batch_size = 10
-
-    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 = "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 = "cpu"
-        return torch.device(device)
-
-    def _set_trainer_on_callbacks(self):
-        """Gives each callback a reference to this trainer instance."""
-        for callback in self.callbacks:
-            callback.set_trainer(self)
-
-    def _load_checkpoint(self, path: Union[str, Path]):
-        """Loads a training checkpoint to resume training."""
-        p = make_fullpath(path, enforce="file")
-        _LOGGER.info(f"Loading checkpoint from '{p.name}'...")
-        
-        try:
-            checkpoint = torch.load(p, map_location=self.device)
-            
-            if PyTorchCheckpointKeys.MODEL_STATE not in checkpoint or PyTorchCheckpointKeys.OPTIMIZER_STATE not in checkpoint:
-                _LOGGER.error(f"Checkpoint file '{p.name}' is invalid. Missing 'model_state_dict' or 'optimizer_state_dict'.")
-                raise KeyError()
-
-            self.model.load_state_dict(checkpoint[PyTorchCheckpointKeys.MODEL_STATE])
-            self.optimizer.load_state_dict(checkpoint[PyTorchCheckpointKeys.OPTIMIZER_STATE])
-            self.epoch = checkpoint.get(PyTorchCheckpointKeys.EPOCH, 0)
-            self.start_epoch = self.epoch + 1 # Resume on the *next* epoch
-            
-            # --- Load History ---
-            if PyTorchCheckpointKeys.HISTORY in checkpoint:
-                self.history = checkpoint[PyTorchCheckpointKeys.HISTORY]
-                _LOGGER.info(f"Restored training history up to epoch {self.epoch}.")
-            else:
-                _LOGGER.warning("No 'history' found in checkpoint. A new history will be started.")
-                self.history = {} # Ensure it's at least an empty dict
-            
-            # --- Scheduler State Loading Logic ---
-            scheduler_state_exists = PyTorchCheckpointKeys.SCHEDULER_STATE in checkpoint
-            scheduler_object_exists = self.scheduler is not None
-
-            if scheduler_object_exists and scheduler_state_exists:
-                # Case 1: Both exist. Attempt to load.
-                try:
-                    self.scheduler.load_state_dict(checkpoint[PyTorchCheckpointKeys.SCHEDULER_STATE]) # type: ignore
-                    scheduler_name = self.scheduler.__class__.__name__
-                    _LOGGER.info(f"Restored LR scheduler state for: {scheduler_name}")
-                except Exception as e:
-                    # Loading failed, likely a mismatch
-                    scheduler_name = self.scheduler.__class__.__name__
-                    _LOGGER.error(f"Failed to load scheduler state for '{scheduler_name}'. A different scheduler type might have been used.")
-                    raise e
-
-            elif scheduler_object_exists and not scheduler_state_exists:
-                # Case 2: Scheduler provided, but no state in checkpoint.
-                scheduler_name = self.scheduler.__class__.__name__
-                _LOGGER.warning(f"'{scheduler_name}' was provided, but no scheduler state was found in the checkpoint. The scheduler will start from its initial state.")
-            
-            elif not scheduler_object_exists and scheduler_state_exists:
-                # Case 3: State in checkpoint, but no scheduler provided.
-                _LOGGER.error("Checkpoint contains an LR scheduler state, but no LRScheduler callback was provided.")
-                raise ValueError()
-            
-            # Restore callback states
-            for cb in self.callbacks:
-                if isinstance(cb, DragonModelCheckpoint) and PyTorchCheckpointKeys.BEST_SCORE in checkpoint:
-                    cb.best = checkpoint[PyTorchCheckpointKeys.BEST_SCORE]
-                    _LOGGER.info(f"Restored {cb.__class__.__name__} 'best' score to: {cb.best:.4f}")
-            
-            _LOGGER.info(f"Checkpoint loaded. Resuming training from epoch {self.start_epoch}.")
-            
-        except Exception as e:
-            _LOGGER.error(f"Failed to load checkpoint from '{p}': {e}")
-            raise
-
-    def fit(self, 
-            save_dir: Union[str,Path],
-            epochs: int = 100, 
-            batch_size: int = 10, 
-            shuffle: bool = True,
-            resume_from_checkpoint: Optional[Union[str, Path]] = None):
-        """
-        Starts the training-validation process of the model.
-        
-        Returns the "History" callback dictionary.
-
-        Args:
-            save_dir (str | Path): Directory to save the loss plot.
-            epochs (int): The total number of epochs to train for.
-            batch_size (int): The number of samples per batch.
-            shuffle (bool): Whether to shuffle the training data at each epoch.
-            resume_from_checkpoint (str | Path | None): Optional path to a checkpoint to resume training.
-        """
-        self.epochs = epochs
-        self._batch_size = batch_size
-        self._create_dataloaders(self._batch_size, shuffle) # type: ignore
-        self.model.to(self.device)
-        
-        if resume_from_checkpoint:
-            self._load_checkpoint(resume_from_checkpoint)
-        
-        # Reset stop_training flag on the trainer
-        self.stop_training = False
-
-        self._callbacks_hook('on_train_begin')
-        
-        if not self.train_loader:
-            _LOGGER.error("Train loader is not initialized.")
-            raise ValueError()
-        
-        if not self.validation_loader:
-            _LOGGER.error("Validation loader is not initialized.")
-            raise ValueError()
-        
-        for epoch in range(self.start_epoch, self.epochs + 1):
-            self.epoch = epoch
-            epoch_logs: Dict[str, Any] = {}
-            self._callbacks_hook('on_epoch_begin', epoch, logs=epoch_logs)
-            
-            train_logs = self._train_step()
-            epoch_logs.update(train_logs)
-
-            val_logs = self._validation_step()
-            epoch_logs.update(val_logs)
-            
-            self._callbacks_hook('on_epoch_end', epoch, logs=epoch_logs)
-            
-            # Check the early stopping flag
-            if self.stop_training:
-                break
-
-        self._callbacks_hook('on_train_end')
-        
-        # Training History
-        plot_losses(self.history, save_dir=save_dir)
-        
-        return self.history
-
-    def _callbacks_hook(self, method_name: str, *args, **kwargs):
-        """Calls the specified method on all callbacks."""
-        for callback in self.callbacks:
-            method = getattr(callback, method_name)
-            method(*args, **kwargs)
-            
-    def to_cpu(self):
-        """
-        Moves the model to the CPU and updates the trainer's device setting.
-        
-        This is useful for running operations that require the CPU.
-        """
-        self.device = torch.device('cpu')
-        self.model.to(self.device)
-        _LOGGER.info("Trainer and model moved to CPU.")
-    
-    def to_device(self, device: str):
-        """
-        Moves the model to the specified device and updates the trainer's device setting.
-
-        Args:
-            device (str): The target device (e.g., 'cuda', 'mps', 'cpu').
-        """
-        self.device = self._validate_device(device)
-        self.model.to(self.device)
-        _LOGGER.info(f"Trainer and model moved to {self.device}.")
-        
-    def _load_model_state_for_finalizing(self, model_checkpoint: Union[Path, Literal['best', 'current']]):
-        """
-        Private helper to load the correct model state_dict based on user's choice.
-        This is called by finalize_model_training() in subclasses.
-        """
-        if isinstance(model_checkpoint, Path):
-            self._load_checkpoint(path=model_checkpoint)
-        elif model_checkpoint == MagicWords.BEST and self._checkpoint_callback:
-            path_to_latest = self._checkpoint_callback.best_checkpoint_path
-            self._load_checkpoint(path_to_latest)
-        elif model_checkpoint == MagicWords.BEST and self._checkpoint_callback is None:
-            _LOGGER.error(f"'model_checkpoint' set to '{MagicWords.BEST}' but no checkpoint callback was found.")
-            raise ValueError()
-        elif model_checkpoint == MagicWords.CURRENT:
-            pass
-        else:
-            _LOGGER.error(f"Unknown 'model_checkpoint' received '{model_checkpoint}'.")
-            raise ValueError()
-        
-    # --- Abstract Methods ---
-    # These must be implemented by subclasses
-
-    @abstractmethod
-    def _create_dataloaders(self, batch_size: int, shuffle: bool):
-        """Initializes the DataLoaders."""
-        raise NotImplementedError
-
-    @abstractmethod
-    def _train_step(self) -> Dict[str, float]:
-        """Runs a single training epoch."""
-        raise NotImplementedError
-
-    @abstractmethod
-    def _validation_step(self) -> Dict[str, float]:
-        """Runs a single validation epoch."""
-        raise NotImplementedError
-
-    @abstractmethod
-    def evaluate(self, *args, **kwargs):
-        """Runs the full model evaluation."""
-        raise NotImplementedError
-    
-    @abstractmethod
-    def _evaluate(self, *args, **kwargs):
-        """Internal evaluation helper."""
-        raise NotImplementedError
-
-    @abstractmethod
-    def finalize_model_training(self, *args, **kwargs):
-        """Saves the finalized model for inference."""
-        raise NotImplementedError
-
-
-# --- DragonTrainer ----
-class DragonTrainer(_BaseDragonTrainer):
-    def __init__(self, 
-                 model: nn.Module, 
-                 train_dataset: Dataset, 
-                 validation_dataset: Dataset, 
-                 kind: Literal["regression", "binary classification", "multiclass classification", 
-                               "multitarget regression", "multilabel binary classification", 
-                               "binary segmentation", "multiclass segmentation", "binary image classification", "multiclass image classification"],
-                 optimizer: torch.optim.Optimizer, 
-                 device: Union[Literal['cuda', 'mps', 'cpu'],str], 
-                 checkpoint_callback: Optional[DragonModelCheckpoint],
-                 early_stopping_callback: Optional[_DragonEarlyStopping],
-                 lr_scheduler_callback: Optional[_DragonLRScheduler],
-                 extra_callbacks: Optional[List[_Callback]] = None,
-                 criterion: Union[nn.Module,Literal["auto"]] = "auto", 
-                 dataloader_workers: int = 2):
-        """
-        Automates the training process of a PyTorch Model.
-        
-        Built-in Callbacks: `History`, `TqdmProgressBar`
-
-        Args:
-            model (nn.Module): The PyTorch model to train.
-            train_dataset (Dataset): The training dataset.
-            validation_dataset (Dataset): The validation dataset.
-            kind (str): Used to redirect to the correct process. 
-            criterion (nn.Module | "auto"): The loss function to use. If "auto", it will be inferred from the selected task
-            optimizer (torch.optim.Optimizer): The optimizer.
-            device (str): The device to run training on ('cpu', 'cuda', 'mps').
-            dataloader_workers (int): Subprocesses for data loading.
-            extra_callbacks (List[Callback] | None): A list of extra callbacks to use during training.
-            
-        Note:
-            - For **regression** and **multi_target_regression** tasks, suggested criterions include `nn.MSELoss` or `nn.L1Loss`. The model should output as many logits as existing targets.
-            
-            - For **single-label, binary classification**, `nn.BCEWithLogitsLoss` is the standard choice. The model should output a single logit.
-    
-            - For **single-label, multi-class classification** tasks, `nn.CrossEntropyLoss` is the standard choice. The model should output as many logits as existing classes.
-    
-            - For **multi-label, binary classification** tasks (where each label is a 0 or 1), `nn.BCEWithLogitsLoss` is the correct choice as it treats each output as an independent binary problem. The model should output 1 logit per binary target.
-        
-            - For **binary segmentation** tasks, `nn.BCEWithLogitsLoss` is common. The model should output a single logit.
-            
-            - for **multiclass segmentation** tasks, `nn.CrossEntropyLoss` is the standard. The model should output as many logits as existing classes.
-        """
-        # Call the base class constructor with common parameters
-        super().__init__(
-            model=model,
-            optimizer=optimizer,
-            device=device,
-            dataloader_workers=dataloader_workers,
-            checkpoint_callback=checkpoint_callback,
-            early_stopping_callback=early_stopping_callback,
-            lr_scheduler_callback=lr_scheduler_callback,
-            extra_callbacks=extra_callbacks
-        )
-        
-        if kind not in [MLTaskKeys.REGRESSION,
-                        MLTaskKeys.BINARY_CLASSIFICATION,
-                        MLTaskKeys.MULTICLASS_CLASSIFICATION,
-                        MLTaskKeys.MULTILABEL_BINARY_CLASSIFICATION,
-                        MLTaskKeys.MULTITARGET_REGRESSION,
-                        MLTaskKeys.BINARY_SEGMENTATION,
-                        MLTaskKeys.MULTICLASS_SEGMENTATION,
-                        MLTaskKeys.BINARY_IMAGE_CLASSIFICATION,
-                        MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION]:
-            raise ValueError(f"'{kind}' is not a valid task type.")
-
-        self.train_dataset = train_dataset
-        self.validation_dataset = validation_dataset
-        self.kind = kind
-        self._classification_threshold: float = 0.5
-        
-        # loss function
-        if criterion == "auto":
-            if kind in [MLTaskKeys.REGRESSION, MLTaskKeys.MULTITARGET_REGRESSION]:
-                self.criterion = nn.MSELoss()
-            elif kind in [MLTaskKeys.BINARY_CLASSIFICATION, MLTaskKeys.BINARY_IMAGE_CLASSIFICATION, MLTaskKeys.BINARY_SEGMENTATION, MLTaskKeys.MULTILABEL_BINARY_CLASSIFICATION]:
-                self.criterion = nn.BCEWithLogitsLoss()
-            elif kind in [MLTaskKeys.MULTICLASS_CLASSIFICATION, MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION, MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION, MLTaskKeys.MULTICLASS_SEGMENTATION]:
-                self.criterion = nn.CrossEntropyLoss()
-        else:
-            self.criterion = criterion
-
-    def _create_dataloaders(self, batch_size: int, shuffle: bool):
-        """Initializes the DataLoaders."""
-        # Ensure stability on MPS devices by setting num_workers to 0
-        loader_workers = 0 if self.device.type == 'mps' else self.dataloader_workers
-        
-        self.train_loader = DataLoader(
-            dataset=self.train_dataset, 
-            batch_size=batch_size, 
-            shuffle=shuffle, 
-            num_workers=loader_workers, 
-            pin_memory=("cuda" in self.device.type),
-            drop_last=True  # Drops the last batch if incomplete, selecting a good batch size is key.
-        )
-        
-        self.validation_loader = DataLoader(
-            dataset=self.validation_dataset, 
-            batch_size=batch_size, 
-            shuffle=False, 
-            num_workers=loader_workers, 
-            pin_memory=("cuda" in self.device.type)
-        )
-
-    def _train_step(self):
-        self.model.train()
-        running_loss = 0.0
-        total_samples = 0
-        
-        for batch_idx, (features, target) in enumerate(self.train_loader): # type: ignore
-            # Create a log dictionary for the batch
-            batch_logs = {
-                PyTorchLogKeys.BATCH_INDEX: batch_idx, 
-                PyTorchLogKeys.BATCH_SIZE: features.size(0)
-            }
-            self._callbacks_hook('on_batch_begin', batch_idx, logs=batch_logs)
-
-            features, target = features.to(self.device), target.to(self.device)
-            self.optimizer.zero_grad()
-            
-            output = self.model(features)
-            
-            # --- Label Type/Shape Correction ---
-            # Cast target to float for BCE-based losses
-            if self.kind in MLTaskKeys.ALL_BINARY_TASKS:
-                target = target.float()
-
-            # Reshape output to match target for single-logit tasks
-            if self.kind in [MLTaskKeys.REGRESSION, MLTaskKeys.BINARY_CLASSIFICATION, MLTaskKeys.BINARY_IMAGE_CLASSIFICATION]:
-                # If model outputs [N, 1] and target is [N], squeeze output
-                if output.ndim == 2 and output.shape[1] == 1 and target.ndim == 1:
-                    output = output.squeeze(1)
-            
-            if self.kind == MLTaskKeys.BINARY_SEGMENTATION:
-                # If model outputs [N, 1, H, W] and target is [N, H, W], squeeze output
-                if output.ndim == 4 and output.shape[1] == 1 and target.ndim == 3:
-                    output = output.squeeze(1)
-                
-            loss = self.criterion(output, target)
-            
-            loss.backward()
-            self.optimizer.step()
-
-            # Calculate batch loss and update running loss for the epoch
-            batch_loss = loss.item()
-            batch_size = features.size(0)
-            running_loss += batch_loss * batch_size  # Accumulate total loss
-            total_samples += batch_size # total samples
-            
-            # Add the batch loss to the logs and call the end-of-batch hook
-            batch_logs[PyTorchLogKeys.BATCH_LOSS] = batch_loss
-            self._callbacks_hook('on_batch_end', batch_idx, logs=batch_logs)
-        
-        if total_samples == 0:
-            _LOGGER.warning("No samples processed in a train_step. Returning 0 loss.")
-            return {PyTorchLogKeys.TRAIN_LOSS: 0.0}
-
-        return {PyTorchLogKeys.TRAIN_LOSS: running_loss / total_samples} # type: ignore
-
-    def _validation_step(self):
-        self.model.eval()
-        running_loss = 0.0
-        
-        with torch.no_grad():
-            for features, target in self.validation_loader: # type: ignore
-                features, target = features.to(self.device), target.to(self.device)
-                
-                output = self.model(features)
-                
-                # --- Label Type/Shape Correction ---
-                # Cast target to float for BCE-based losses
-                if self.kind in MLTaskKeys.ALL_BINARY_TASKS:
-                    target = target.float()
-
-                # Reshape output to match target for single-logit tasks
-                if self.kind in [MLTaskKeys.REGRESSION, MLTaskKeys.BINARY_CLASSIFICATION, MLTaskKeys.BINARY_IMAGE_CLASSIFICATION]:
-                    # If model outputs [N, 1] and target is [N], squeeze output
-                    if output.ndim == 2 and output.shape[1] == 1 and target.ndim == 1:
-                        output = output.squeeze(1)
-                
-                if self.kind == MLTaskKeys.BINARY_SEGMENTATION:
-                    # If model outputs [N, 1, H, W] and target is [N, H, W], squeeze output
-                    if output.ndim == 4 and output.shape[1] == 1 and target.ndim == 3:
-                        output = output.squeeze(1)
-                
-                loss = self.criterion(output, target)
-                
-                running_loss += loss.item() * features.size(0)
-                
-        if not self.validation_loader.dataset: # type: ignore
-            _LOGGER.warning("No samples processed in _validation_step. Returning 0 loss.")
-            return {PyTorchLogKeys.VAL_LOSS: 0.0}
-        
-        logs = {PyTorchLogKeys.VAL_LOSS: running_loss / len(self.validation_loader.dataset)} # type: ignore
-        return logs
-    
-    def _predict_for_eval(self, dataloader: DataLoader):
-        """
-        Private method to yield model predictions batch by batch for evaluation.
-        
-        Automatically detects if `target_scaler` is present in the training dataset
-        and applies inverse transformation for Regression tasks.
-        
-        Yields:
-            tuple: A tuple containing (y_pred_batch, y_prob_batch, y_true_batch).
-                   
-                - y_prob_batch is None for regression tasks.
-        """
-        self.model.eval()
-        self.model.to(self.device)
-        
-        # --- Check for Target Scaler (for Regression Un-scaling) ---
-        target_scaler = None
-        if self.kind in [MLTaskKeys.REGRESSION, MLTaskKeys.MULTITARGET_REGRESSION]:
-            # Try to get the scaler from the dataset attached to the trainer
-            if hasattr(self.train_dataset, ScalerKeys.TARGET_SCALER):
-                target_scaler = getattr(self.train_dataset, ScalerKeys.TARGET_SCALER)
-                if target_scaler is not None:
-                     _LOGGER.debug("Target scaler detected. Un-scaling predictions and targets for metric calculation.")
-        
-        with torch.no_grad():
-            for features, target in dataloader:
-                features = features.to(self.device)
-                # Keep target on device initially for potential un-scaling
-                target = target.to(self.device) 
-                
-                output = self.model(features)
-
-                y_pred_batch = None
-                y_prob_batch = None
-                y_true_batch = None
-
-                if self.kind in [MLTaskKeys.REGRESSION, MLTaskKeys.MULTITARGET_REGRESSION]:
-                    
-                    # --- Automatic Un-scaling Logic ---
-                    if target_scaler:
-                        # 1. Reshape output/target if flattened (common in single regression)
-                        # Scaler expects [N, Features]
-                        original_out_shape = output.shape
-                        original_target_shape = target.shape
-                        
-                        if output.ndim == 1: output = output.reshape(-1, 1)
-                        if target.ndim == 1: target = target.reshape(-1, 1)
-                            
-                        # 2. Apply Inverse Transform
-                        output = target_scaler.inverse_transform(output)
-                        target = target_scaler.inverse_transform(target)
-                        
-                        # 3. Restore shapes (optional, but good for consistency)
-                        if len(original_out_shape) == 1: output = output.flatten()
-                        if len(original_target_shape) == 1: target = target.flatten()
-
-                    y_pred_batch = output.cpu().numpy()
-                    y_true_batch = target.cpu().numpy()
-                    
-                elif self.kind in [MLTaskKeys.BINARY_CLASSIFICATION, MLTaskKeys.BINARY_IMAGE_CLASSIFICATION]:
-                    if output.ndim == 2 and output.shape[1] == 1:
-                        output = output.squeeze(1)
-                        
-                    probs_pos = torch.sigmoid(output) 
-                    preds = (probs_pos >= self._classification_threshold).int()
-                    y_pred_batch = preds.cpu().numpy()
-                    
-                    probs_neg = 1.0 - probs_pos
-                    y_prob_batch = torch.stack([probs_neg, probs_pos], dim=1).cpu().numpy()
-                    y_true_batch = target.cpu().numpy()
-
-                elif self.kind in [MLTaskKeys.MULTICLASS_CLASSIFICATION, MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION]:
-                    probs = torch.softmax(output, dim=1)
-                    preds = torch.argmax(probs, dim=1)
-                    y_pred_batch = preds.cpu().numpy()
-                    y_prob_batch = probs.cpu().numpy()
-                    y_true_batch = target.cpu().numpy()
-
-                elif self.kind == MLTaskKeys.MULTILABEL_BINARY_CLASSIFICATION:
-                    probs = torch.sigmoid(output)
-                    preds = (probs >= self._classification_threshold).int()
-                    y_pred_batch = preds.cpu().numpy()
-                    y_prob_batch = probs.cpu().numpy()
-                    y_true_batch = target.cpu().numpy()
-                
-                elif self.kind == MLTaskKeys.BINARY_SEGMENTATION:
-                    probs_pos = torch.sigmoid(output) 
-                    preds = (probs_pos >= self._classification_threshold).int() 
-                    y_pred_batch = preds.squeeze(1).cpu().numpy()
-
-                    probs_neg = 1.0 - probs_pos
-                    y_prob_batch = torch.cat([probs_neg, probs_pos], dim=1).cpu().numpy()
-
-                    if target.ndim == 4 and target.shape[1] == 1:
-                        target = target.squeeze(1)
-                    y_true_batch = target.cpu().numpy()
-                    
-                elif self.kind == MLTaskKeys.MULTICLASS_SEGMENTATION:
-                    probs = torch.softmax(output, dim=1)
-                    preds = torch.argmax(probs, dim=1) 
-                    y_pred_batch = preds.cpu().numpy()
-                    y_prob_batch = probs.cpu().numpy() 
-                    
-                    if target.ndim == 4 and target.shape[1] == 1:
-                        target = target.squeeze(1)
-                    y_true_batch = target.cpu().numpy()
-
-                yield y_pred_batch, y_prob_batch, y_true_batch
-                
-    def evaluate(self, 
-                 save_dir: Union[str, Path], 
-                 model_checkpoint: Union[Path, Literal["best", "current"]],
-                 classification_threshold: Optional[float] = None,
-                 test_data: Optional[Union[DataLoader, Dataset]] = None,
-                 val_format_configuration: Optional[Union[
-                        RegressionMetricsFormat, 
-                        MultiTargetRegressionMetricsFormat,
-                        BinaryClassificationMetricsFormat,
-                        MultiClassClassificationMetricsFormat,
-                        BinaryImageClassificationMetricsFormat,
-                        MultiClassImageClassificationMetricsFormat,
-                        MultiLabelBinaryClassificationMetricsFormat,
-                        BinarySegmentationMetricsFormat,
-                        MultiClassSegmentationMetricsFormat
-                    ]]=None,
-                 test_format_configuration: Optional[Union[
-                        RegressionMetricsFormat, 
-                        MultiTargetRegressionMetricsFormat,
-                        BinaryClassificationMetricsFormat,
-                        MultiClassClassificationMetricsFormat,
-                        BinaryImageClassificationMetricsFormat,
-                        MultiClassImageClassificationMetricsFormat,
-                        MultiLabelBinaryClassificationMetricsFormat,
-                        BinarySegmentationMetricsFormat,
-                        MultiClassSegmentationMetricsFormat,
-                    ]]=None):
-        """
-        Evaluates the model, routing to the correct evaluation function based on task `kind`.
-
-        Args:
-            model_checkpoint (Path | "best" | "current"): 
-                - Path to a valid checkpoint for the model. The state of the trained model will be overwritten in place.
-                - If 'best', the best checkpoint will be loaded if a DragonModelCheckpoint was provided. The state of the trained model will be overwritten in place.
-                - If 'current', use the current state of the trained model up the latest trained epoch.
-            save_dir (str | Path): Directory to save all reports and plots.
-            classification_threshold (float | None): Used for tasks using a binary approach (binary classification, binary segmentation, multilabel binary classification)
-            test_data (DataLoader | Dataset | None): Optional Test data to evaluate the model performance. Validation and Test metrics will be saved to subdirectories.
-            val_format_configuration (object): Optional configuration for metric format output for the validation set.
-            test_format_configuration (object): Optional configuration for metric format output for the test set.
-        """
-        # Validate model checkpoint
-        if isinstance(model_checkpoint, Path):
-            checkpoint_validated = make_fullpath(model_checkpoint, enforce="file")
-        elif model_checkpoint in [MagicWords.BEST, MagicWords.CURRENT]:
-            checkpoint_validated = model_checkpoint
-        else:
-            _LOGGER.error(f"'model_checkpoint' must be a Path object, or the string '{MagicWords.BEST}', or the string '{MagicWords.CURRENT}'.")
-            raise ValueError()
-        
-        # Validate classification threshold
-        if self.kind not in MLTaskKeys.ALL_BINARY_TASKS:
-            # dummy value for tasks that do not need it
-            threshold_validated = 0.5
-        elif classification_threshold is None:
-            # it should have been provided for binary tasks
-            _LOGGER.error(f"The classification threshold must be provided for '{self.kind}'.")
-            raise ValueError()
-        elif classification_threshold <= 0.0 or classification_threshold >= 1.0:
-            # Invalid float
-            _LOGGER.error(f"A classification threshold of {classification_threshold} is invalid. Must be in the range (0.0 - 1.0).")
-            raise ValueError()
-        else:
-            threshold_validated = classification_threshold
-        
-        # Validate val configuration
-        if val_format_configuration is not None:
-            if not isinstance(val_format_configuration, (RegressionMetricsFormat, 
-                                                        MultiTargetRegressionMetricsFormat,
-                                                        BinaryClassificationMetricsFormat,
-                                                        MultiClassClassificationMetricsFormat,
-                                                        BinaryImageClassificationMetricsFormat,
-                                                        MultiClassImageClassificationMetricsFormat,
-                                                        MultiLabelBinaryClassificationMetricsFormat,
-                                                        BinarySegmentationMetricsFormat,
-                                                        MultiClassSegmentationMetricsFormat)):
-                _LOGGER.error(f"Invalid 'format_configuration': '{type(val_format_configuration)}'.")
-                raise ValueError()
-            else:
-                val_configuration_validated = val_format_configuration
-        else: # config is None
-            val_configuration_validated = None
-        
-        # Validate directory
-        save_path = make_fullpath(save_dir, make=True, enforce="directory")
-        
-        # Validate test data and dispatch
-        if test_data is not None:
-            if not isinstance(test_data, (DataLoader, Dataset)):
-                _LOGGER.error(f"Invalid type for 'test_data': '{type(test_data)}'.")
-                raise ValueError()
-            test_data_validated = test_data
-                
-            validation_metrics_path = save_path / DragonTrainerKeys.VALIDATION_METRICS_DIR
-            test_metrics_path = save_path / DragonTrainerKeys.TEST_METRICS_DIR
-            
-            # Dispatch validation set
-            _LOGGER.info(f"🔎 Evaluating on validation dataset. Metrics will be saved to '{DragonTrainerKeys.VALIDATION_METRICS_DIR}'")
-            self._evaluate(save_dir=validation_metrics_path,
-                           model_checkpoint=checkpoint_validated,
-                           classification_threshold=threshold_validated,
-                           data=None,
-                           format_configuration=val_configuration_validated)
-            
-            # Validate test configuration
-            if test_format_configuration is not None:
-                if not isinstance(test_format_configuration, (RegressionMetricsFormat, 
-                                                        MultiTargetRegressionMetricsFormat,
-                                                        BinaryClassificationMetricsFormat,
-                                                        MultiClassClassificationMetricsFormat,
-                                                        BinaryImageClassificationMetricsFormat,
-                                                        MultiClassImageClassificationMetricsFormat,
-                                                        MultiLabelBinaryClassificationMetricsFormat,
-                                                        BinarySegmentationMetricsFormat,
-                                                        MultiClassSegmentationMetricsFormat)):
-                    warning_message_type = f"Invalid test_format_configuration': '{type(test_format_configuration)}'."
-                    if val_configuration_validated is not None:
-                        warning_message_type += " 'val_format_configuration' will be used for the test set metrics output."
-                        test_configuration_validated = val_configuration_validated
-                    else:
-                        warning_message_type += " Using default format."
-                        test_configuration_validated = None
-                    _LOGGER.warning(warning_message_type)
-                else:
-                    test_configuration_validated = test_format_configuration
-            else: #config is None
-                test_configuration_validated = None
-            
-            # Dispatch test set
-            _LOGGER.info(f"🔎 Evaluating on test dataset. Metrics will be saved to '{DragonTrainerKeys.TEST_METRICS_DIR}'")
-            self._evaluate(save_dir=test_metrics_path,
-                           model_checkpoint="current",
-                           classification_threshold=threshold_validated,
-                           data=test_data_validated,
-                           format_configuration=test_configuration_validated)
-        else:
-            # Dispatch validation set
-            _LOGGER.info(f"Evaluating on validation dataset. Metrics will be saved to '{save_path.name}'")
-            self._evaluate(save_dir=save_path,
-                           model_checkpoint=checkpoint_validated,
-                           classification_threshold=threshold_validated,
-                           data=None,
-                           format_configuration=val_configuration_validated)
-        
-    def _evaluate(self, 
-                 save_dir: Union[str, Path], 
-                 model_checkpoint: Union[Path, Literal["best", "current"]],
-                 classification_threshold: float,
-                 data: Optional[Union[DataLoader, Dataset]],
-                 format_configuration: Optional[Union[
-                        RegressionMetricsFormat, 
-                        MultiTargetRegressionMetricsFormat,
-                        BinaryClassificationMetricsFormat,
-                        MultiClassClassificationMetricsFormat,
-                        BinaryImageClassificationMetricsFormat,
-                        MultiClassImageClassificationMetricsFormat,
-                        MultiLabelBinaryClassificationMetricsFormat,
-                        BinarySegmentationMetricsFormat,
-                        MultiClassSegmentationMetricsFormat
-                    ]]=None):
-        """
-        Changed to a private helper function.
-        """
-        dataset_for_artifacts = None
-        eval_loader = None
-        
-        # set threshold
-        self._classification_threshold = classification_threshold
-        
-        # load model checkpoint
-        if isinstance(model_checkpoint, Path):
-            self._load_checkpoint(path=model_checkpoint)
-        elif model_checkpoint == MagicWords.BEST and self._checkpoint_callback:
-            path_to_latest = self._checkpoint_callback.best_checkpoint_path
-            self._load_checkpoint(path_to_latest)
-        elif model_checkpoint == MagicWords.BEST and self._checkpoint_callback is None:
-            _LOGGER.error(f"'model_checkpoint' set to '{MagicWords.BEST}' but no checkpoint callback was found.")
-            raise ValueError()
-        
-        # Dataloader
-        if isinstance(data, DataLoader):
-            eval_loader = data
-            # Try to get the dataset from the loader for fetching target names
-            if hasattr(data, 'dataset'):
-                dataset_for_artifacts = data.dataset # type: ignore
-        elif isinstance(data, Dataset):
-            # Create a new loader from the provided dataset
-            eval_loader = DataLoader(data, 
-                                     batch_size=self._batch_size, 
-                                     shuffle=False, 
-                                     num_workers=0 if self.device.type == 'mps' else self.dataloader_workers,
-                                     pin_memory=(self.device.type == "cuda"))
-            dataset_for_artifacts = data
-        else: # data is None, use the trainer's default test dataset
-            if self.validation_dataset is None:
-                _LOGGER.error("Cannot evaluate. No data provided and no validation dataset available in the trainer.")
-                raise ValueError()
-            # Create a fresh DataLoader from the test_dataset
-            eval_loader = DataLoader(self.validation_dataset, 
-                                     batch_size=self._batch_size, 
-                                     shuffle=False, 
-                                     num_workers=0 if self.device.type == 'mps' else self.dataloader_workers,
-                                     pin_memory=(self.device.type == "cuda"))
-            
-            dataset_for_artifacts = self.validation_dataset
-
-        if eval_loader is None:
-            _LOGGER.error("Cannot evaluate. No valid data was provided or found.")
-            raise ValueError()
-
-        # print("\n--- Model Evaluation ---")
-
-        all_preds, all_probs, all_true = [], [], []
-        for y_pred_b, y_prob_b, y_true_b in self._predict_for_eval(eval_loader):
-            if y_pred_b is not None: all_preds.append(y_pred_b)
-            if y_prob_b is not None: all_probs.append(y_prob_b)
-            if y_true_b is not None: all_true.append(y_true_b)
-
-        if not all_true:
-            _LOGGER.error("Evaluation failed: No data was processed.")
-            return
-
-        y_pred = np.concatenate(all_preds)
-        y_true = np.concatenate(all_true)
-        y_prob = np.concatenate(all_probs) if all_probs else None
-
-        # --- Routing Logic ---
-        # Single-target regression
-        if self.kind == MLTaskKeys.REGRESSION:
-            # Check configuration
-            config = None
-            if format_configuration and isinstance(format_configuration, RegressionMetricsFormat):
-                config = format_configuration
-            elif format_configuration:
-                _LOGGER.warning(f"Wrong configuration type: Received '{type(format_configuration).__name__}'.")
-            
-            regression_metrics(y_true=y_true.flatten(), 
-                               y_pred=y_pred.flatten(), 
-                               save_dir=save_dir,
-                               config=config)
-        
-        # single target classification
-        elif self.kind in [MLTaskKeys.BINARY_CLASSIFICATION, 
-                           MLTaskKeys.BINARY_IMAGE_CLASSIFICATION, 
-                           MLTaskKeys.MULTICLASS_CLASSIFICATION, 
-                           MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION]:
-            # get the class map if it exists
-            try:
-                class_map = dataset_for_artifacts.class_map # type: ignore
-            except AttributeError:
-                _LOGGER.warning(f"Dataset has no 'class_map' attribute. Using generics.")
-                class_map = None
-            else:
-                if not isinstance(class_map, dict):
-                    _LOGGER.warning(f"Dataset has a 'class_map' attribute, but it is not a dictionary: '{type(class_map)}'.")
-                    class_map = None
-            
-            # Check configuration
-            config = None
-            if format_configuration:
-                if self.kind == MLTaskKeys.BINARY_CLASSIFICATION and isinstance(format_configuration, BinaryClassificationMetricsFormat):
-                    config = format_configuration
-                elif self.kind == MLTaskKeys.BINARY_IMAGE_CLASSIFICATION and isinstance(format_configuration, BinaryImageClassificationMetricsFormat):
-                    config = format_configuration
-                elif self.kind == MLTaskKeys.MULTICLASS_CLASSIFICATION and isinstance(format_configuration, MultiClassClassificationMetricsFormat):
-                    config = format_configuration
-                elif self.kind == MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION and isinstance(format_configuration, MultiClassImageClassificationMetricsFormat):
-                    config = format_configuration
-                else:
-                    _LOGGER.warning(f"Wrong configuration type: Received '{type(format_configuration).__name__}'.")
-  
-            classification_metrics(save_dir=save_dir,
-                                   y_true=y_true,
-                                   y_pred=y_pred,
-                                   y_prob=y_prob,
-                                   class_map=class_map,
-                                   config=config)
-        
-        # multitarget regression
-        elif self.kind == MLTaskKeys.MULTITARGET_REGRESSION:
-            try:
-                target_names = dataset_for_artifacts.target_names # type: ignore
-            except AttributeError:
-                num_targets = y_true.shape[1]
-                target_names = [f"target_{i}" for i in range(num_targets)]
-                _LOGGER.warning(f"Dataset has no 'target_names' attribute. Using generic names.")
-                
-            # Check configuration
-            config = None
-            if format_configuration and isinstance(format_configuration, MultiTargetRegressionMetricsFormat):
-                config = format_configuration
-            elif format_configuration:
-                _LOGGER.warning(f"Wrong configuration type: Received '{type(format_configuration).__name__}'.")
-            
-            multi_target_regression_metrics(y_true=y_true, 
-                                            y_pred=y_pred,
-                                            target_names=target_names, 
-                                            save_dir=save_dir,
-                                            config=config)
-            
-        # multi-label binary classification
-        elif self.kind == MLTaskKeys.MULTILABEL_BINARY_CLASSIFICATION:
-            try:
-                target_names = dataset_for_artifacts.target_names # type: ignore
-            except AttributeError:
-                num_targets = y_true.shape[1]
-                target_names = [f"label_{i}" for i in range(num_targets)]
-                _LOGGER.warning(f"Dataset has no 'target_names' attribute. Using generic names.")
-            
-            if y_prob is None:
-                _LOGGER.error("Evaluation for multi_label_classification requires probabilities (y_prob).")
-                return
-            
-            # Check configuration
-            config = None
-            if format_configuration and isinstance(format_configuration, MultiLabelBinaryClassificationMetricsFormat):
-                config = format_configuration
-            elif format_configuration:
-                _LOGGER.warning(f"Wrong configuration type: Received '{type(format_configuration).__name__}'.")
-
-            multi_label_classification_metrics(y_true=y_true,
-                                               y_pred=y_pred,
-                                               y_prob=y_prob,
-                                               target_names=target_names,
-                                               save_dir=save_dir,
-                                               config=config)
-        
-        # Segmentation tasks
-        elif self.kind in [MLTaskKeys.BINARY_SEGMENTATION, MLTaskKeys.MULTICLASS_SEGMENTATION]:
-            class_names = None
-            try:
-                # Try to get 'classes' from VisionDatasetMaker
-                if hasattr(dataset_for_artifacts, 'classes'):
-                    class_names = dataset_for_artifacts.classes # type: ignore
-                # Fallback for Subset
-                elif hasattr(dataset_for_artifacts, 'dataset') and hasattr(dataset_for_artifacts.dataset, 'classes'): # type: ignore
-                     class_names = dataset_for_artifacts.dataset.classes # type: ignore
-            except AttributeError:
-                pass # class_names is still None
-
-            if class_names is None:
-                try:
-                    # Fallback to 'target_names'
-                    class_names = dataset_for_artifacts.target_names # type: ignore
-                except AttributeError:
-                    # Fallback to inferring from labels
-                    labels = np.unique(y_true)
-                    class_names = [f"Class {i}" for i in labels]
-                    _LOGGER.warning(f"Dataset has no 'classes' or 'target_names' attribute. Using generic names.")
-            
-            # Check configuration
-            config = None
-            if format_configuration and isinstance(format_configuration, (BinarySegmentationMetricsFormat, MultiClassSegmentationMetricsFormat)):
-                config = format_configuration
-            elif format_configuration:
-                _LOGGER.warning(f"Wrong configuration type: Received '{type(format_configuration).__name__}'.")
-            
-            segmentation_metrics(y_true=y_true,
-                                 y_pred=y_pred,
-                                 save_dir=save_dir,
-                                 class_names=class_names,
-                                 config=config)
-    
-    def explain_shap(self,
-                save_dir: Union[str,Path], 
-                explain_dataset: Optional[Dataset] = None, 
-                n_samples: int = 300,
-                feature_names: Optional[List[str]] = None,
-                target_names: Optional[List[str]] = None,
-                explainer_type: Literal['deep', 'kernel'] = 'kernel'):
-        """
-        Explains model predictions using SHAP and saves all artifacts.
-        
-        NOTE: SHAP support is limited to single-target tasks (Regression, Binary/Multiclass Classification).
-        For complex tasks (Multi-target, Multi-label, Sequences, Images), please use `explain_captum()`.
-
-        The background data is automatically sampled from the trainer's training dataset.
-        
-        This method automatically routes to the appropriate SHAP summary plot
-        function based on the task. If `feature_names` or `target_names` (multi-target) are not provided,
-        it will attempt to extract them from the dataset.
-
-        Args:
-            explain_dataset (Dataset | None): A specific dataset to explain. 
-                                                 If None, the trainer's test dataset is used.
-            n_samples (int): The number of samples to use for both background and explanation.
-            feature_names (list[str] | None): Feature names. If None, the names will be extracted from the Dataset and raise an error on failure.
-            target_names (list[str] | None): Target names for multi-target tasks.
-            save_dir (str | Path): Directory to save all SHAP artifacts.
-            explainer_type (Literal['deep', 'kernel']): The explainer to use.
-                - 'deep': Uses shap.DeepExplainer. Fast and efficient for PyTorch models.
-                - 'kernel': Uses shap.KernelExplainer. Model-agnostic but EXTREMELY slow and memory-intensive. Use with a very low 'n_samples'< 100.
-        """
-        # --- 1. Compatibility Guard ---
-        valid_shap_tasks = [
-            MLTaskKeys.REGRESSION, 
-            MLTaskKeys.BINARY_CLASSIFICATION, 
-            MLTaskKeys.MULTICLASS_CLASSIFICATION
-        ]
-        
-        if self.kind not in valid_shap_tasks:
-            _LOGGER.warning(f"SHAP explanation is deprecated for task '{self.kind}' due to instability. Please use 'explain_captum()' instead.")
-            return
-        
-        # memory efficient helper
-        def _get_random_sample(dataset: Dataset, num_samples: int):
-            """
-            Memory-efficiently samples data from a dataset.
-            """
-            if dataset is None:
-                return None
-            
-            dataset_len = len(dataset) # type: ignore
-            if dataset_len == 0:
-                return None
-            
-            # For MPS devices, num_workers must be 0 to ensure stability
-            loader_workers = 0 if self.device.type == 'mps' else self.dataloader_workers
-            
-            # Ensure batch_size is not larger than the dataset itself
-            batch_size = min(num_samples, 64, dataset_len) 
-            
-            loader = DataLoader(
-                dataset, 
-                batch_size=batch_size,
-                shuffle=True, # Shuffle to get random samples
-                num_workers=loader_workers
-            )
-            
-            collected_features = []
-            num_collected = 0
-            
-            for features, _ in loader:
-                collected_features.append(features)
-                num_collected += features.size(0)
-                if num_collected >= num_samples:
-                    break # Stop once we have enough samples
-            
-            if not collected_features:
-                return None
-            
-            full_data = torch.cat(collected_features, dim=0)
-            
-            # If we collected more than needed, trim it down
-            if full_data.size(0) > num_samples:
-                return full_data[:num_samples]
-            
-            return full_data
-        
-        # print(f"\n--- Preparing SHAP Data (sampling up to {n_samples} instances) ---")
-
-        # 1. Get background data from the trainer's train_dataset
-        background_data = _get_random_sample(self.train_dataset, n_samples)
-        if background_data is None:
-            _LOGGER.error("Trainer's train_dataset is empty or invalid. Skipping SHAP analysis.")
-            return
-
-        # 2. Determine target dataset and get explanation instances
-        target_dataset = explain_dataset if explain_dataset is not None else self.validation_dataset
-        instances_to_explain = _get_random_sample(target_dataset, n_samples)
-        if instances_to_explain is None:
-            _LOGGER.error("Explanation dataset is empty or invalid. Skipping SHAP analysis.")
-            return
-        
-        # attempt to get feature names
-        if feature_names is None:
-            # _LOGGER.info("`feature_names` not provided. Attempting to extract from dataset...")
-            if hasattr(target_dataset, DatasetKeys.FEATURE_NAMES):
-                feature_names = target_dataset.feature_names # type: ignore
-            else:
-                _LOGGER.error(f"Could not extract `feature_names` from the dataset. It must be provided if the dataset object does not have a '{DatasetKeys.FEATURE_NAMES}' attribute.")
-                raise ValueError()
-            
-        # move model to device
-        self.model.to(self.device)
-
-        # 3. Call the plotting function
-        if self.kind in [MLTaskKeys.REGRESSION, MLTaskKeys.BINARY_CLASSIFICATION, MLTaskKeys.MULTICLASS_CLASSIFICATION]:
-            shap_summary_plot(
-                model=self.model,
-                background_data=background_data,
-                instances_to_explain=instances_to_explain,
-                feature_names=feature_names,
-                save_dir=save_dir,
-                explainer_type=explainer_type,
-                device=self.device
-            )
-        # DEPRECATED: Multi-target SHAP support is unstable; recommend Captum instead.
-        elif self.kind in [MLTaskKeys.MULTITARGET_REGRESSION, MLTaskKeys.MULTILABEL_BINARY_CLASSIFICATION]:
-            # try to get target names
-            if target_names is None:
-                target_names = []
-                if hasattr(target_dataset, DatasetKeys.TARGET_NAMES):
-                    target_names = target_dataset.target_names # type: ignore
-                else:
-                    # Infer number of targets from the model's output layer
-                    try:
-                        num_targets = self.model.output_layer.out_features # type: ignore
-                        target_names = [f"target_{i}" for i in range(num_targets)] # type: ignore
-                        _LOGGER.warning("Dataset has no 'target_names' attribute. Using generic names.")
-                    except AttributeError:
-                        _LOGGER.error("Cannot determine target names for multi-target SHAP plot. Skipping.")
-                        return
-
-            multi_target_shap_summary_plot(
-                model=self.model,
-                background_data=background_data,
-                instances_to_explain=instances_to_explain,
-                feature_names=feature_names, # type: ignore
-                target_names=target_names, # type: ignore
-                save_dir=save_dir,
-                explainer_type=explainer_type,
-                device=self.device
-            )
-
-    def explain_captum(self,
-                       save_dir: Union[str, Path],
-                       explain_dataset: Optional[Dataset] = None,
-                       n_samples: int = 100,
-                       feature_names: Optional[List[str]] = None,
-                       target_names: Optional[List[str]] = None,
-                       n_steps: int = 50):
-        """
-        Explains model predictions using Captum's Integrated Gradients.
-        
-        - **Tabular/Classification:** Generates Feature Importance Bar Charts.
-        - **Segmentation:** Generates Spatial Heatmaps for each class.
-        
-        Args:
-            save_dir (str | Path): Directory to save artifacts.
-            explain_dataset (Dataset | None): Dataset to sample from. Defaults to validation set.
-            n_samples (int): Number of samples to evaluate.
-            feature_names (list[str] | None): Feature names. 
-                - Required for Tabular tasks.
-                - Ignored/Optional for Image tasks (defaults to Channel names).
-            target_names (list[str] | None): Names for the model outputs (or Class names).
-                - If None, attempts to extract from dataset attributes (`target_names`, `classes`, or `class_map`).
-                - If extraction fails, generates generic names (e.g. "Output_0").
-            n_steps (int): Number of interpolation steps.
-        """
-        # 1. Check availability
-        if not _is_captum_available():
-            _LOGGER.error("Captum is not installed or could not be imported.")
-            return
-         
-        # 2. Prepare Data
-        dataset_to_use = explain_dataset if explain_dataset is not None else self.validation_dataset
-        if dataset_to_use is None:
-            _LOGGER.error("No dataset available for explanation.")
-            return
-
-        # Efficient sampling helper
-        def _get_samples(ds, n):
-            # Use num_workers=0 for stability during ad-hoc sampling
-            loader = DataLoader(ds, batch_size=n, shuffle=True, num_workers=0)
-            data_iter = iter(loader)
-            features, targets = next(data_iter)
-            return features, targets
-
-        input_data, _ = _get_samples(dataset_to_use, n_samples)
-        
-        # 3. Get Feature Names (Only if NOT segmentation AND NOT image classification)
-        # Image tasks generally don't have explicit feature names; Captum will default to "Channel_X"
-        is_segmentation = self.kind in [MLTaskKeys.BINARY_SEGMENTATION, MLTaskKeys.MULTICLASS_SEGMENTATION]
-        is_image_classification = self.kind in [MLTaskKeys.BINARY_IMAGE_CLASSIFICATION, MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION]
-        
-        if feature_names is None and not is_segmentation and not is_image_classification:
-            if hasattr(dataset_to_use, DatasetKeys.FEATURE_NAMES):
-                feature_names = dataset_to_use.feature_names # type: ignore
-            else:
-                _LOGGER.error(f"Could not extract `feature_names`. It must be provided if the dataset does not have it.")
-                raise ValueError()
-
-        # 4. Handle Target Names (or Class Names)
-        if target_names is None:
-            # A. Try dataset attributes first
-            if hasattr(dataset_to_use, DatasetKeys.TARGET_NAMES):
-                target_names = dataset_to_use.target_names # type: ignore
-            elif hasattr(dataset_to_use, "classes"): 
-                 target_names = dataset_to_use.classes # type: ignore
-            elif hasattr(dataset_to_use, "class_map") and isinstance(dataset_to_use.class_map, dict): # type: ignore
-                 # Sort by value (index) to ensure correct order: {name: index} -> [name_at_0, name_at_1...]
-                 sorted_items = sorted(dataset_to_use.class_map.items(), key=lambda item: item[1]) # type: ignore
-                 target_names = [k for k, v in sorted_items]
-
-            # B. Infer based on task
-            if target_names is None:
-                if self.kind in [MLTaskKeys.REGRESSION, MLTaskKeys.BINARY_CLASSIFICATION, MLTaskKeys.BINARY_IMAGE_CLASSIFICATION]:
-                    target_names = ["Output"]
-                elif self.kind == MLTaskKeys.BINARY_SEGMENTATION:
-                    target_names = ["Foreground"]
-                
-                # For multiclass/multitarget without names, leave it None and let the evaluation function generate generics.
-
-        # 5. Dispatch based on Task
-        if is_segmentation:
-            # lower n_steps for segmentation to save memory
-            if n_steps > 30:
-                n_steps = 30
-                _LOGGER.warning(f"Segmentation task detected: Reducing Captum n_steps to {n_steps} to prevent OOM. If you encounter OOM errors, consider lowering this further.")
-            
-            captum_segmentation_heatmap(
-                model=self.model,
-                input_data=input_data,
-                save_dir=save_dir,
-                target_names=target_names, # Can be None, helper handles it
-                n_steps=n_steps,
-                device=self.device
-            )
-        
-        elif is_image_classification:
-            captum_image_heatmap(
-                model=self.model,
-                input_data=input_data,
-                save_dir=save_dir,
-                target_names=target_names,
-                n_steps=n_steps,
-                device=self.device
-            )
-            
-        else:
-            # Standard Tabular/Image Classification
-            captum_feature_importance(
-                model=self.model,
-                input_data=input_data,
-                feature_names=feature_names,
-                save_dir=save_dir,
-                target_names=target_names,
-                n_steps=n_steps,
-                device=self.device
-            )
-
-    def _attention_helper(self, dataloader: DataLoader):
-        """
-        Private method to yield model attention weights batch by batch for evaluation.
-
-        Args:
-            dataloader (DataLoader): The dataloader to predict on.
-
-        Yields:
-            (torch.Tensor): Attention weights
-        """
-        self.model.eval()
-        self.model.to(self.device)
-        
-        with torch.no_grad():
-            for features, target in dataloader:
-                features = features.to(self.device)
-                attention_weights = None
-                
-                # Get model output
-                # Unpack logits and weights from the special forward method
-                _output, attention_weights = self.model.forward_attention(features) # type: ignore
-                
-                if attention_weights is not None:
-                    attention_weights = attention_weights.cpu()
-
-                yield attention_weights
-    
-    def explain_attention(self, save_dir: Union[str, Path], 
-                          feature_names: Optional[List[str]] = None, 
-                          explain_dataset: Optional[Dataset] = None,
-                          plot_n_features: int = 10):
-        """
-        Generates and saves a feature importance plot based on attention weights.
-
-        This method only works for models with models with 'has_interpretable_attention'.
-
-        Args:
-            save_dir (str | Path): Directory to save the plot and summary data.
-            feature_names (List[str] | None): Names for the features for plot labeling. If None, the names will be extracted from the Dataset and raise an error on failure.
-            explain_dataset (Dataset, optional): A specific dataset to explain. If None, the trainer's test dataset is used.
-            plot_n_features (int): Number of top features to plot.
-        """
-        
-        # print("\n--- Attention Analysis ---")
-        
-        # --- Step 1: Check if the model supports this explanation ---
-        if not getattr(self.model, 'has_interpretable_attention', False):
-            _LOGGER.warning("Model is not compatible with interpretable attention analysis. Skipping.")
-            return
-
-        # --- Step 2: Set up the dataloader ---
-        dataset_to_use = explain_dataset if explain_dataset is not None else self.validation_dataset
-        if not isinstance(dataset_to_use, Dataset):
-            _LOGGER.error("The explanation dataset is empty or invalid. Skipping attention analysis.")
-            return
-        
-        # Get feature names
-        if feature_names is None:
-            if hasattr(dataset_to_use, DatasetKeys.FEATURE_NAMES):
-                feature_names = dataset_to_use.feature_names # type: ignore
-            else:
-                _LOGGER.error(f"Could not extract `feature_names` from the dataset for attention plot. It must be provided if the dataset object does not have a '{DatasetKeys.FEATURE_NAMES}' attribute.")
-                raise ValueError()
-        
-        explain_loader = DataLoader(
-            dataset=dataset_to_use, batch_size=32, shuffle=False,
-            num_workers=0 if self.device.type == 'mps' else self.dataloader_workers,
-            pin_memory=("cuda" in self.device.type)
-        )
-        
-        # --- Step 3: Collect weights ---
-        all_weights = []
-        for att_weights_b in self._attention_helper(explain_loader):
-            if att_weights_b is not None:
-                all_weights.append(att_weights_b)
-
-        # --- Step 4: Call the plotting function ---
-        if all_weights:
-            plot_attention_importance(
-                weights=all_weights,
-                feature_names=feature_names,
-                save_dir=save_dir,
-                top_n=plot_n_features
-            )
-        else:
-            _LOGGER.error("No attention weights were collected from the model.")
-        
-    def finalize_model_training(self, 
-                                model_checkpoint: Union[Path, Literal['best', 'current']],
-                                save_dir: Union[str, Path], 
-                                finalize_config: Union[FinalizeRegression,
-                                                       FinalizeMultiTargetRegression,
-                                                       FinalizeBinaryClassification,
-                                                       FinalizeBinaryImageClassification,
-                                                       FinalizeMultiClassClassification,
-                                                       FinalizeMultiClassImageClassification,
-                                                       FinalizeBinarySegmentation,
-                                                       FinalizeMultiClassSegmentation,
-                                                       FinalizeMultiLabelBinaryClassification]):
-        """
-        Saves a finalized, "inference-ready" model state to a .pth file.
-
-        This method saves the model's `state_dict`, the final epoch number, and optional configuration for the task at hand.
-
-        Args:
-            model_checkpoint (Path | "best" | "current"):
-                - Path: Loads the model state from a specific checkpoint file.
-                - "best": Loads the best model state saved by the `DragonModelCheckpoint` callback.
-                - "current": Uses the model's state as it is.
-            save_dir (str | Path): The directory to save the finalized model.
-            finalize_config (object): A data class instance specific to the ML task containing task-specific metadata required for inference.
-        """
-        if self.kind == MLTaskKeys.REGRESSION and not isinstance(finalize_config, FinalizeRegression):
-            _LOGGER.error(f"For task {self.kind}, expected finalize_config of type 'FinalizeRegression', but got {type(finalize_config).__name__}.")
-            raise TypeError()
-        elif self.kind == MLTaskKeys.MULTITARGET_REGRESSION and not isinstance(finalize_config, FinalizeMultiTargetRegression):
-            _LOGGER.error(f"For task {self.kind}, expected finalize_config of type 'FinalizeMultiTargetRegression', but got {type(finalize_config).__name__}.")
-            raise TypeError()
-        elif self.kind == MLTaskKeys.BINARY_CLASSIFICATION and not isinstance(finalize_config, FinalizeBinaryClassification):
-            _LOGGER.error(f"For task {self.kind}, expected finalize_config of type 'FinalizeBinaryClassification', but got {type(finalize_config).__name__}.")
-            raise TypeError()
-        elif self.kind == MLTaskKeys.BINARY_IMAGE_CLASSIFICATION and not isinstance(finalize_config, FinalizeBinaryImageClassification):
-            _LOGGER.error(f"For task {self.kind}, expected finalize_config of type 'FinalizeBinaryImageClassification', but got {type(finalize_config).__name__}.")
-            raise TypeError()
-        elif self.kind == MLTaskKeys.MULTICLASS_CLASSIFICATION and not isinstance(finalize_config, FinalizeMultiClassClassification):
-            _LOGGER.error(f"For task {self.kind}, expected finalize_config of type 'FinalizeMultiClassClassification', but got {type(finalize_config).__name__}.")
-            raise TypeError()
-        elif self.kind == MLTaskKeys.MULTICLASS_IMAGE_CLASSIFICATION and not isinstance(finalize_config, FinalizeMultiClassImageClassification):
-            _LOGGER.error(f"For task {self.kind}, expected finalize_config of type 'FinalizeMultiClassImageClassification', but got {type(finalize_config).__name__}.")
-            raise TypeError()
-        elif self.kind == MLTaskKeys.BINARY_SEGMENTATION and not isinstance(finalize_config, FinalizeBinarySegmentation):
-            _LOGGER.error(f"For task {self.kind}, expected finalize_config of type 'FinalizeBinarySegmentation', but got {type(finalize_config).__name__}.")
-            raise TypeError()
-        elif self.kind == MLTaskKeys.MULTICLASS_SEGMENTATION and not isinstance(finalize_config, FinalizeMultiClassSegmentation):
-            _LOGGER.error(f"For task {self.kind}, expected finalize_config of type 'FinalizeMultiClassSegmentation', but got {type(finalize_config).__name__}.")
-            raise TypeError()
-        elif self.kind == MLTaskKeys.MULTILABEL_BINARY_CLASSIFICATION and not isinstance(finalize_config, FinalizeMultiLabelBinaryClassification):
-            _LOGGER.error(f"For task {self.kind}, expected finalize_config of type 'FinalizeMultiLabelBinaryClassification', but got {type(finalize_config).__name__}.")
-            raise TypeError()
-        
-        # handle save path
-        dir_path = make_fullpath(save_dir, make=True, enforce="directory")
-        full_path = dir_path / finalize_config.filename
-        
-        # handle checkpoint
-        self._load_model_state_for_finalizing(model_checkpoint)
-        
-        # Create finalized data
-        finalized_data = {
-            PyTorchCheckpointKeys.EPOCH: self.epoch,
-            PyTorchCheckpointKeys.MODEL_STATE: self.model.state_dict(),
-            PyTorchCheckpointKeys.TASK: finalize_config.task
-        }
-
-        # Parse config
-        if finalize_config.target_name is not None:
-            finalized_data[PyTorchCheckpointKeys.TARGET_NAME] = finalize_config.target_name
-        if finalize_config.target_names is not None:
-            finalized_data[PyTorchCheckpointKeys.TARGET_NAMES] = finalize_config.target_names
-        if finalize_config.classification_threshold is not None:
-            finalized_data[PyTorchCheckpointKeys.CLASSIFICATION_THRESHOLD] = finalize_config.classification_threshold
-        if finalize_config.class_map is not None:
-            finalized_data[PyTorchCheckpointKeys.CLASS_MAP] = finalize_config.class_map
-
-        # Save model file
-        torch.save(finalized_data, full_path)
-        
-        _LOGGER.info(f"Finalized model file saved to '{full_path}'")
-
-
-# Object Detection Trainer
-class DragonDetectionTrainer(_BaseDragonTrainer):
-    def __init__(self, model: nn.Module, 
-                 train_dataset: Dataset, 
-                 validation_dataset: Dataset, 
-                 collate_fn: Callable, optimizer: torch.optim.Optimizer, 
-                 device: Union[Literal['cuda', 'mps', 'cpu'],str], 
-                 checkpoint_callback: Optional[DragonModelCheckpoint],
-                 early_stopping_callback: Optional[_DragonEarlyStopping],
-                 lr_scheduler_callback: Optional[_DragonLRScheduler],
-                 extra_callbacks: Optional[List[_Callback]] = None,
-                 dataloader_workers: int = 2):
-        """
-        Automates the training process of an Object Detection Model (e.g., DragonFastRCNN).
-        
-        Built-in Callbacks: `History`, `TqdmProgressBar`
-
-        Args:
-            model (nn.Module): The PyTorch object detection model to train.
-            train_dataset (Dataset): The training dataset.
-            validation_dataset (Dataset): The testing/validation dataset.
-            collate_fn (Callable): The collate function from `ObjectDetectionDatasetMaker.collate_fn`.
-            optimizer (torch.optim.Optimizer): The optimizer.
-            device (str): The device to run training on ('cpu', 'cuda', 'mps').
-            dataloader_workers (int): Subprocesses for data loading.
-            checkpoint_callback (DragonModelCheckpoint | None): Callback to save the model.
-            early_stopping_callback (DragonEarlyStopping | None): Callback to stop training early.
-            lr_scheduler_callback (DragonLRScheduler | None): Callback to manage the LR scheduler.
-            extra_callbacks (List[Callback] | None): A list of extra callbacks to use during training.
-            
-        ## Note:
-            This trainer is specialized. It does not take a `criterion` because object detection models like Faster R-CNN return a dictionary of losses directly from their forward pass during training.
-        """
-        # Call the base class constructor with common parameters
-        super().__init__(
-            model=model,
-            optimizer=optimizer,
-            device=device,
-            dataloader_workers=dataloader_workers,
-            checkpoint_callback=checkpoint_callback,
-            early_stopping_callback=early_stopping_callback,
-            lr_scheduler_callback=lr_scheduler_callback,
-            extra_callbacks=extra_callbacks
-        )
-        
-        self.train_dataset = train_dataset
-        self.validation_dataset = validation_dataset # <-- Renamed
-        self.kind = MLTaskKeys.OBJECT_DETECTION
-        self.collate_fn = collate_fn
-        self.criterion = None # Criterion is handled inside the model
-
-    def _create_dataloaders(self, batch_size: int, shuffle: bool):
-        """Initializes the DataLoaders with the object detection collate_fn."""
-        # Ensure stability on MPS devices by setting num_workers to 0
-        loader_workers = 0 if self.device.type == 'mps' else self.dataloader_workers
-        
-        self.train_loader = DataLoader(
-            dataset=self.train_dataset, 
-            batch_size=batch_size, 
-            shuffle=shuffle, 
-            num_workers=loader_workers, 
-            pin_memory=("cuda" in self.device.type), 
-            collate_fn=self.collate_fn, # Use the provided collate function
-            drop_last=True 
-        )
-        
-        self.validation_loader = DataLoader(
-            dataset=self.validation_dataset, 
-            batch_size=batch_size, 
-            shuffle=False, 
-            num_workers=loader_workers, 
-            pin_memory=("cuda" in self.device.type),
-            collate_fn=self.collate_fn # Use the provided collate function
-        )
-
-    def _train_step(self):
-        self.model.train()
-        running_loss = 0.0
-        total_samples = 0
-        
-        for batch_idx, (images, targets) in enumerate(self.train_loader): # type: ignore
-            # images is a tuple of tensors, targets is a tuple of dicts
-            batch_size = len(images)
-            
-            # Create a log dictionary for the batch
-            batch_logs = {
-                PyTorchLogKeys.BATCH_INDEX: batch_idx, 
-                PyTorchLogKeys.BATCH_SIZE: batch_size
-            }
-            self._callbacks_hook('on_batch_begin', batch_idx, logs=batch_logs)
-
-            # Move data to device
-            images = list(img.to(self.device) for img in images)
-            targets = [{k: v.to(self.device) for k, v in t.items()} for t in targets]
-            
-            self.optimizer.zero_grad()
-            
-            # Model returns a loss dict when in train() mode and targets are passed
-            loss_dict = self.model(images, targets)
-            
-            if not loss_dict:
-                # No losses returned, skip batch
-                _LOGGER.warning(f"Model returned no losses for batch {batch_idx}. Skipping.")
-                batch_logs[PyTorchLogKeys.BATCH_LOSS] = 0
-                self._callbacks_hook('on_batch_end', batch_idx, logs=batch_logs)
-                continue
-            
-            # Sum all losses
-            loss: torch.Tensor = sum(l for l in loss_dict.values()) # type: ignore
-            
-            loss.backward()
-            self.optimizer.step()
-
-            # Calculate batch loss and update running loss for the epoch
-            batch_loss = loss.item()
-            running_loss += batch_loss * batch_size
-            total_samples += batch_size # <-- Accumulate total samples
-            
-            # Add the batch loss to the logs and call the end-of-batch hook
-            batch_logs[PyTorchLogKeys.BATCH_LOSS] = batch_loss # type: ignore
-            self._callbacks_hook('on_batch_end', batch_idx, logs=batch_logs)
-        
-        # Calculate loss using the correct denominator
-        if total_samples == 0:
-            _LOGGER.warning("No samples processed in _train_step. Returning 0 loss.")
-            return {PyTorchLogKeys.TRAIN_LOSS: 0.0}
-
-        return {PyTorchLogKeys.TRAIN_LOSS: running_loss / total_samples}
-
-    def _validation_step(self):
-        self.model.train() # Set to train mode even for validation loss calculation
-                           # as model internals (e.g., proposals) might differ, but we still need loss_dict.
-                           # use torch.no_grad() to prevent gradient updates.
-        running_loss = 0.0
-        total_samples = 0 
-        
-        with torch.no_grad():
-            for images, targets in self.validation_loader: # type: ignore
-                batch_size = len(images)
-                
-                # Move data to device
-                images = list(img.to(self.device) for img in images)
-                targets = [{k: v.to(self.device) for k, v in t.items()} for t in targets]
-                
-                # Get loss dict
-                loss_dict = self.model(images, targets)
-                
-                if not loss_dict:
-                    _LOGGER.warning("Model returned no losses during validation step. Skipping batch.")
-                    continue # Skip if no losses
-                
-                # Sum all losses
-                loss: torch.Tensor = sum(l for l in loss_dict.values()) # type: ignore
-                
-                running_loss += loss.item() * batch_size
-                total_samples += batch_size # <-- Accumulate total samples
-        
-        # Calculate loss using the correct denominator
-        if total_samples == 0:
-            _LOGGER.warning("No samples processed in _validation_step. Returning 0 loss.")
-            return {PyTorchLogKeys.VAL_LOSS: 0.0}
-        
-        logs = {PyTorchLogKeys.VAL_LOSS: running_loss / total_samples}
-        return logs
-    
-    def evaluate(self, 
-                 save_dir: Union[str, Path], 
-                 model_checkpoint: Union[Path, Literal["best", "current"]],
-                 test_data: Optional[Union[DataLoader, Dataset]] = None):
-        """
-        Evaluates the model using object detection mAP metrics.
-
-        Args:
-            save_dir (str | Path): Directory to save all reports and plots.
-            model_checkpoint (Path | "best" | "current"): 
-                - Path to a valid checkpoint for the model. The state of the trained model will be overwritten in place.
-                - If 'best', the best checkpoint will be loaded if a DragonModelCheckpoint was provided. The state of the trained model will be overwritten in place.
-                - If 'current', use the current state of the trained model up the latest trained epoch.
-            test_data (DataLoader | Dataset | None): Optional Test data to evaluate the model performance. Validation and Test metrics will be saved to subdirectories.
-        """
-        # Validate model checkpoint
-        if isinstance(model_checkpoint, Path):
-            checkpoint_validated = make_fullpath(model_checkpoint, enforce="file")
-        elif model_checkpoint in [MagicWords.BEST, MagicWords.CURRENT]:
-            checkpoint_validated = model_checkpoint
-        else:
-            _LOGGER.error(f"'model_checkpoint' must be a Path object, or the string '{MagicWords.BEST}', or the string '{MagicWords.CURRENT}'.")
-            raise ValueError()
-        
-        # Validate directory
-        save_path = make_fullpath(save_dir, make=True, enforce="directory")
-        
-        # Validate test data and dispatch
-        if test_data is not None:
-            if not isinstance(test_data, (DataLoader, Dataset)):
-                _LOGGER.error(f"Invalid type for 'test_data': '{type(test_data)}'.")
-                raise ValueError()
-            test_data_validated = test_data
-            
-            validation_metrics_path = save_path / DragonTrainerKeys.VALIDATION_METRICS_DIR
-            test_metrics_path = save_path / DragonTrainerKeys.TEST_METRICS_DIR
-            
-            # Dispatch validation set
-            _LOGGER.info(f"Evaluating on validation dataset. Metrics will be saved to '{DragonTrainerKeys.VALIDATION_METRICS_DIR}'")
-            self._evaluate(save_dir=validation_metrics_path,
-                           model_checkpoint=checkpoint_validated,
-                           data=None) # 'None' triggers use of self.test_dataset
-            
-            # Dispatch test set
-            _LOGGER.info(f"Evaluating on test dataset. Metrics will be saved to '{DragonTrainerKeys.TEST_METRICS_DIR}'")
-            self._evaluate(save_dir=test_metrics_path,
-                           model_checkpoint="current", # Use 'current' state after loading checkpoint once
-                           data=test_data_validated)
-        else:
-            # Dispatch validation set
-            _LOGGER.info(f"Evaluating on validation dataset. Metrics will be saved to '{save_path.name}'")
-            self._evaluate(save_dir=save_path,
-                           model_checkpoint=checkpoint_validated,
-                           data=None) # 'None' triggers use of self.test_dataset
-    
-    def _evaluate(self, 
-                 save_dir: Union[str, Path], 
-                 model_checkpoint: Union[Path, Literal["best", "current"]],
-                 data: Optional[Union[DataLoader, Dataset]]):
-        """
-        Changed to a private helper method
-        Evaluates the model using object detection mAP metrics.
-
-        Args:
-            save_dir (str | Path): Directory to save all reports and plots.
-            data (DataLoader | Dataset | None): The data to evaluate on. If None, defaults to the trainer's internal test_dataset.
-            model_checkpoint ('auto' | Path | None): 
-                - Path to a valid checkpoint for the model. The state of the trained model will be overwritten in place.
-                - If 'best', the best checkpoint will be loaded if a DragonModelCheckpoint was provided. The state of the trained model will be overwritten in place.
-                - If 'current', use the current state of the trained model up the latest trained epoch.
-        """
-        dataset_for_artifacts = None
-        eval_loader = None
-        
-        # load model checkpoint
-        if isinstance(model_checkpoint, Path):
-            self._load_checkpoint(path=model_checkpoint)
-        elif model_checkpoint == MagicWords.BEST and self._checkpoint_callback:
-            path_to_latest = self._checkpoint_callback.best_checkpoint_path
-            self._load_checkpoint(path_to_latest)
-        elif model_checkpoint == MagicWords.BEST and self._checkpoint_callback is None:
-            _LOGGER.error(f"'model_checkpoint' set to '{MagicWords.BEST}' but no checkpoint callback was found.")
-            raise ValueError()
-
-        # Dataloader
-        if isinstance(data, DataLoader):
-            eval_loader = data
-            if hasattr(data, 'dataset'):
-                dataset_for_artifacts = data.dataset # type: ignore
-        elif isinstance(data, Dataset):
-            # Create a new loader from the provided dataset
-            eval_loader = DataLoader(data, 
-                                     batch_size=self._batch_size, 
-                                     shuffle=False, 
-                                     num_workers=0 if self.device.type == 'mps' else self.dataloader_workers,
-                                     pin_memory=(self.device.type == "cuda"),
-                                     collate_fn=self.collate_fn)
-            dataset_for_artifacts = data
-        else: # data is None, use the trainer's default test dataset
-            if self.validation_dataset is None:
-                _LOGGER.error("Cannot evaluate. No data provided and no test_dataset available in the trainer.")
-                raise ValueError()
-            # Create a fresh DataLoader from the test_dataset
-            eval_loader = DataLoader(
-                self.validation_dataset, 
-                batch_size=self._batch_size, 
-                shuffle=False, 
-                num_workers=0 if self.device.type == 'mps' else self.dataloader_workers,
-                pin_memory=(self.device.type == "cuda"),
-                collate_fn=self.collate_fn
-            )
-            dataset_for_artifacts = self.validation_dataset
-
-        if eval_loader is None:
-            _LOGGER.error("Cannot evaluate. No valid data was provided or found.")
-            raise ValueError()
-
-        # print("\n--- Model Evaluation ---")
-
-        all_predictions = []
-        all_targets = []
-        
-        self.model.eval() # Set model to evaluation mode
-        self.model.to(self.device)
-        
-        with torch.no_grad():
-            for images, targets in eval_loader:
-                # Move images to device
-                images = list(img.to(self.device) for img in images)
-                
-                # Model returns predictions when in eval() mode
-                predictions = self.model(images)
-                
-                # Move predictions and targets to CPU for aggregation
-                cpu_preds = [{k: v.to('cpu') for k, v in p.items()} for p in predictions]
-                cpu_targets = [{k: v.to('cpu') for k, v in t.items()} for t in targets]
-                
-                all_predictions.extend(cpu_preds)
-                all_targets.extend(cpu_targets)
-
-        if not all_targets:
-            _LOGGER.error("Evaluation failed: No data was processed.")
-            return
-        
-        # Get class names from the dataset for the report
-        class_names = None
-        try:
-            # Try to get 'classes' from ObjectDetectionDatasetMaker
-            if hasattr(dataset_for_artifacts, 'classes'):
-                class_names = dataset_for_artifacts.classes # type: ignore
-            # Fallback for Subset
-            elif hasattr(dataset_for_artifacts, 'dataset') and hasattr(dataset_for_artifacts.dataset, 'classes'): # type: ignore
-                 class_names = dataset_for_artifacts.dataset.classes # type: ignore
-        except AttributeError:
-            _LOGGER.warning("Could not find 'classes' attribute on dataset. Per-class metrics will not be named.")
-            pass # class_names is still None
-
-        # --- Routing Logic ---
-        object_detection_metrics(
-            preds=all_predictions, 
-            targets=all_targets, 
-            save_dir=save_dir,
-            class_names=class_names,
-            print_output=False
-        )
-    
-    def finalize_model_training(self, 
-                                save_dir: Union[str, Path], 
-                                model_checkpoint: Union[Path, Literal['best', 'current']],
-                                finalize_config: FinalizeObjectDetection
-                                ):
-        """
-        Saves a finalized, "inference-ready" model state to a .pth file.
-
-        This method saves the model's `state_dict` and the final epoch number.
-
-        Args:
-            save_dir (Union[str, Path]): The directory to save the finalized model.
-            model_checkpoint (Union[Path, Literal["best", "current"]]):
-                - Path: Loads the model state from a specific checkpoint file.
-                - "best": Loads the best model state saved by the `DragonModelCheckpoint` callback.
-                - "current": Uses the model's state as it is.
-            finalize_config (FinalizeObjectDetection): A data class instance specific to the ML task containing task-specific metadata required for inference.
-        """
-        if not isinstance(finalize_config, FinalizeObjectDetection):
-            _LOGGER.error(f"For task {self.kind}, expected finalize_config of type 'FinalizeObjectDetection', but got {type(finalize_config).__name__}.")
-            raise TypeError()
-        
-        # handle save path
-        dir_path = make_fullpath(save_dir, make=True, enforce="directory")
-        full_path = dir_path / finalize_config.filename
-        
-        # handle checkpoint
-        self._load_model_state_for_finalizing(model_checkpoint)
-        
-        # Create finalized data
-        finalized_data = {
-            PyTorchCheckpointKeys.EPOCH: self.epoch,
-            PyTorchCheckpointKeys.MODEL_STATE: self.model.state_dict(),
-            PyTorchCheckpointKeys.TASK: finalize_config.task
-        }
-        
-        if finalize_config.class_map is not None:
-            finalized_data[PyTorchCheckpointKeys.CLASS_MAP] = finalize_config.class_map
-        
-        torch.save(finalized_data, full_path)
-        
-        _LOGGER.info(f"Finalized model file saved to '{full_path}'")
-
-# --- DragonSequenceTrainer ----
-class DragonSequenceTrainer(_BaseDragonTrainer):
-    def __init__(self, 
-                 model: nn.Module, 
-                 train_dataset: Dataset, 
-                 validation_dataset: Dataset, 
-                 kind: Literal["sequence-to-sequence", "sequence-to-value"],
-                 optimizer: torch.optim.Optimizer, 
-                 device: Union[Literal['cuda', 'mps', 'cpu'],str], 
-                 checkpoint_callback: Optional[DragonModelCheckpoint],
-                 early_stopping_callback: Optional[_DragonEarlyStopping],
-                 lr_scheduler_callback: Optional[_DragonLRScheduler],
-                 extra_callbacks: Optional[List[_Callback]] = None,
-                 criterion: Union[nn.Module,Literal["auto"]] = "auto", 
-                 dataloader_workers: int = 2):
-        """
-        Automates the training process of a PyTorch Sequence Model.
-        
-        Built-in Callbacks: `History`, `TqdmProgressBar`
-
-        Args:
-            model (nn.Module): The PyTorch model to train.
-            train_dataset (Dataset): The training dataset.
-            validation_dataset (Dataset): The validation dataset.
-            kind (str): Used to redirect to the correct process ('sequence-to-sequence' or 'sequence-to-value'). 
-            criterion (nn.Module | "auto"): The loss function to use. If "auto", it will be inferred from the selected task
-            optimizer (torch.optim.Optimizer): The optimizer.
-            device (str): The device to run training on ('cpu', 'cuda', 'mps').
-            dataloader_workers (int): Subprocesses for data loading.
-            extra_callbacks (List[Callback] | None): A list of extra callbacks to use during training.
-        """
-        # Call the base class constructor with common parameters
-        super().__init__(
-            model=model,
-            optimizer=optimizer,
-            device=device,
-            dataloader_workers=dataloader_workers,
-            checkpoint_callback=checkpoint_callback,
-            early_stopping_callback=early_stopping_callback,
-            lr_scheduler_callback=lr_scheduler_callback,
-            extra_callbacks=extra_callbacks
-        )
-        
-        if kind not in [MLTaskKeys.SEQUENCE_SEQUENCE, MLTaskKeys.SEQUENCE_VALUE]:
-            raise ValueError(f"'{kind}' is not a valid task type for DragonSequenceTrainer.")
-
-        self.train_dataset = train_dataset
-        self.validation_dataset = validation_dataset
-        self.kind = kind
-        
-        # try to validate against Dragon Sequence model
-        if hasattr(self.model, "prediction_mode"):
-            key_to_check: str = self.model.prediction_mode # type: ignore
-            if not key_to_check == self.kind:
-                _LOGGER.error(f"Trainer was set for '{self.kind}', but model architecture '{self.model}' is built for '{key_to_check}'.")
-                raise RuntimeError()
-        
-        # loss function
-        if criterion == "auto":
-            # Both sequence tasks are treated as regression problems
-            self.criterion = nn.MSELoss()
-        else:
-            self.criterion = criterion
-
-    def _create_dataloaders(self, batch_size: int, shuffle: bool):
-        """Initializes the DataLoaders."""
-        # Ensure stability on MPS devices by setting num_workers to 0
-        loader_workers = 0 if self.device.type == 'mps' else self.dataloader_workers
-        
-        self.train_loader = DataLoader(
-            dataset=self.train_dataset, 
-            batch_size=batch_size, 
-            shuffle=shuffle, 
-            num_workers=loader_workers, 
-            pin_memory=("cuda" in self.device.type),
-            drop_last=True  # Drops the last batch if incomplete, selecting a good batch size is key.
-        )
-        
-        self.validation_loader = DataLoader(
-            dataset=self.validation_dataset, 
-            batch_size=batch_size, 
-            shuffle=False, 
-            num_workers=loader_workers, 
-            pin_memory=("cuda" in self.device.type)
-        )
-
-    def _train_step(self):
-        self.model.train()
-        running_loss = 0.0
-        total_samples = 0
-        
-        for batch_idx, (features, target) in enumerate(self.train_loader): # type: ignore
-            # Create a log dictionary for the batch
-            batch_logs = {
-                PyTorchLogKeys.BATCH_INDEX: batch_idx, 
-                PyTorchLogKeys.BATCH_SIZE: features.size(0)
-            }
-            self._callbacks_hook('on_batch_begin', batch_idx, logs=batch_logs)
-
-            features, target = features.to(self.device), target.to(self.device)
-            self.optimizer.zero_grad()
-            
-            output = self.model(features)
-            
-            # --- Label Type/Shape Correction ---
-            # Ensure target is float for MSELoss
-            target = target.float()
-
-            # For seq-to-val, models might output [N, 1] but target is [N].
-            if self.kind == MLTaskKeys.SEQUENCE_VALUE:
-                if output.ndim == 2 and output.shape[1] == 1 and target.ndim == 1:
-                    output = output.squeeze(1)
-            
-            # For seq-to-seq, models might output [N, Seq, 1] but target is [N, Seq].
-            elif self.kind == MLTaskKeys.SEQUENCE_SEQUENCE:
-                if output.ndim == 3 and output.shape[2] == 1 and target.ndim == 2:
-                    output = output.squeeze(-1)
-            
-            loss = self.criterion(output, target)
-            
-            loss.backward()
-            self.optimizer.step()
-
-            # Calculate batch loss and update running loss for the epoch
-            batch_loss = loss.item()
-            batch_size = features.size(0)
-            running_loss += batch_loss * batch_size  # Accumulate total loss
-            total_samples += batch_size # total samples
-            
-            # Add the batch loss to the logs and call the end-of-batch hook
-            batch_logs[PyTorchLogKeys.BATCH_LOSS] = batch_loss
-            self._callbacks_hook('on_batch_end', batch_idx, logs=batch_logs)
-        
-        if total_samples == 0:
-            _LOGGER.warning("No samples processed in a train_step. Returning 0 loss.")
-            return {PyTorchLogKeys.TRAIN_LOSS: 0.0}
-
-        return {PyTorchLogKeys.TRAIN_LOSS: running_loss / total_samples} # type: ignore
-
-    def _validation_step(self):
-        self.model.eval()
-        running_loss = 0.0
-        
-        with torch.no_grad():
-            for features, target in self.validation_loader: # type: ignore
-                features, target = features.to(self.device), target.to(self.device)
-                
-                output = self.model(features)
-                
-                # --- Label Type/Shape Correction ---
-                target = target.float()
-                
-                # For seq-to-val, models might output [N, 1] but target is [N].
-                if self.kind == MLTaskKeys.SEQUENCE_VALUE:
-                    if output.ndim == 2 and output.shape[1] == 1 and target.ndim == 1:
-                        output = output.squeeze(1)
-                        
-                # For seq-to-seq, models might output [N, Seq, 1] but target is [N, Seq].
-                elif self.kind == MLTaskKeys.SEQUENCE_SEQUENCE:
-                    if output.ndim == 3 and output.shape[2] == 1 and target.ndim == 2:
-                        output = output.squeeze(-1)
-                
-                loss = self.criterion(output, target)
-                
-                running_loss += loss.item() * features.size(0)
-                
-        if not self.validation_loader.dataset: # type: ignore
-            _LOGGER.warning("No samples processed in _validation_step. Returning 0 loss.")
-            return {PyTorchLogKeys.VAL_LOSS: 0.0}
-        
-        logs = {PyTorchLogKeys.VAL_LOSS: running_loss / len(self.validation_loader.dataset)} # type: ignore
-        return logs
-    
-    def _predict_for_eval(self, dataloader: DataLoader):
-        """
-        Private method to yield model predictions batch by batch for evaluation.
-        
-        Automatically checks for 'scaler'.
-        
-        Yields:
-            tuple: A tuple containing (y_pred_batch, y_prob_batch, y_true_batch).
-                   y_prob_batch is always None for sequence tasks.
-        """
-        self.model.eval()
-        self.model.to(self.device)
-        
-        # --- Check for Scaler ---
-        # DragonDatasetSequence stores it as 'scaler'
-        scaler = None
-        if hasattr(self.train_dataset, ScalerKeys.TARGET_SCALER):
-            scaler = getattr(self.train_dataset, ScalerKeys.TARGET_SCALER)
-            if scaler is not None:
-                _LOGGER.debug("Sequence scaler detected. Un-scaling predictions and targets.")
-        
-        with torch.no_grad():
-            for features, target in dataloader:
-                features = features.to(self.device)
-                target = target.to(self.device)
-                
-                output = self.model(features)
-
-                # --- Automatic Un-scaling Logic ---
-                if scaler:
-                    # 1. Reshape for scaler (N, 1) or (N*Seq, 1)
-                    original_out_shape = output.shape
-                    original_target_shape = target.shape
-                    
-                    # Flatten sequence dims
-                    output_flat = output.reshape(-1, 1)
-                    target_flat = target.reshape(-1, 1)
-                    
-                    # 2. Inverse Transform
-                    output_flat = scaler.inverse_transform(output_flat)
-                    target_flat = scaler.inverse_transform(target_flat)
-                    
-                    # 3. Restore
-                    output = output_flat.reshape(original_out_shape)
-                    target = target_flat.reshape(original_target_shape)
-
-                # Move to CPU
-                y_pred_batch = output.cpu().numpy()
-                y_true_batch = target.cpu().numpy()
-                y_prob_batch = None
-
-                yield y_pred_batch, y_prob_batch, y_true_batch
-                
-    def evaluate(self, 
-                 save_dir: Union[str, Path], 
-                 model_checkpoint: Union[Path, Literal["best", "current"]],
-                 test_data: Optional[Union[DataLoader, Dataset]] = None,
-                 val_format_configuration: Optional[Union[SequenceValueMetricsFormat, 
-                                                          SequenceSequenceMetricsFormat]]=None,
-                 test_format_configuration: Optional[Union[SequenceValueMetricsFormat, 
-                                                           SequenceSequenceMetricsFormat]]=None):
-        """
-        Evaluates the model, routing to the correct evaluation function.
-
-        Args:
-            model_checkpoint (Path | "best" | "current"): 
-                - Path to a valid checkpoint for the model.
-                - If 'best', the best checkpoint will be loaded.
-                - If 'current', use the current state of the trained model.
-            save_dir (str | Path): Directory to save all reports and plots.
-            test_data (DataLoader | Dataset | None): Optional Test data.
-            val_format_configuration: Optional configuration for validation metrics.
-            test_format_configuration: Optional configuration for test metrics.
-        """
-        # Validate model checkpoint
-        if isinstance(model_checkpoint, Path):
-            checkpoint_validated = make_fullpath(model_checkpoint, enforce="file")
-        elif model_checkpoint in [MagicWords.BEST, MagicWords.CURRENT]:
-            checkpoint_validated = model_checkpoint
-        else:
-            _LOGGER.error(f"'model_checkpoint' must be a Path object, or '{MagicWords.BEST}', or '{MagicWords.CURRENT}'.")
-            raise ValueError()
-        
-        # Validate val configuration
-        if val_format_configuration is not None:
-            if not isinstance(val_format_configuration, (SequenceValueMetricsFormat, SequenceSequenceMetricsFormat)):
-                _LOGGER.error(f"Invalid 'val_format_configuration': '{type(val_format_configuration)}'.")
-                raise ValueError()
-        
-        # Validate directory
-        save_path = make_fullpath(save_dir, make=True, enforce="directory")
-        
-        # Validate test data and dispatch
-        if test_data is not None:
-            if not isinstance(test_data, (DataLoader, Dataset)):
-                _LOGGER.error(f"Invalid type for 'test_data': '{type(test_data)}'.")
-                raise ValueError()
-            test_data_validated = test_data
-    
-            validation_metrics_path = save_path / DragonTrainerKeys.VALIDATION_METRICS_DIR
-            test_metrics_path = save_path / DragonTrainerKeys.TEST_METRICS_DIR
-            
-            # Dispatch validation set
-            _LOGGER.info(f"Evaluating on validation dataset. Metrics will be saved to '{DragonTrainerKeys.VALIDATION_METRICS_DIR}'")
-            self._evaluate(save_dir=validation_metrics_path,
-                           model_checkpoint=checkpoint_validated,
-                           data=None,
-                           format_configuration=val_format_configuration)
-            
-            # Validate test configuration
-            test_configuration_validated = None
-            if test_format_configuration is not None:
-                if not isinstance(test_format_configuration, (SequenceValueMetricsFormat, SequenceSequenceMetricsFormat)):
-                    warning_message_type = f"Invalid test_format_configuration': '{type(test_format_configuration)}'."
-                    if val_format_configuration is not None:
-                        warning_message_type += " 'val_format_configuration' will be used."
-                        test_configuration_validated = val_format_configuration
-                    else:
-                        warning_message_type += " Using default format."
-                    _LOGGER.warning(warning_message_type)
-                else:
-                    test_configuration_validated = test_format_configuration
-            
-            # Dispatch test set
-            _LOGGER.info(f"Evaluating on test dataset. Metrics will be saved to '{DragonTrainerKeys.TEST_METRICS_DIR}'")
-            self._evaluate(save_dir=test_metrics_path,
-                           model_checkpoint="current",
-                           data=test_data_validated,
-                           format_configuration=test_configuration_validated)
-        else:
-            # Dispatch validation set
-            _LOGGER.info(f"Evaluating on validation dataset. Metrics will be saved to '{save_path.name}'")
-            self._evaluate(save_dir=save_path,
-                           model_checkpoint=checkpoint_validated,
-                           data=None,
-                           format_configuration=val_format_configuration)
-        
-    def _evaluate(self, 
-                 save_dir: Union[str, Path], 
-                 model_checkpoint: Union[Path, Literal["best", "current"]],
-                 data: Optional[Union[DataLoader, Dataset]],
-                 format_configuration: object):
-        """
-        Private evaluation helper.
-        """
-        eval_loader = None
-        
-        # load model checkpoint
-        if isinstance(model_checkpoint, Path):
-            self._load_checkpoint(path=model_checkpoint)
-        elif model_checkpoint == MagicWords.BEST and self._checkpoint_callback:
-            path_to_latest = self._checkpoint_callback.best_checkpoint_path
-            self._load_checkpoint(path_to_latest)
-        elif model_checkpoint == MagicWords.BEST and self._checkpoint_callback is None:
-            _LOGGER.error(f"'model_checkpoint' set to '{MagicWords.BEST}' but no checkpoint callback was found.")
-            raise ValueError()
-        
-        # Dataloader
-        if isinstance(data, DataLoader):
-            eval_loader = data
-        elif isinstance(data, Dataset):
-            # Create a new loader from the provided dataset
-            eval_loader = DataLoader(data, 
-                                     batch_size=self._batch_size, 
-                                     shuffle=False, 
-                                     num_workers=0 if self.device.type == 'mps' else self.dataloader_workers,
-                                     pin_memory=(self.device.type == "cuda"))
-        else: # data is None, use the trainer's default validation dataset
-            if self.validation_dataset is None:
-                _LOGGER.error("Cannot evaluate. No data provided and no validation_dataset available in the trainer.")
-                raise ValueError()
-            eval_loader = DataLoader(self.validation_dataset, 
-                                     batch_size=self._batch_size, 
-                                     shuffle=False, 
-                                     num_workers=0 if self.device.type == 'mps' else self.dataloader_workers,
-                                     pin_memory=(self.device.type == "cuda"))
-
-        if eval_loader is None:
-            _LOGGER.error("Cannot evaluate. No valid data was provided or found.")
-            raise ValueError()
-
-        all_preds, _, all_true = [], [], []
-        for y_pred_b, y_prob_b, y_true_b in self._predict_for_eval(eval_loader):
-            if y_pred_b is not None: all_preds.append(y_pred_b)
-            if y_true_b is not None: all_true.append(y_true_b)
-
-        if not all_true:
-            _LOGGER.error("Evaluation failed: No data was processed.")
-            return
-
-        y_pred = np.concatenate(all_preds)
-        y_true = np.concatenate(all_true)
-
-        # --- Routing Logic ---
-        if self.kind == MLTaskKeys.SEQUENCE_VALUE:
-            config = None
-            if format_configuration and isinstance(format_configuration, SequenceValueMetricsFormat):
-                config = format_configuration
-            elif format_configuration:
-                _LOGGER.warning(f"Wrong config type: Received {type(format_configuration).__name__}, expected SequenceValueMetricsFormat.")
-            
-            sequence_to_value_metrics(y_true=y_true, 
-                                      y_pred=y_pred, 
-                                      save_dir=save_dir,
-                                      config=config)
-
-        elif self.kind == MLTaskKeys.SEQUENCE_SEQUENCE:
-            config = None
-            if format_configuration and isinstance(format_configuration, SequenceSequenceMetricsFormat):
-                config = format_configuration
-            elif format_configuration:
-                _LOGGER.warning(f"Wrong config type: Received {type(format_configuration).__name__}, expected SequenceSequenceMetricsFormat.")
-
-            sequence_to_sequence_metrics(y_true=y_true, 
-                                         y_pred=y_pred, 
-                                         save_dir=save_dir,
-                                         config=config)
-            
-    def explain_captum(self,
-                       save_dir: Union[str, Path],
-                       explain_dataset: Optional[Dataset] = None,
-                       n_samples: int = 100,
-                       feature_names: Optional[List[str]] = None,
-                       target_names: Optional[List[str]] = None,
-                       n_steps: int = 50):
-        """
-        Explains sequence model predictions using Captum's Integrated Gradients.
-
-        This method calculates global feature importance by aggregating attributions across 
-        the time dimension. 
-        - For **multivariate** sequences, it highlights which variables (channels) are most influential.
-        - For **univariate** sequences, it attributes importance to the single signal feature.
-
-        Args:
-            save_dir (str | Path): Directory to save the importance plots and CSV reports.
-            explain_dataset (Dataset | None): A specific dataset to sample from. If None, the 
-                                            trainer's validation dataset is used.
-            n_samples (int): The number of samples to use for the explanation (background + inputs).
-            feature_names (List[str] | None): Names of the features (signals). If None, attempts to extract them from the dataset attribute.
-            target_names (List[str] | None): Names of the model outputs (e.g., for Seq2Seq or Multivariate output). If None, attempts to extract them from the dataset attribute.
-            n_steps (int): Number of integral approximation steps.
-
-        Note:
-            For univariate data (Shape: N, Seq_Len), the 'feature' is the signal itself. 
-        """
-        if not _is_captum_available():
-            _LOGGER.error("Captum is not installed.")
-            return
-            
-        dataset_to_use = explain_dataset if explain_dataset is not None else self.validation_dataset
-        if dataset_to_use is None:
-            _LOGGER.error("No dataset available for explanation.")
-            return
-
-        # Helper to sample data (same as DragonTrainer)
-        def _get_samples(ds, n):
-            loader = DataLoader(ds, batch_size=n, shuffle=True, num_workers=0)
-            data_iter = iter(loader)
-            features, targets = next(data_iter)
-            return features, targets
-
-        input_data, _ = _get_samples(dataset_to_use, n_samples)
-        
-        if feature_names is None:
-             if hasattr(dataset_to_use, DatasetKeys.FEATURE_NAMES):
-                feature_names = dataset_to_use.feature_names # type: ignore
-             else:
-                # If retrieval fails, leave it as None. 
-                _LOGGER.warning("'feature_names' not provided and not found in dataset. Generic names will be used.")
-            
-        if target_names is None:
-            if hasattr(dataset_to_use, DatasetKeys.TARGET_NAMES):
-                target_names = dataset_to_use.target_names # type: ignore
-            else:
-                # If retrieval fails, leave it as None. 
-                _LOGGER.warning("'target_names' not provided and not found in dataset. Generic names will be used.")
-
-        # Sequence models usually output [N, 1] (Value) or [N, Seq, 1] (Seq2Seq)
-        # captum_feature_importance handles the aggregation.
-        
-        captum_feature_importance(
-            model=self.model,
-            input_data=input_data,
-            feature_names=feature_names,
-            save_dir=save_dir,
-            target_names=target_names,
-            n_steps=n_steps,
-            device=self.device
-        )
-    
-    def finalize_model_training(self, 
-                                save_dir: Union[str, Path], 
-                                model_checkpoint: Union[Path, Literal['best', 'current']],
-                                finalize_config: Union[FinalizeSequenceSequencePrediction, FinalizeSequenceValuePrediction]):
-        """
-        Saves a finalized, "inference-ready" model state to a .pth file.
-
-        This method saves the model's `state_dict` and the final epoch number.
-
-        Args:
-            save_dir (Union[str, Path]): The directory to save the finalized model.
-            model_checkpoint (Union[Path, Literal["best", "current"]]):
-                - Path: Loads the model state from a specific checkpoint file.
-                - "best": Loads the best model state saved by the `DragonModelCheckpoint` callback.
-                - "current": Uses the model's state as it is.
-            finalize_config (FinalizeSequencePrediction): A data class instance specific to the ML task containing task-specific metadata required for inference.
-        """
-        if self.kind == MLTaskKeys.SEQUENCE_SEQUENCE and not isinstance(finalize_config, FinalizeSequenceSequencePrediction):
-            _LOGGER.error(f"Received a wrong finalize configuration for task {self.kind}: {type(finalize_config).__name__}.")
-            raise TypeError()
-        elif self.kind == MLTaskKeys.SEQUENCE_VALUE and not isinstance(finalize_config, FinalizeSequenceValuePrediction):
-            _LOGGER.error(f"Received a wrong finalize configuration for task {self.kind}: {type(finalize_config).__name__}.")
-            raise TypeError()
-
-        # handle save path
-        dir_path = make_fullpath(save_dir, make=True, enforce="directory")
-        full_path = dir_path / finalize_config.filename
-        
-        # handle checkpoint
-        self._load_model_state_for_finalizing(model_checkpoint)
-        
-        # Create finalized data
-        finalized_data = {
-            PyTorchCheckpointKeys.EPOCH: self.epoch,
-            PyTorchCheckpointKeys.MODEL_STATE: self.model.state_dict(),
-            PyTorchCheckpointKeys.TASK: finalize_config.task
-        }
-        
-        if finalize_config.sequence_length is not None:
-            finalized_data[PyTorchCheckpointKeys.SEQUENCE_LENGTH] = finalize_config.sequence_length
-        if finalize_config.initial_sequence is not None:
-            finalized_data[PyTorchCheckpointKeys.INITIAL_SEQUENCE] = finalize_config.initial_sequence
-        
-        torch.save(finalized_data, full_path)
-        
-        _LOGGER.info(f"Finalized model file saved to '{full_path}'")
-
-
-def info():
-    _script_info(__all__)