Upload modeling_fastesm.py with huggingface_hub

modeling_fastesm.py

@@ -364,7 +364,6 @@ from typing import Optional, Tuple, Union, Dict, Any
 from einops import rearrange
 from dataclasses import dataclass
 from transformers import PreTrainedModel, PretrainedConfig, EsmTokenizer
-from transformers import initialization as init
 from transformers.modeling_outputs import (
     ModelOutput,
     BaseModelOutputWithPastAndCrossAttentions,
@@ -399,9 +398,9 @@ def get_attention_mask(
     attention_mask: Optional[torch.Tensor] = None
 ) -> torch.Tensor:
     if attention_mask is None:
-        token_attention_mask = torch.ones((batch_size, seq_len), device=device).bool() 
+        attention_mask_2d = torch.ones((batch_size, seq_len), device=device).bool() 
     else:
-        token_attention_mask = attention_mask.bool()
+        attention_mask_2d = attention_mask.bool()
     
     if attn_backend == "flex":
         assert create_block_mask is not None, "Flex attention backend requested but torch.create_block_mask is unavailable."
@@ -409,8 +408,10 @@ def get_attention_mask(
         if attention_mask is None:
             flex_block_mask = None
         else:
+            valid_lens = attention_mask_2d.sum(dim=-1)
+
             def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
-                return (token_attention_mask[batch_idx, q_idx] == token_attention_mask[batch_idx, kv_idx]) & (token_attention_mask[batch_idx, q_idx] != 0)
+                return (q_idx < valid_lens[batch_idx]) & (kv_idx < valid_lens[batch_idx])
     
             flex_block_mask = create_block_mask(
                 mask_mod,
@@ -420,12 +421,12 @@ def get_attention_mask(
                 seq_len,
                 device=device,
             )
-        extended_attention_mask = None
+        attention_mask_4d = None
     else:
         flex_block_mask = None
-        extended_attention_mask = token_attention_mask[:, None, :, None] & token_attention_mask[:, None, None, :]
+        attention_mask_4d = attention_mask_2d[:, None, :, None] & attention_mask_2d[:, None, None, :]
 
-    return extended_attention_mask, flex_block_mask
+    return attention_mask_4d, flex_block_mask
 
 
 @dataclass
@@ -763,16 +764,19 @@ class FastEsmPreTrainedModel(PreTrainedModel):
         return True
 
     @torch.no_grad()
-    def _init_weights(self, module):
-        """Initialize the weights"""
-        super()._init_weights(module)
-        if isinstance(module, EsmLMHead):
-            init.zeros_(module.bias)
-        elif isinstance(module, EsmEmbeddings):
-            init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
-        elif isinstance(module, RotaryEmbedding):
-            inv_freq = 1.0 / (10000 ** (torch.arange(0, module.dim, 2, dtype=torch.int64).float() / module.dim))
-            init.copy_(module.inv_freq, inv_freq)
+    def _init_weights(self, module: nn.Module) -> None:
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def post_init(self) -> None:
+        super().post_init()
 
     def get_output_embeddings(self):
         # NOTE: get_output_embeddings() must return None to prevent accidental weight tying.
@@ -809,7 +813,7 @@ class FAST_ESM_ENCODER(FastEsmPreTrainedModel, EmbeddingMixin):
 
     def _embed(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
         token_embedding_output = self.embeddings(input_ids, attention_mask=attention_mask)
-        attention_mask, flex_block_mask = get_attention_mask(
+        attention_mask_4d, flex_block_mask = get_attention_mask(
             attn_backend=self.config.attn_backend,
             batch_size=input_ids.shape[0],
             seq_len=input_ids.shape[1],
@@ -818,7 +822,7 @@ class FAST_ESM_ENCODER(FastEsmPreTrainedModel, EmbeddingMixin):
         )
         encoder_outputs = self.encoder(
             token_embedding_output,
-            attention_mask=attention_mask,
+            attention_mask=attention_mask_4d,
             flex_block_mask=flex_block_mask,
             output_hidden_states=False,
             output_attentions=False,
@@ -874,7 +878,7 @@ class FAST_ESM_ENCODER(FastEsmPreTrainedModel, EmbeddingMixin):
             attention_mask=attention_mask,
             inputs_embeds=inputs_embeds,
         )
-        attention_mask, flex_block_mask = get_attention_mask(
+        attention_mask_4d, flex_block_mask = get_attention_mask(
             attn_backend=self.config.attn_backend,
             batch_size=input_ids.shape[0],
             seq_len=input_ids.shape[1],
@@ -883,7 +887,7 @@ class FAST_ESM_ENCODER(FastEsmPreTrainedModel, EmbeddingMixin):
         )
         encoder_outputs = self.encoder(
             token_embedding_output,
-            attention_mask=attention_mask,
+            attention_mask=attention_mask_4d,
             flex_block_mask=flex_block_mask,
             output_hidden_states=output_hidden_states,
             output_attentions=output_attentions,