Update modeling_custom_llama.py

modeling_custom_llama.py

@@ -1,4 +1,5 @@
 from transformers.models.llama.modeling_llama import LlamaForCausalLM
+from transformers import MODEL_FOR_MASKED_LM_MAPPING
 from transformers import PretrainedConfig
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from transformers import GPT2TokenizerFast
@@ -122,26 +123,6 @@ class CustomLlamaAttention(LlamaAttention):
 
         attn_scores = self._compute_metric(q_start, q_dir, k_start, k_dir)
 
-        # # Handle attention mask and causality
-        # if attention_mask is not None:
-        #     # Convert padding mask [batch_size, seq_len] to [batch_size, 1, 1, seq_len]
-        #     padding_mask = attention_mask.unsqueeze(1).unsqueeze(2)
-        #     padding_mask = (1.0 - padding_mask) * torch.finfo(attn_scores.dtype).min
-        #     if is_causal is not None:
-        #         causal_mask = self._get_causal_mask(seq_len, seq_len, attn_scores.device)
-        #         is_causal_expanded = is_causal.view(-1, 1, 1, 1)
-        #         attention_mask = padding_mask + (causal_mask * is_causal_expanded)
-        #     else:
-        #         attention_mask = padding_mask
-        # else:
-        #     if is_causal is not None:
-        #         causal_mask = self._get_causal_mask(seq_len, seq_len, attn_scores.device)
-        #         is_causal_expanded = is_causal.view(-1, 1, 1, 1)
-        #         attention_mask = causal_mask * is_causal_expanded
-        #     else:
-        #         attention_mask = torch.zeros_like(attn_scores)
-
-        # attn_scores = attn_scores + attention_mask
         # Replace existing mask logic with:
         if attention_mask is not None:
             padding_mask = (attention_mask == 0).view(batch_size, 1, 1, -1)
@@ -324,8 +305,6 @@ class CustomLlamaForCausalLM(LlamaForCausalLM):
               )
 
         return ModelOutput(loss=loss, logits=logits)
-        # return {"loss": loss, "logits": logits}
-        # return {"loss": loss, "logits": logits} if return_dict else (loss, logits)
 
 class CustomLlamaForMaskedLM(CustomLlamaForCausalLM):
     config_class = CustomLlamaConfig  # Add this line
@@ -369,6 +348,6 @@ MODEL_MAPPING.update({"custom_llama": CustomLlamaForMaskedLM})
 def _register():
     from transformers import AutoConfig, AutoModelForCausalLM
     AutoConfig.register("custom_llama", CustomLlamaConfig)
-    # AutoModelForCausalLM.register(CustomLlamaConfig, CustomLlamaForCausalLM)
+    MODEL_FOR_MASKED_LM_MAPPING.register(CustomLlamaConfig, CustomLlamaForMaskedLM)
 
 _register()