Class FlexBertPaddedEncoder

src/bert_layers/layers.py:665–684

A stack of BERT layers providing the backbone of FlexBERT. This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.modeling_bert.BertAlibiEncoder`, but with substantial modifications to implement unpadding and ALiBi. Compared to the analogous Hugging Face

class FlexBertPaddedEncoder(FlexBertEncoderBase):
    """A stack of BERT layers providing the backbone of FlexBERT.

    This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.modeling_bert.BertAlibiEncoder`,
    but with substantial modifications to implement unpadding and ALiBi.

    Compared to the analogous Hugging Face BERT module, this module handles unpadding to reduce unnecessary computation
    at padded tokens, and pre-computes attention biases to implement ALiBi.
    """

    def __init__(self, config: FlexBertConfig):
        super().__init__()
        self.layers = nn.ModuleList([get_bert_layer(config, layer_id=i) for i in range(config.num_hidden_layers)])
        self.num_attention_heads = config.num_attention_heads

    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> torch.Tensor:
        for layer_module in self.layers:
            hidden_states = layer_module(hidden_states, attn_mask=attention_mask)

        return hidden_states


ENC2CLS = {