Function bert_model_config

axlearn/common/bert.py:463–511

Builds configs for BERT model. Defaults are from the BERT-BASE model. Args: vocab_size: Vocab size. hidden_dim: Hidden dim. dropout_rate: Dropout rate. dtype: Model dtype. embedding_cfg: Optional embedding config. Defaults to a BERT-base Embedding.

(
    *,
    vocab_size: int,
    hidden_dim: int = 768,
    dropout_rate: float = 0.0,
    dtype: jnp.dtype = jnp.float32,
    embedding_cfg: Optional[Embedding.Config] = None,
    stack_cfg: Optional[BaseStackedTransformerLayer.Config] = None,
    head_cfg: Optional[BaseClassificationHead.Config] = None,
    encoder_cfg: Optional[Encoder.Config] = None,
    base_cfg: Optional[BertModel.Config] = None,
)

def bert_model_config(
    *,
    vocab_size: int,
    hidden_dim: int = 768,
    dropout_rate: float = 0.0,
    dtype: jnp.dtype = jnp.float32,
    embedding_cfg: Optional[Embedding.Config] = None,
    stack_cfg: Optional[BaseStackedTransformerLayer.Config] = None,
    head_cfg: Optional[BaseClassificationHead.Config] = None,
    encoder_cfg: Optional[Encoder.Config] = None,
    base_cfg: Optional[BertModel.Config] = None,
) -> BertModel.Config:
    """Builds configs for BERT model.

    Defaults are from the BERT-BASE model.

    Args:
        vocab_size: Vocab size.
        hidden_dim: Hidden dim.
        dropout_rate: Dropout rate.
        dtype: Model dtype.
        embedding_cfg: Optional embedding config. Defaults to a BERT-base Embedding.
        stack_cfg: Optional transformer stack config. Defaults to a StackedTransformerLayer.
        head_cfg: Optional head config. Defaults to a BertLMHead.Config.
        encoder_cfg: Optional encoder config. Defaults to a Encoder.Config.
        base_cfg: Optional base config. Will be cloned.

    Returns:
        The stack config.
    """
    base_cfg = base_cfg.clone() if base_cfg else BertModel.default_config()

    encoder_cfg = encoder_cfg or Encoder.default_config().set(
        dim=hidden_dim,
        vocab_size=vocab_size,
        dropout_rate=dropout_rate,
        emb=embedding_cfg or bert_embedding_config(),
        transformer=stack_cfg or bert_transformer_config(),
        pad_token_id=0,
    )
    head_cfg = head_cfg or bert_lm_head_config(
        base_cfg=base_cfg.head,  # pylint: disable=no-member
        vocab_size=vocab_size,
    )
    base_cfg.set(
        dtype=dtype, vocab_size=vocab_size, dim=hidden_dim, encoder=encoder_cfg, head=head_cfg
    )
    set_layer_norm_eps_recursively(base_cfg, bert_layer_norm_epsilon(dtype=dtype))
    return base_cfg


def bert_layer_norm_epsilon(dtype=None):