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, )