fix causalLM

modeling_ltgbert.py

@@ -318,6 +318,7 @@ class LtgbertModel(LtgbertPreTrainedModel):
         self.transformer = Encoder(config)
         self.classifier = MaskClassifier(config, self.embedding.word_embedding.weight) if add_mlm_layer else None
 
+
     def get_input_embeddings(self):
         return self.embedding.word_embedding
 
@@ -414,7 +415,7 @@ class LtgbertForMaskedLM(LtgbertModel):
 
         sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
         subword_prediction = self.classifier(sequence_output)
-        subword_prediction[:, :, :16+1] = float("-inf")
+        # subword_prediction[:, :, :16+1] = float("-inf")
 
         masked_lm_loss = None
         if labels is not None:
@@ -443,7 +444,6 @@ class Classifier(nn.Module):
         super().__init__()
 
         self.temperature = config.temperature
-
         drop_out = getattr(config, "cls_dropout", None)
         drop_out = config.hidden_dropout_prob if drop_out is None else drop_out
 
@@ -494,6 +494,7 @@ class LtgbertForCausalLM(LtgbertModel):
         input_ids: torch.LongTensor = None,
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
         past_key_values = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         labels: Optional[torch.LongTensor] = None,
@@ -511,7 +512,7 @@ class LtgbertForCausalLM(LtgbertModel):
 
         sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
         subword_prediction = self.classifier(sequence_output)
-        subword_prediction[:, :, :16+1] = float("-inf")
+        # subword_prediction[:, :, :16+1] = float("-inf")
 
         masked_lm_loss = None
         if labels is not None: