Function bert_embedding_config

axlearn/common/bert.py:351–383

Builds configs for BERT Embedding layer. Defaults are from the BERT-BASE model. Args: max_position_embeddings: Number of positional embeddings. If None, position embedding is not used. layer_norm_epsilon: Optional epsilon for layer norm. type_vocab_size:

(
    *,
    max_position_embeddings: Optional[int] = 512,
    layer_norm_epsilon: Optional[float] = None,
    type_vocab_size: Optional[int] = None,
)

def bert_embedding_config(
    *,
    max_position_embeddings: Optional[int] = 512,
    layer_norm_epsilon: Optional[float] = None,
    type_vocab_size: Optional[int] = None,
) -> TransformerTextEmbeddings.Config:
    """Builds configs for BERT Embedding layer.

    Defaults are from the BERT-BASE model.

    Args:
        max_position_embeddings: Number of positional embeddings.
            If None, position embedding is not used.
        layer_norm_epsilon: Optional epsilon for layer norm.
        type_vocab_size: Optional number of token type embeddings.

    Returns:
        The embedding configs.
    """
    cfg = TransformerTextEmbeddings.default_config().set(
        norm=LayerNorm.default_config().set(eps=layer_norm_epsilon),
    )
    if max_position_embeddings is not None:
        cfg.set(
            pos_emb=LearnedPositionalEmbedding.default_config().set(
                shape=(max_position_embeddings,)
            )
        )
    if type_vocab_size is not None:
        cfg.set(
            type_emb=Embedding.default_config().set(num_embeddings=type_vocab_size),
        )
    return cfg


# Linter thinks `ignored_target_id` is an unused param and complains that we document it, so we