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