OCDocker.OCScore.DNN.future.models module

Models for the future DNN pipeline (shared encoder + multi-head).

class OCDocker.OCScore.DNN.future.models.MLP(*args, **kwargs)[source]

Bases: Module

Simple MLP with optional batch norm and dropout.

Parameters:

  • input_size (int) – Input feature dimension.

  • layer_sizes (list[int]) – Hidden layer sizes (final size is last).

  • activations (list[tuple[str, dict]] | str) – Activation configuration for each layer or a single activation name.

  • dropout (float, optional) – Dropout probability applied after activation, by default 0.0.

  • batch_norm (bool, optional) – Use BatchNorm1d after Linear, by default True.

__init__(input_size, layer_sizes, activations='GELU', dropout=0.0, batch_norm=True)[source]

Initialize MLP.

Parameters:

  • input_size (int) – Input feature dimension.

  • layer_sizes (list[int]) – Layer sizes.

  • activations (str | list[tuple[str, dict]], optional) – Activation configuration, by default “GELU”.

  • dropout (float, optional) – Dropout probability, by default 0.0.

  • batch_norm (bool, optional) – Use BatchNorm1d, by default True.

Return type:

None

forward(x)[source]

Forward pass.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor.

Return type:

torch.Tensor

class OCDocker.OCScore.DNN.future.models.MultiTaskModel(*args, **kwargs)[source]

Bases: Module

Shared encoder with heads for energy, activity, embedding and reconstruction.

Parameters:

  • input_size (int) – Input feature dimension.

  • encoder_params (dict | None) – Encoder params (old-style).

  • shared_sizes (list[int]) – Sizes for shared encoder when encoder_params is None.

  • shared_activation (str, optional) – Activation for shared encoder, by default GELU.

  • decoder_sizes (list[int] | None, optional) – Decoder sizes for reconstruction. If None, decoder is disabled.

  • head_sizes (list[int]) – Hidden sizes for heads.

  • embedding_dim (int | None) – Output dimension for embedding head.

  • dropout (float, optional) – Dropout probability.

  • batch_norm (bool, optional) – Use BatchNorm1d.

  • mask (torch.Tensor | None) – Feature mask.

__init__(input_size, encoder_params, shared_sizes, shared_activation='GELU', decoder_sizes=None, head_sizes=[128, 64], embedding_dim=64, dropout=0.0, batch_norm=True, mask=None)[source]

Initialize multi-task model.

Parameters:

  • input_size (int) – Input feature dimension.

  • encoder_params (dict | None) – Encoder parameters (legacy format).

  • shared_sizes (list[int]) – Shared encoder layer sizes.

  • shared_activation (str, optional) – Activation name, by default “GELU”.

  • decoder_sizes (list[int] | None, optional) – Decoder sizes for reconstruction, by default None.

  • head_sizes (list[int], optional) – Head hidden sizes, by default [128, 64].

  • embedding_dim (int | None, optional) – Embedding dimension, by default 64.

  • dropout (float, optional) – Dropout probability, by default 0.0.

  • batch_norm (bool, optional) – Use BatchNorm1d, by default True.

  • mask (torch.Tensor | None, optional) – Feature mask, by default None.

Return type:

None

forward(x, return_reconstruction=False)[source]

Forward pass.

Parameters:

  • x (torch.Tensor) – Input tensor.

  • return_reconstruction (bool, optional) – Whether to return reconstruction, by default False.

Returns:

Dictionary with latent, energy, activity, embedding and reconstruction tensors.

Return type:

dict[str, torch.Tensor]

OCDocker.OCScore.DNN.future.models.parse_encoder_params(encoder_params)[source]

Parse old-style encoder params into layer sizes and activation configs.

Parameters:

encoder_params (Dict[str, Any]) – Encoder parameters dictionary.

Returns:

Layer sizes and activation configs.

Return type:

tuple[list[int], list[tuple[str, dict]]]