modeling_bert.py

import math
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint
from transformers.activations import ACT2FN
from transformers import PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput

from .configuration_bert_updated import BertUpdatedConfig


class BertSelfAttention(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = config.hidden_size // config.num_attention_heads
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)
        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        return x.view(B, T, self.num_attention_heads, self.attention_head_size).permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        head_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        q = self._split_heads(self.query(hidden_states))
        k = self._split_heads(self.key(hidden_states))
        v = self._split_heads(self.value(hidden_states))

        scale = math.sqrt(self.attention_head_size)
        scores = torch.matmul(q, k.transpose(-1, -2)) / scale
        if key_padding_mask is not None:
            scores = scores.masked_fill(key_padding_mask[:, None, None, :], float("-inf"))
        probs = torch.nan_to_num(
            F.softmax(scores.float(), dim=-1),
            nan=0.0,
        )
        context_probs = self.dropout(probs)
        if head_mask is not None:
            context_probs = context_probs * head_mask
        context = torch.matmul(context_probs.to(v.dtype), v)

        B, _, T, _ = context.shape
        context = context.permute(0, 2, 1, 3).contiguous().view(B, T, self.all_head_size)

        if output_attentions:
            return context, probs
        return context, None


class BertSdpaSelfAttention(BertSelfAttention):

    def forward(
        self,
        hidden_states: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        head_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        if output_attentions or head_mask is not None:
            return super().forward(
                hidden_states,
                key_padding_mask,
                output_attentions=output_attentions,
                head_mask=head_mask,
            )

        B, T, _ = hidden_states.shape
        q = self._split_heads(self.query(hidden_states))
        k = self._split_heads(self.key(hidden_states))
        v = self._split_heads(self.value(hidden_states))

        attn_mask = None
        if key_padding_mask is not None:
            attn_mask = torch.zeros(B, 1, 1, T, dtype=q.dtype, device=q.device)
            attn_mask = attn_mask.masked_fill(key_padding_mask[:, None, None, :], float("-inf"))

        context = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=attn_mask,
            dropout_p=self.dropout.p if self.training else 0.0,
        )
        context = context.permute(0, 2, 1, 3).contiguous().view(B, T, self.all_head_size)
        return context, None


class BertFlashSelfAttention(BertSelfAttention):

    def forward(
        self,
        hidden_states: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        head_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        if (
            output_attentions
            or head_mask is not None
            or (
                key_padding_mask is not None
                and (~key_padding_mask).sum(dim=-1).eq(0).any()
            )
        ):
            return super().forward(
                hidden_states,
                key_padding_mask,
                output_attentions=output_attentions,
                head_mask=head_mask,
            )

        try:
            from flash_attn import flash_attn_func, flash_attn_varlen_func
            from flash_attn.bert_padding import pad_input, unpad_input
        except ImportError as e:
            raise ImportError(
                "flash_attn is required for attn_implementation='flash_attention_2'. "
                "Install with: pip install flash-attn --no-build-isolation"
            ) from e

        B, T, _ = hidden_states.shape
        q = self._split_heads(self.query(hidden_states)).permute(0, 2, 1, 3)
        k = self._split_heads(self.key(hidden_states)).permute(0, 2, 1, 3)
        v = self._split_heads(self.value(hidden_states)).permute(0, 2, 1, 3)

        if q.dtype not in (torch.float16, torch.bfloat16):
            raise ValueError(
                "flash_attention_2 requires float16 or bfloat16 model weights. "
                f"Received {q.dtype}."
            )

        if key_padding_mask is not None and key_padding_mask.any():
            attend = ~key_padding_mask
            q_u, indices, cu_seqlens, max_seqlen, _ = unpad_input(q, attend)
            k_u, _, _, _, _ = unpad_input(k, attend)
            v_u, _, _, _, _ = unpad_input(v, attend)
            out_u = flash_attn_varlen_func(
                q_u, k_u, v_u,
                cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
                max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen,
                dropout_p=self.dropout.p if self.training else 0.0,
                causal=False,
            )
            out = pad_input(out_u, indices, B, T)
        else:
            out = flash_attn_func(
                q,
                k,
                v,
                dropout_p=self.dropout.p if self.training else 0.0,
                causal=False,
            )

        out = out.reshape(B, T, self.all_head_size)
        return out, None


BERT_SELF_ATTENTION_CLASSES = {
    "eager": BertSelfAttention,
    "sdpa": BertSdpaSelfAttention,
    "flash_attention_2": BertFlashSelfAttention,
}


class BertSelfOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dropout(self.dense(hidden_states))
        return self.LayerNorm(hidden_states + input_tensor)


class BertAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        attn_cls = BERT_SELF_ATTENTION_CLASSES[getattr(config, "_attn_implementation", "eager")]
        self.self = attn_cls(config)
        self.output = BertSelfOutput(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor],
        output_attentions: bool = False,
        head_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        self_out, attn_weights = self.self(
            hidden_states,
            key_padding_mask,
            output_attentions,
            head_mask,
        )
        return self.output(self_out, hidden_states), attn_weights


class BertIntermediate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if config.hidden_act in (None, "identity", "linear"):
            self.intermediate_act_fn = lambda value: value
        elif isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.intermediate_act_fn(self.dense(hidden_states))


class BertOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dropout(self.dense(hidden_states))
        return self.LayerNorm(hidden_states + input_tensor)


class BertLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attention = BertAttention(config)
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor],
        output_attentions: bool = False,
        head_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        attn_out, attn_weights = self.attention(
            hidden_states,
            key_padding_mask,
            output_attentions,
            head_mask,
        )
        return self.output(self.intermediate(attn_out), attn_out), attn_weights


class BertEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor],
        output_hidden_states: bool = False,
        output_attentions: bool = False,
        head_mask: Optional[torch.Tensor] = None,
    ) -> Tuple:
        all_hidden_states = (hidden_states,) if output_hidden_states else None
        all_attentions = () if output_attentions else None

        for layer_index, layer in enumerate(self.layer):
            layer_head_mask = (
                head_mask[layer_index] if head_mask is not None else None
            )
            if self.gradient_checkpointing and self.training:
                hidden_states, attn_weights = torch.utils.checkpoint.checkpoint(
                    layer.__call__,
                    hidden_states,
                    key_padding_mask,
                    output_attentions,
                    layer_head_mask,
                    use_reentrant=False,
                )
            else:
                hidden_states, attn_weights = layer(
                    hidden_states,
                    key_padding_mask,
                    output_attentions,
                    layer_head_mask,
                )
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)
            if output_attentions:
                all_attentions = all_attentions + (attn_weights,)

        return hidden_states, all_hidden_states, all_attentions


class BertEmbeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
    ) -> torch.Tensor:
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
        B, T = (
            input_ids.shape
            if input_ids is not None
            else inputs_embeds.shape[:2]
        )
        if token_type_ids is None:
            token_type_ids = torch.zeros(
                (B, T),
                dtype=torch.long,
                device=(
                    input_ids.device
                    if input_ids is not None
                    else inputs_embeds.device
                ),
            )
        if position_ids is None:
            position_ids = self.position_ids[:, :T]
        x = (
            self.word_embeddings(input_ids)
            if inputs_embeds is None
            else inputs_embeds
        )
        x = x + self.token_type_embeddings(token_type_ids)
        x = x + self.position_embeddings(position_ids)
        return self.dropout(self.LayerNorm(x))


class BertPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.activation(self.dense(hidden_states[:, 0]))


class BertPredictionHeadTransform(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        hidden_act = (
            config.mlm_hidden_act
            if getattr(config, "mlm_hidden_act", None) is not None
            else config.hidden_act
        )
        if hidden_act in (None, "identity", "linear"):
            self.transform_act_fn = lambda value: value
        elif isinstance(hidden_act, str):
            self.transform_act_fn = ACT2FN[hidden_act]
        else:
            self.transform_act_fn = hidden_act

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.LayerNorm(self.transform_act_fn(self.dense(hidden_states)))


class BertModel(PreTrainedModel):
    config_class = BertUpdatedConfig
    base_model_prefix = "bert"
    _supports_sdpa = True
    _supports_flash_attn_2 = True
    supports_gradient_checkpointing = True
    _keys_to_ignore_on_load_missing = [r"pooler\."]

    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.embeddings = BertEmbeddings(config)
        self.encoder = BertEncoder(config)
        self.pooler = BertPooler(config) if add_pooling_layer else None
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
        input_shape = (
            input_ids.shape
            if input_ids is not None
            else inputs_embeds.shape[:2]
        )
        input_device = (
            input_ids.device
            if input_ids is not None
            else inputs_embeds.device
        )
        if attention_mask is None:
            attention_mask = torch.ones(input_shape, device=input_device)
        key_padding_mask = attention_mask.eq(0)
        if not key_padding_mask.any():
            key_padding_mask = None
        head_mask = self.get_head_mask(
            head_mask,
            self.config.num_hidden_layers,
        )
        if all(mask is None for mask in head_mask):
            head_mask = None

        x = self.embeddings(
            input_ids=input_ids,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
        )
        last_hidden_state, all_hidden_states, all_attentions = self.encoder(
            x, key_padding_mask,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            head_mask=head_mask,
        )
        pooled = self.pooler(last_hidden_state) if self.pooler is not None else None

        if not return_dict:
            output = (last_hidden_state, pooled)
            if output_hidden_states:
                output += (all_hidden_states,)
            if output_attentions:
                output += (all_attentions,)
            return output

        return BaseModelOutputWithPooling(
            last_hidden_state=last_hidden_state,
            pooler_output=pooled,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
        )


class BertForMaskedLM(PreTrainedModel):
    config_class = BertUpdatedConfig
    base_model_prefix = "bert"
    _supports_sdpa = True
    _supports_flash_attn_2 = True
    supports_gradient_checkpointing = True
    _keys_to_ignore_on_load_unexpected = [r"bert\.pooler\."]

    def __init__(self, config):
        super().__init__(config)
        self.bert = BertModel(config, add_pooling_layer=False)
        self.transform = BertPredictionHeadTransform(config)
        self.cls = nn.Linear(config.hidden_size, config.vocab_size)
        self.post_init()

    def get_input_embeddings(self):
        return self.bert.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.bert.embeddings.word_embeddings = value

    def get_output_embeddings(self):
        return self.cls

    def set_output_embeddings(self, value):
        self.cls = value

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, MaskedLMOutput]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            return_dict=True,
        )
        logits = self.cls(self.transform(outputs.last_hidden_state))

        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100)

        if not return_dict:
            output = (logits,)
            if output_hidden_states:
                output += (outputs.hidden_states,)
            if output_attentions:
                output += (outputs.attentions,)
            return (loss,) + output if loss is not None else output

        return MaskedLMOutput(
            loss=loss, logits=logits,
            hidden_states=outputs.hidden_states, attentions=outputs.attentions,
        )