configuration_autoencoder.py

"""
Autoencoder configuration for Hugging Face Transformers.
"""

from dataclasses import dataclass
from typing import Union

# Import PretrainedConfig in a way that avoids circular imports in some environments (e.g., Databricks)
try:
    from transformers.configuration_utils import PretrainedConfig
except Exception:  # fallback
    from transformers import PretrainedConfig
from typing import List, Optional

# Support both package-relative and flat import in HF remote code context
try:
    from . import __version__ as _pkg_version  # type: ignore
except Exception:  # pragma: no cover
    _pkg_version = None

@dataclass
class BlockConfig:
    type: str


@dataclass
class LinearBlockConfig(BlockConfig):
    input_dim: int
    output_dim: int
    activation: str = "relu"
    normalization: Optional[str] = "batch"  # batch|layer|group|instance|none
    dropout_rate: float = 0.0
    use_residual: bool = False
    residual_scale: float = 1.0

    def __init__(self, input_dim: int, output_dim: int, activation: str = "relu", normalization: Optional[str] = "batch", dropout_rate: float = 0.0, use_residual: bool = False, residual_scale: float = 1.0):
        super().__init__(type="linear")
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.activation = activation
        self.normalization = normalization
        self.dropout_rate = dropout_rate
        self.use_residual = use_residual
        self.residual_scale = residual_scale


@dataclass
class AttentionBlockConfig(BlockConfig):
    input_dim: int
    num_heads: int = 8
    ffn_dim: Optional[int] = None
    dropout_rate: float = 0.0

    def __init__(self, input_dim: int, num_heads: int = 8, ffn_dim: Optional[int] = None, dropout_rate: float = 0.0):
        super().__init__(type="attention")
        self.input_dim = input_dim
        self.num_heads = num_heads
        self.ffn_dim = ffn_dim
        self.dropout_rate = dropout_rate

@dataclass
class RecurrentBlockConfig(BlockConfig):
    input_dim: int
    hidden_size: int
    num_layers: int = 1
    rnn_type: str = "lstm"  # lstm|gru|rnn
    bidirectional: bool = False
    dropout_rate: float = 0.0
    output_dim: Optional[int] = None  # if None, use hidden_size * directions

    def __init__(self, input_dim: int, hidden_size: int, num_layers: int = 1, rnn_type: str = "lstm", bidirectional: bool = False, dropout_rate: float = 0.0, output_dim: Optional[int] = None):
        super().__init__(type="recurrent")
        self.input_dim = input_dim
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.rnn_type = rnn_type
        self.bidirectional = bidirectional
        self.dropout_rate = dropout_rate
        self.output_dim = output_dim


@dataclass
class ConvolutionalBlockConfig(BlockConfig):
    input_dim: int  # channels in (features)
    output_dim: int  # channels out
    kernel_size: int = 3
    padding: Union[int, str] = "same"  # "same" or int
    activation: str = "relu"
    normalization: Optional[str] = "batch"
    dropout_rate: float = 0.0

    def __init__(self, input_dim: int, output_dim: int, kernel_size: int = 3, padding: Union[int, str] = "same", activation: str = "relu", normalization: Optional[str] = "batch", dropout_rate: float = 0.0):
        super().__init__(type="conv1d")
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.kernel_size = kernel_size
        self.padding = padding
        self.activation = activation
        self.normalization = normalization
        self.dropout_rate = dropout_rate

@dataclass
class VariationalBlockConfig(BlockConfig):
    input_dim: int
    latent_dim: int

    def __init__(self, input_dim: int, latent_dim: int):
        super().__init__(type="variational")
        self.input_dim = input_dim
        self.latent_dim = latent_dim


class AutoencoderConfig(PretrainedConfig):
    """
    Configuration class for Autoencoder models.

    This configuration class stores the configuration of an autoencoder model. It is used to instantiate
    an autoencoder model according to the specified arguments, defining the model architecture.

    Args:
        input_dim (int, optional): Dimensionality of the input data. Defaults to 784.
        hidden_dims (List[int], optional): Legacy: List of hidden layer dims for simple MLP encoder.
        encoder_blocks (List[dict], optional): New: List of block configs for encoder.
        decoder_blocks (List[dict], optional): New: List of block configs for decoder.
        latent_dim (int, optional): Dimensionality of the latent space. Defaults to 64.
        activation (str, optional): Default activation for Linear blocks. See supported list below.
        dropout_rate (float, optional): Default dropout for Linear blocks. Defaults to 0.1.
        use_batch_norm (bool, optional): Default normalization for Linear blocks (batch vs none). Defaults to True.
        tie_weights (bool, optional): Whether to tie encoder and decoder weights. Defaults to False.
        reconstruction_loss (str, optional): Type of reconstruction loss. Options: "mse", "bce", "l1",
            "huber", "smooth_l1", "kl_div", "cosine", "focal", "dice", "tversky", "ssim", "perceptual".
            Defaults to "mse".
        autoencoder_type (str, optional): Type of autoencoder architecture. Options: "classic",
            "variational", "beta_vae", "denoising", "sparse", "contractive", "recurrent". Defaults to "classic".
        beta (float, optional): Beta parameter for beta-VAE. Defaults to 1.0.
        temperature (float, optional): Temperature parameter for Gumbel softmax or other operations. Defaults to 1.0.
        noise_factor (float, optional): Noise factor for denoising autoencoders. Defaults to 0.1.
        rnn_type (str, optional): Type of RNN cell for recurrent autoencoders. Options: "lstm", "gru", "rnn".
            Defaults to "lstm".
        num_layers (int, optional): Number of RNN layers for recurrent autoencoders. Defaults to 2.
        bidirectional (bool, optional): Whether to use bidirectional RNN for encoding. Defaults to True.
        sequence_length (int, optional): Fixed sequence length. If None, supports variable length sequences.
            Defaults to None.
        teacher_forcing_ratio (float, optional): Ratio of teacher forcing during training for recurrent decoders.
            Defaults to 0.5.
        use_learnable_preprocessing (bool, optional): Whether to use learnable preprocessing. Defaults to False.
        preprocessing_type (str, optional): Type of learnable preprocessing. Options: "none", "neural_scaler",
            "normalizing_flow", "minmax_scaler", "robust_scaler", "yeo_johnson". Defaults to "none".
        preprocessing_hidden_dim (int, optional): Hidden dimension for preprocessing networks. Defaults to 64.
        preprocessing_num_layers (int, optional): Number of layers in preprocessing networks. Defaults to 2.
        learn_inverse_preprocessing (bool, optional): Whether to learn inverse preprocessing for reconstruction.
            Defaults to True.
        flow_coupling_layers (int, optional): Number of coupling layers for normalizing flows. Defaults to 4.
        **kwargs: Additional keyword arguments passed to the parent class.
    """

    model_type = "autoencoder"

    def __init__(
        self,
        input_dim: int = 784,
        hidden_dims: List[int] = None,
        encoder_blocks: Optional[List[dict]] = None,
        decoder_blocks: Optional[List[dict]] = None,
        latent_dim: int = 64,
        activation: str = "relu",
        dropout_rate: float = 0.1,
        use_batch_norm: bool = True,
        tie_weights: bool = False,
        reconstruction_loss: str = "mse",
        autoencoder_type: str = "classic",
        beta: float = 1.0,
        temperature: float = 1.0,
        noise_factor: float = 0.1,
        # Recurrent autoencoder parameters
        rnn_type: str = "lstm",
        num_layers: int = 2,
        bidirectional: bool = True,
        sequence_length: Optional[int] = None,
        teacher_forcing_ratio: float = 0.5,
        # Deep learning preprocessing parameters
        use_learnable_preprocessing: bool = False,
        preprocessing_type: str = "none",
        preprocessing_hidden_dim: int = 64,
        preprocessing_num_layers: int = 2,
        learn_inverse_preprocessing: bool = True,
        flow_coupling_layers: int = 4,
        **kwargs,
    ):
        # Validate parameters
        if hidden_dims is None:
            hidden_dims = [512, 256, 128]

        # Extended activation functions
        valid_activations = [
            "relu", "tanh", "sigmoid", "leaky_relu", "gelu", "swish", "silu",
            "elu", "prelu", "relu6", "hardtanh", "hardsigmoid", "hardswish",
            "mish", "softplus", "softsign", "tanhshrink", "threshold"
        ]
        if activation not in valid_activations:
            raise ValueError(
                f"`activation` must be one of {valid_activations}, got {activation}."
            )

        # Extended loss functions
        valid_losses = [
            "mse", "bce", "l1", "huber", "smooth_l1", "kl_div", "cosine",
            "focal", "dice", "tversky", "ssim", "perceptual"
        ]
        if reconstruction_loss not in valid_losses:
            raise ValueError(
                f"`reconstruction_loss` must be one of {valid_losses}, got {reconstruction_loss}."
            )

        # Autoencoder types
        valid_types = ["classic", "variational", "beta_vae", "denoising", "sparse", "contractive", "recurrent"]
        if autoencoder_type not in valid_types:
            raise ValueError(
                f"`autoencoder_type` must be one of {valid_types}, got {autoencoder_type}."
            )

        # RNN types for recurrent autoencoders
        valid_rnn_types = ["lstm", "gru", "rnn"]
        if rnn_type not in valid_rnn_types:
            raise ValueError(
                f"`rnn_type` must be one of {valid_rnn_types}, got {rnn_type}."
            )

        if not (0.0 <= dropout_rate <= 1.0):
            raise ValueError(f"`dropout_rate` must be between 0.0 and 1.0, got {dropout_rate}.")

        if input_dim <= 0:
            raise ValueError(f"`input_dim` must be positive, got {input_dim}.")

        if latent_dim <= 0:
            raise ValueError(f"`latent_dim` must be positive, got {latent_dim}.")

        if not all(dim > 0 for dim in hidden_dims):
            raise ValueError("All dimensions in `hidden_dims` must be positive.")

        if beta <= 0:
            raise ValueError(f"`beta` must be positive, got {beta}.")

        if num_layers <= 0:
            raise ValueError(f"`num_layers` must be positive, got {num_layers}.")

        if not (0.0 <= teacher_forcing_ratio <= 1.0):
            raise ValueError(f"`teacher_forcing_ratio` must be between 0.0 and 1.0, got {teacher_forcing_ratio}.")

        if sequence_length is not None and sequence_length <= 0:
            raise ValueError(f"`sequence_length` must be positive when specified, got {sequence_length}.")

        # Preprocessing validation
        valid_preprocessing = [
            "none",
            "neural_scaler",
            "normalizing_flow",
            "minmax_scaler",
            "robust_scaler",
            "yeo_johnson",
        ]
        if preprocessing_type not in valid_preprocessing:
            raise ValueError(
                f"`preprocessing_type` must be one of {valid_preprocessing}, got {preprocessing_type}."
            )

        if preprocessing_hidden_dim <= 0:
            raise ValueError(f"`preprocessing_hidden_dim` must be positive, got {preprocessing_hidden_dim}.")

        if preprocessing_num_layers <= 0:
            raise ValueError(f"`preprocessing_num_layers` must be positive, got {preprocessing_num_layers}.")

        if flow_coupling_layers <= 0:
            raise ValueError(f"`flow_coupling_layers` must be positive, got {flow_coupling_layers}.")

        # Set configuration attributes
        self.input_dim = input_dim
        self.hidden_dims = hidden_dims
        self.encoder_blocks = encoder_blocks
        self.decoder_blocks = decoder_blocks
        self.latent_dim = latent_dim
        self.activation = activation
        self.dropout_rate = dropout_rate
        self.use_batch_norm = use_batch_norm
        self.tie_weights = tie_weights
        self.reconstruction_loss = reconstruction_loss
        self.autoencoder_type = autoencoder_type
        self.beta = beta
        self.temperature = temperature
        self.noise_factor = noise_factor
        self.rnn_type = rnn_type
        self.num_layers = num_layers
        self.bidirectional = bidirectional
        self.sequence_length = sequence_length
        self.teacher_forcing_ratio = teacher_forcing_ratio
        self.use_learnable_preprocessing = use_learnable_preprocessing
        self.preprocessing_type = preprocessing_type
        self.preprocessing_hidden_dim = preprocessing_hidden_dim
        self.preprocessing_num_layers = preprocessing_num_layers
        self.learn_inverse_preprocessing = learn_inverse_preprocessing
        self.flow_coupling_layers = flow_coupling_layers

        # Call parent constructor
        super().__init__(**kwargs)

    @property
    def decoder_dims(self) -> List[int]:
        """Get decoder dimensions (reverse of encoder hidden dims)."""
        return list(reversed(self.hidden_dims))

    @property
    def has_block_lists(self) -> bool:
        """Whether explicit encoder/decoder block configs are provided."""
        return (self.encoder_blocks is not None) or (self.decoder_blocks is not None)

    @property
    def is_variational(self) -> bool:
        """Check if this is a variational autoencoder."""
        return self.autoencoder_type in ["variational", "beta_vae"]

    @property
    def is_denoising(self) -> bool:
        """Check if this is a denoising autoencoder."""
        return self.autoencoder_type == "denoising"

    @property
    def is_sparse(self) -> bool:
        """Check if this is a sparse autoencoder."""
        return self.autoencoder_type == "sparse"

    @property
    def is_contractive(self) -> bool:
        """Check if this is a contractive autoencoder."""
        return self.autoencoder_type == "contractive"

    @property
    def is_recurrent(self) -> bool:
        """Check if this is a recurrent autoencoder."""
        return self.autoencoder_type == "recurrent"

    @property
    def rnn_hidden_size(self) -> int:
        """Get the RNN hidden size (same as latent_dim for recurrent AE)."""
        return self.latent_dim

    @property
    def rnn_output_size(self) -> int:
        """Get the RNN output size considering bidirectionality."""
        return self.latent_dim * (2 if self.bidirectional else 1)

    @property
    def has_preprocessing(self) -> bool:
        """Check if learnable preprocessing is enabled."""
        return self.use_learnable_preprocessing and self.preprocessing_type != "none"

    @property
    def is_neural_scaler(self) -> bool:
        """Check if using neural scaler preprocessing."""
        return self.preprocessing_type == "neural_scaler"

    @property
    def is_normalizing_flow(self) -> bool:
        """Check if using normalizing flow preprocessing."""
        return self.preprocessing_type == "normalizing_flow"

    @property
    def is_minmax_scaler(self) -> bool:
        """Check if using learnable MinMax scaler preprocessing."""
        return self.preprocessing_type == "minmax_scaler"

    @property
    def is_robust_scaler(self) -> bool:
        """Check if using learnable Robust scaler preprocessing."""
        return self.preprocessing_type == "robust_scaler"

    @property
    def is_yeo_johnson(self) -> bool:
        """Check if using learnable Yeo-Johnson power transform preprocessing."""
        return self.preprocessing_type == "yeo_johnson"

    def to_dict(self):
        """
        Serializes this instance to a Python dictionary.
        """
        output = super().to_dict()
        return output