configuration_utils

multimolecule.models.configuration_utils

HeadConfig

Bases: BaseHeadConfig

Configuration class for a prediction head.

Parameters:

Name	Type	Description	Default
num_labels		Number of labels to use in the last layer added to the model, typically for a classification task. Head should look for Config.num_labels if is None.	required
problem_type		Problem type for XxxForYyyPrediction models. Can be one of "binary", "regression", "multiclass" or "multilabel". Head should look for Config.problem_type if is None.	required
hidden_size		Dimensionality of the encoder layers and the pooler layer. Head should look for Config.hidden_size if is None.	required
dropout		The dropout ratio for the hidden states.	required
transform		The transform operation applied to hidden states.	required
transform_act		The activation function of transform applied to hidden states.	required
bias		Whether to apply bias to the final prediction layer.	required
act		The activation function of the final prediction output.	required
layer_norm_eps		The epsilon used by the layer normalization layers.	required
output_name		The name of the tensor required in model outputs. If is None, will use the default output name of the corresponding head.	required
type		The type of the head in the model. This is used by [MultiMoleculeModel][multimolecule.MultiMoleculeModel] to construct heads.	required

Source code in multimolecule/module/heads/config.py

class HeadConfig(BaseHeadConfig):
    r"""
    Configuration class for a prediction head.

    Args:
        num_labels:
            Number of labels to use in the last layer added to the model, typically for a classification task.

            Head should look for [`Config.num_labels`][multimolecule.PreTrainedConfig] if is `None`.
        problem_type:
            Problem type for `XxxForYyyPrediction` models. Can be one of `"binary"`, `"regression"`,
            `"multiclass"` or `"multilabel"`.

            Head should look for [`Config.problem_type`][multimolecule.PreTrainedConfig] if is `None`.
        hidden_size:
            Dimensionality of the encoder layers and the pooler layer.

            Head should look for [`Config.hidden_size`][multimolecule.PreTrainedConfig] if is `None`.
        dropout:
            The dropout ratio for the hidden states.
        transform:
            The transform operation applied to hidden states.
        transform_act:
            The activation function of transform applied to hidden states.
        bias:
            Whether to apply bias to the final prediction layer.
        act:
            The activation function of the final prediction output.
        layer_norm_eps:
            The epsilon used by the layer normalization layers.
        output_name:
            The name of the tensor required in model outputs.

            If is `None`, will use the default output name of the corresponding head.
        type:
            The type of the head in the model.

            This is used by [`MultiMoleculeModel`][multimolecule.MultiMoleculeModel] to construct heads.
    """

    num_labels: Optional[int] = None
    problem_type: Optional[str] = None
    hidden_size: Optional[int] = None
    dropout: float = 0.0
    transform: Optional[str] = None
    transform_act: Optional[str] = "gelu"
    bias: bool = True
    act: Optional[str] = None
    layer_norm_eps: float = 1e-12
    output_name: Optional[str] = None
    type: Optional[str] = None

MaskedLMHeadConfig

Bases: BaseHeadConfig

Configuration class for a Masked Language Modeling head.

Parameters:

Name	Type	Description	Default
hidden_size		Dimensionality of the encoder layers and the pooler layer. Head should look for Config.hidden_size if is None.	required
dropout		The dropout ratio for the hidden states.	required
transform		The transform operation applied to hidden states.	required
transform_act		The activation function of transform applied to hidden states.	required
bias		Whether to apply bias to the final prediction layer.	required
act		The activation function of the final prediction output.	required
layer_norm_eps		The epsilon used by the layer normalization layers.	required
output_name		The name of the tensor required in model outputs. If is None, will use the default output name of the corresponding head.	required

Source code in multimolecule/module/heads/config.py

class MaskedLMHeadConfig(BaseHeadConfig):
    r"""
    Configuration class for a Masked Language Modeling head.

    Args:
        hidden_size:
            Dimensionality of the encoder layers and the pooler layer.

            Head should look for [`Config.hidden_size`][multimolecule.PreTrainedConfig] if is `None`.
        dropout:
            The dropout ratio for the hidden states.
        transform:
            The transform operation applied to hidden states.
        transform_act:
            The activation function of transform applied to hidden states.
        bias:
            Whether to apply bias to the final prediction layer.
        act:
            The activation function of the final prediction output.
        layer_norm_eps:
            The epsilon used by the layer normalization layers.
        output_name:
            The name of the tensor required in model outputs.

            If is `None`, will use the default output name of the corresponding head.
    """

    hidden_size: Optional[int] = None
    dropout: float = 0.0
    transform: Optional[str] = "nonlinear"
    transform_act: Optional[str] = "gelu"
    bias: bool = True
    act: Optional[str] = None
    layer_norm_eps: float = 1e-12
    output_name: Optional[str] = None

PreTrainedConfig

Bases: PretrainedConfig

Base class for all model configuration classes.

Source code in multimolecule/models/configuration_utils.py

class PreTrainedConfig(PretrainedConfig):
    r"""
    Base class for all model configuration classes.
    """

    head: HeadConfig | None
    num_labels: int = 1

    hidden_size: int

    pad_token_id: int = 0
    bos_token_id: int = 1
    eos_token_id: int = 2
    unk_token_id: int = 3
    mask_token_id: int = 4
    null_token_id: int = 5

    def __init__(
        self,
        pad_token_id: int = 0,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        unk_token_id: int = 3,
        mask_token_id: int = 4,
        null_token_id: int = 5,
        num_labels: int = 1,
        **kwargs,
    ):
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            unk_token_id=unk_token_id,
            mask_token_id=mask_token_id,
            null_token_id=null_token_id,
            num_labels=num_labels,
            **kwargs,
        )

    def to_dict(self):
        output = super().to_dict()
        for k, v in output.items():
            if hasattr(v, "to_dict"):
                output[k] = v.to_dict()
            if is_dataclass(v):
                output[k] = asdict(v)
        return output