Upload modeling_dplm2.py with huggingface_hub

modeling_dplm2.py

@@ -419,9 +419,9 @@ def get_attention_mask(
     attention_mask: Optional[torch.Tensor] = None,
 ) -> Tuple[Optional[torch.Tensor], Optional[object]]:
     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."
@@ -429,8 +429,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,
@@ -440,12 +442,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
 
 
 def _infer_modality_type(input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
@@ -989,12 +991,12 @@ class FAST_DPLM2_ENCODER(DPLM2PreTrainedModel, EmbeddingMixin):
         past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
 
         if attention_mask is None:
-            token_attention_mask = torch.ones((batch_size, seq_length + past_key_values_length), device=device).bool()
+            attention_mask_2d = torch.ones((batch_size, seq_length + past_key_values_length), device=device).bool()
         elif attention_mask.dim() == 2:
-            token_attention_mask = attention_mask.bool()
+            attention_mask_2d = attention_mask.bool()
         elif attention_mask.dim() == 4:
             assert input_ids is not None, "4D attention_mask requires input_ids to infer token-level mask."
-            token_attention_mask = input_ids.ne(self.config.pad_token_id)
+            attention_mask_2d = input_ids.ne(self.config.pad_token_id)
         else:
             raise ValueError(f"Unsupported attention_mask shape: {attention_mask.shape}")
 
@@ -1009,19 +1011,19 @@ class FAST_DPLM2_ENCODER(DPLM2PreTrainedModel, EmbeddingMixin):
 
         head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
 
-        embedding_attention_mask = token_attention_mask
+        embedding_attention_mask = attention_mask_2d
         if embedding_attention_mask is None and input_ids is not None:
             embedding_attention_mask = input_ids.ne(self.config.pad_token_id)
 
         if self.config.attn_backend == "flex" and output_attentions:
             raise AssertionError("output_attentions=True is not supported with attn_backend='flex'.")
 
-        extended_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=batch_size,
             seq_len=seq_length,
             device=device,
-            attention_mask=token_attention_mask,
+            attention_mask=attention_mask_2d,
         )
 
         embedding_output = self.embeddings(
@@ -1032,7 +1034,7 @@ class FAST_DPLM2_ENCODER(DPLM2PreTrainedModel, EmbeddingMixin):
         )
         encoder_outputs = self.encoder(
             embedding_output,
-            attention_mask=extended_attention_mask,
+            attention_mask=attention_mask_4d,
             head_mask=head_mask,
             encoder_hidden_states=encoder_hidden_states,
             encoder_attention_mask=encoder_extended_attention_mask,
@@ -1134,7 +1136,7 @@ class DPLM2ForMaskedLM(DPLM2PreTrainedModel, EmbeddingMixin):
         if vocab_size is not None:
             config.vocab_size = vocab_size
         DPLM2PreTrainedModel.__init__(self, config)
-        self.esm = DPLM2Model(config, add_pooling_layer=False)
+        self.esm = FAST_DPLM2_ENCODER(config)
         self.lm_head = EsmLMHead(config)
         self.loss_fct = nn.CrossEntropyLoss()
         self.post_init()
@@ -1235,7 +1237,7 @@ class DPLM2ForSequenceClassification(DPLM2PreTrainedModel, EmbeddingMixin):
     def __init__(self, config):
         DPLM2PreTrainedModel.__init__(self, config)
         self.num_labels = config.num_labels
-        self.esm = DPLM2Model(config, add_pooling_layer=False)
+        self.esm = FAST_DPLM2_ENCODER(config)
         self.classifier = EsmClassificationHead(config)
         self.mse = nn.MSELoss()
         self.ce = nn.CrossEntropyLoss()
@@ -1312,7 +1314,7 @@ class DPLM2ForTokenClassification(DPLM2PreTrainedModel, EmbeddingMixin):
     def __init__(self, config):
         DPLM2PreTrainedModel.__init__(self, config)
         self.num_labels = config.num_labels
-        self.esm = DPLM2Model(config, add_pooling_layer=False)
+        self.esm = FAST_DPLM2_ENCODER(config)
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
         self.loss_fct = nn.CrossEntropyLoss()