""" Ready-to-use configuration templates for the block-based Autoencoder. These helpers demonstrate how to assemble encoder_blocks and decoder_blocks for a variety of architectures using the new block system. Each class extends AutoencoderConfig and can be passed directly to AutoencoderModel. Example: from modeling_autoencoder import AutoencoderModel from template import ClassicAutoencoderConfig cfg = ClassicAutoencoderConfig(input_dim=784, latent_dim=64) model = AutoencoderModel(cfg) """ from __future__ import annotations from typing import List # Support both package-relative and flat import try: from .configuration_autoencoder import ( AutoencoderConfig, ) except Exception: # pragma: no cover from configuration_autoencoder import ( AutoencoderConfig, ) # ------------------------------- Helpers ------------------------------- # def _linear_stack(input_dim: int, dims: List[int], activation: str = "relu", normalization: str = "batch", dropout: float = 0.0): """Build a list of Linear block dict configs mapping input_dim -> dims sequentially.""" blocks = [] prev = input_dim for h in dims: blocks.append({ "type": "linear", "input_dim": prev, "output_dim": h, "activation": activation, "normalization": normalization, "dropout_rate": dropout, "use_residual": False, }) prev = h return blocks def _default_decoder(latent_dim: int, hidden: List[int], out_dim: int, activation: str = "relu", normalization: str = "batch", dropout: float = 0.0): """Linear decoder: latent_dim -> hidden -> out_dim (final layer identity).""" blocks = _linear_stack(latent_dim, hidden + [out_dim], activation, normalization, dropout) if blocks: blocks[-1]["activation"] = "identity" blocks[-1]["normalization"] = "none" blocks[-1]["dropout_rate"] = 0.0 return blocks # ---------------------------- Class-based templates ---------------------------- # class ClassicAutoencoderConfig(AutoencoderConfig): """Classic dense autoencoder using Linear blocks. Example: cfg = ClassicAutoencoderConfig(input_dim=784, latent_dim=64) """ def __init__(self, input_dim: int = 784, latent_dim: int = 64, hidden: List[int] = (512, 256, 128), activation: str = "relu", dropout: float = 0.1, use_batch_norm: bool = True, **kwargs): hidden = list(hidden) norm = "batch" if use_batch_norm else "none" enc = _linear_stack(input_dim, hidden, activation, norm, dropout) dec = _default_decoder(latent_dim, list(reversed(hidden)), input_dim, activation, norm, dropout) super().__init__( input_dim=input_dim, latent_dim=latent_dim, activation=activation, dropout_rate=dropout, use_batch_norm=use_batch_norm, autoencoder_type="classic", encoder_blocks=enc, decoder_blocks=dec, **kwargs, ) class VariationalAutoencoderConfig(AutoencoderConfig): """Variational autoencoder (MLP). Uses VariationalBlock in the model. Example: cfg = VariationalAutoencoderConfig(input_dim=784, latent_dim=32) """ def __init__(self, input_dim: int = 784, latent_dim: int = 32, hidden: List[int] = (512, 256, 128), activation: str = "relu", dropout: float = 0.1, use_batch_norm: bool = True, beta: float = 1.0, **kwargs): hidden = list(hidden) norm = "batch" if use_batch_norm else "none" enc = _linear_stack(input_dim, hidden, activation, norm, dropout) dec = _default_decoder(latent_dim, list(reversed(hidden)), input_dim, activation, norm, dropout) super().__init__( input_dim=input_dim, latent_dim=latent_dim, activation=activation, dropout_rate=dropout, use_batch_norm=use_batch_norm, autoencoder_type="variational", beta=beta, encoder_blocks=enc, decoder_blocks=dec, **kwargs, ) class TransformerAutoencoderConfig(AutoencoderConfig): """Transformer-style autoencoder with attention encoder and MLP decoder. Works with (batch, input_dim) or (batch, time, input_dim). Example: cfg = TransformerAutoencoderConfig(input_dim=256, latent_dim=128) """ def __init__(self, input_dim: int = 256, latent_dim: int = 128, num_layers: int = 2, num_heads: int = 4, ffn_mult: int = 4, activation: str = "relu", dropout: float = 0.1, use_batch_norm: bool = False, **kwargs): norm = "batch" if use_batch_norm else "none" enc = [] enc.append({"type": "linear", "input_dim": input_dim, "output_dim": input_dim, "activation": activation, "normalization": norm, "dropout_rate": dropout}) for _ in range(num_layers): enc.append({"type": "attention", "input_dim": input_dim, "num_heads": num_heads, "ffn_dim": ffn_mult * input_dim, "dropout_rate": dropout}) enc.append({"type": "linear", "input_dim": input_dim, "output_dim": input_dim, "activation": activation, "normalization": norm, "dropout_rate": dropout}) dec = _default_decoder(latent_dim, [input_dim], input_dim, activation, norm, dropout) super().__init__( input_dim=input_dim, latent_dim=latent_dim, activation=activation, dropout_rate=dropout, use_batch_norm=use_batch_norm, autoencoder_type="classic", encoder_blocks=enc, decoder_blocks=dec, **kwargs, ) class RecurrentAutoencoderConfig(AutoencoderConfig): """Recurrent encoder (LSTM/GRU/RNN) for sequence data. Expected input: (batch, time, input_dim). Decoder is MLP back to features per step. Example: cfg = RecurrentAutoencoderConfig(input_dim=128, latent_dim=64, rnn_type="lstm") """ def __init__(self, input_dim: int = 128, latent_dim: int = 64, rnn_type: str = "lstm", num_layers: int = 2, bidirectional: bool = False, activation: str = "relu", dropout: float = 0.1, use_batch_norm: bool = False, **kwargs): norm = "batch" if use_batch_norm else "none" enc = [{ "type": "recurrent", "input_dim": input_dim, "hidden_size": latent_dim, "num_layers": num_layers, "rnn_type": rnn_type, "bidirectional": bidirectional, "dropout_rate": dropout, "output_dim": latent_dim, }] dec = _default_decoder(latent_dim, [max(latent_dim, input_dim)], input_dim, activation, norm, dropout) super().__init__( input_dim=input_dim, latent_dim=latent_dim, activation=activation, dropout_rate=dropout, use_batch_norm=use_batch_norm, autoencoder_type="classic", encoder_blocks=enc, decoder_blocks=dec, **kwargs, ) class ConvolutionalAutoencoderConfig(AutoencoderConfig): """1D convolutional encoder for sequence data; decoder is per-step MLP. Expected input: (batch, time, input_dim). Example: cfg = ConvolutionalAutoencoderConfig(input_dim=64, conv_channels=(64, 64)) """ def __init__(self, input_dim: int = 64, latent_dim: int = 64, conv_channels: List[int] = (64, 64), kernel_size: int = 3, activation: str = "relu", dropout: float = 0.0, use_batch_norm: bool = True, **kwargs): norm = "batch" if use_batch_norm else "none" enc = [] prev = input_dim for ch in conv_channels: enc.append({"type": "conv1d", "input_dim": prev, "output_dim": ch, "kernel_size": kernel_size, "padding": "same", "activation": activation, "normalization": norm, "dropout_rate": dropout}) prev = ch enc.append({"type": "linear", "input_dim": prev, "output_dim": latent_dim, "activation": activation, "normalization": norm, "dropout_rate": dropout}) dec = _default_decoder(latent_dim, [prev], input_dim, activation, norm, dropout) super().__init__( input_dim=input_dim, latent_dim=latent_dim, activation=activation, dropout_rate=dropout, use_batch_norm=use_batch_norm, autoencoder_type="classic", encoder_blocks=enc, decoder_blocks=dec, **kwargs, ) class ConvAttentionAutoencoderConfig(AutoencoderConfig): """Mixed Conv + Attention encoder for sequence data. Example: cfg = ConvAttentionAutoencoderConfig(input_dim=64, latent_dim=64) """ def __init__(self, input_dim: int = 64, latent_dim: int = 64, conv_channels: List[int] = (64,), num_heads: int = 4, activation: str = "relu", dropout: float = 0.1, use_batch_norm: bool = True, **kwargs): norm = "batch" if use_batch_norm else "none" enc = [] prev = input_dim for ch in conv_channels: enc.append({"type": "conv1d", "input_dim": prev, "output_dim": ch, "kernel_size": 3, "padding": "same", "activation": activation, "normalization": norm, "dropout_rate": dropout}) prev = ch enc.append({"type": "attention", "input_dim": prev, "num_heads": num_heads, "ffn_dim": 4 * prev, "dropout_rate": dropout}) enc.append({"type": "linear", "input_dim": prev, "output_dim": latent_dim, "activation": activation, "normalization": norm, "dropout_rate": dropout}) dec = _default_decoder(latent_dim, [prev], input_dim, activation, norm, dropout) super().__init__( input_dim=input_dim, latent_dim=latent_dim, activation=activation, dropout_rate=dropout, use_batch_norm=use_batch_norm, autoencoder_type="classic", encoder_blocks=enc, decoder_blocks=dec, **kwargs, ) class LinearRecurrentAutoencoderConfig(AutoencoderConfig): """Linear down-projection then Recurrent encoder. Example: cfg = LinearRecurrentAutoencoderConfig(input_dim=256, latent_dim=64, rnn_type="gru") """ def __init__(self, input_dim: int = 256, latent_dim: int = 64, rnn_type: str = "gru", activation: str = "relu", dropout: float = 0.1, use_batch_norm: bool = False, **kwargs): norm = "batch" if use_batch_norm else "none" enc = [ {"type": "linear", "input_dim": input_dim, "output_dim": latent_dim, "activation": activation, "normalization": norm, "dropout_rate": dropout}, {"type": "recurrent", "input_dim": latent_dim, "hidden_size": latent_dim, "num_layers": 1, "rnn_type": rnn_type, "bidirectional": False, "dropout_rate": dropout, "output_dim": latent_dim}, ] dec = _default_decoder(latent_dim, [], input_dim, activation, norm, dropout) super().__init__( input_dim=input_dim, latent_dim=latent_dim, activation=activation, dropout_rate=dropout, use_batch_norm=use_batch_norm, autoencoder_type="classic", encoder_blocks=enc, decoder_blocks=dec, **kwargs, ) class PreprocessedAutoencoderConfig(AutoencoderConfig): """Classic MLP AE with learnable preprocessing/inverse. Example: cfg = PreprocessedAutoencoderConfig(input_dim=64, preprocessing_type="neural_scaler") """ def __init__(self, input_dim: int = 64, latent_dim: int = 32, preprocessing_type: str = "neural_scaler", hidden: List[int] = (128, 64), activation: str = "relu", dropout: float = 0.0, use_batch_norm: bool = True, **kwargs): norm = "batch" if use_batch_norm else "none" enc = _linear_stack(input_dim, list(hidden), activation, norm, dropout) dec = _default_decoder(latent_dim, list(reversed(list(hidden))), input_dim, activation, norm, dropout) super().__init__( input_dim=input_dim, latent_dim=latent_dim, activation=activation, dropout_rate=dropout, use_batch_norm=use_batch_norm, autoencoder_type="classic", use_learnable_preprocessing=True, preprocessing_type=preprocessing_type, encoder_blocks=enc, decoder_blocks=dec, **kwargs, ) class BetaVariationalAutoencoderConfig(AutoencoderConfig): """Beta-VAE (MLP). Like VAE but with beta > 1 controlling KL weight. Example: cfg = BetaVariationalAutoencoderConfig(input_dim=784, latent_dim=32, beta=4.0) """ def __init__(self, input_dim: int = 784, latent_dim: int = 32, hidden: List[int] = (512, 256, 128), activation: str = "relu", dropout: float = 0.1, use_batch_norm: bool = True, beta: float = 4.0, **kwargs): hidden = list(hidden) norm = "batch" if use_batch_norm else "none" enc = _linear_stack(input_dim, hidden, activation, norm, dropout) dec = _default_decoder(latent_dim, list(reversed(hidden)), input_dim, activation, norm, dropout) super().__init__( input_dim=input_dim, latent_dim=latent_dim, activation=activation, dropout_rate=dropout, use_batch_norm=use_batch_norm, autoencoder_type="beta_vae", beta=beta, encoder_blocks=enc, decoder_blocks=dec, **kwargs, ) class DenoisingAutoencoderConfig(AutoencoderConfig): """Denoising AE: adds noise during training (handled by training loop/model if supported). Example: cfg = DenoisingAutoencoderConfig(input_dim=128, latent_dim=32, noise_factor=0.2) """ def __init__(self, input_dim: int = 128, latent_dim: int = 32, hidden: List[int] = (128, 64), activation: str = "relu", dropout: float = 0.0, use_batch_norm: bool = True, noise_factor: float = 0.2, **kwargs): hidden = list(hidden) norm = "batch" if use_batch_norm else "none" enc = _linear_stack(input_dim, hidden, activation, norm, dropout) dec = _default_decoder(latent_dim, list(reversed(hidden)), input_dim, activation, norm, dropout) super().__init__( input_dim=input_dim, latent_dim=latent_dim, activation=activation, dropout_rate=dropout, use_batch_norm=use_batch_norm, autoencoder_type="denoising", noise_factor=noise_factor, encoder_blocks=enc, decoder_blocks=dec, **kwargs, ) class SparseAutoencoderConfig(AutoencoderConfig): """Sparse AE (typical L1 activation penalty applied in training loop). Example: cfg = SparseAutoencoderConfig(input_dim=256, latent_dim=64) """ def __init__(self, input_dim: int = 256, latent_dim: int = 64, hidden: List[int] = (128, 64), activation: str = "relu", dropout: float = 0.0, use_batch_norm: bool = True, **kwargs): hidden = list(hidden) norm = "batch" if use_batch_norm else "none" enc = _linear_stack(input_dim, hidden, activation, norm, dropout) dec = _default_decoder(latent_dim, list(reversed(hidden)), input_dim, activation, norm, dropout) super().__init__( input_dim=input_dim, latent_dim=latent_dim, activation=activation, dropout_rate=dropout, use_batch_norm=use_batch_norm, autoencoder_type="sparse", encoder_blocks=enc, decoder_blocks=dec, **kwargs, ) class ContractiveAutoencoderConfig(AutoencoderConfig): """Contractive AE (requires Jacobian penalty in training loop). Example: cfg = ContractiveAutoencoderConfig(input_dim=64, latent_dim=16) """ def __init__(self, input_dim: int = 64, latent_dim: int = 16, hidden: List[int] = (64, 32), activation: str = "relu", dropout: float = 0.0, use_batch_norm: bool = True, **kwargs): hidden = list(hidden) norm = "batch" if use_batch_norm else "none" enc = _linear_stack(input_dim, hidden, activation, norm, dropout) dec = _default_decoder(latent_dim, list(reversed(hidden)), input_dim, activation, norm, dropout) super().__init__( input_dim=input_dim, latent_dim=latent_dim, activation=activation, dropout_rate=dropout, use_batch_norm=use_batch_norm, autoencoder_type="contractive", encoder_blocks=enc, decoder_blocks=dec, **kwargs, ) __all__ = [ "ClassicAutoencoderConfig", "VariationalAutoencoderConfig", "TransformerAutoencoderConfig", "RecurrentAutoencoderConfig", "ConvolutionalAutoencoderConfig", "ConvAttentionAutoencoderConfig", "LinearRecurrentAutoencoderConfig", "PreprocessedAutoencoderConfig", "BetaVariationalAutoencoderConfig", "DenoisingAutoencoderConfig", "SparseAutoencoderConfig", "ContractiveAutoencoderConfig", ]