Update modeling_codesage.py

modeling_codesage.py

@@ -11,7 +11,11 @@ from transformers.activations import ACT2FN
 from transformers.modeling_utils import Conv1D, PreTrainedModel
 from transformers.utils import logging
 from .config_codesage import CodeSageConfig
-from transformers.modeling_outputs import BaseModelOutputWithPooling, SequenceClassifierOutput
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPooling,
+    MaskedLMOutput,
+    SequenceClassifierOutput
+)
 
 logger = logging.get_logger(__name__)
 
@@ -151,6 +155,7 @@ class CodeSageBlock(nn.Module):
 
 class CodeSagePreTrainedModel(PreTrainedModel):
     config_class = CodeSageConfig
+    base_model_prefix = "transformer"
 
     def _init_weights(self, module):
         """Initialize the weights."""
@@ -277,7 +282,72 @@ class CodeSageModel(CodeSagePreTrainedModel):
         )
 
 
+class CodeSageForMaskedLM(CodeSagePreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = CodeSageModel(config)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        self.init_weights()
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(
+            self,
+            input_ids=None,
+            attention_mask=None,
+            position_ids=None,
+            head_mask=None,
+            inputs_embeds=None,
+            labels=None,
+            output_attentions=None,
+            output_hidden_states=None,
+            return_dict=None
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        transformer_outputs = self.transformer(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+        hidden_states = transformer_outputs[0]
+        lm_logits = self.lm_head(hidden_states)
+
+        masked_lm_loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            masked_lm_loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
+
+        if not return_dict:
+            output = (lm_logits,) + transformer_outputs[1:]
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=lm_logits,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )
+
+
 class CodeSageForSequenceClassification(CodeSagePreTrainedModel):
+
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels