Update modeling_deberta.py

modeling_deberta.py

@@ -1058,8 +1058,6 @@ class DebertaV2Model(DebertaV2PreTrainedModel):
         )
         encoded_layers = list(encoder_outputs[1])
 
-#        print(self.z_steps)
-
         if self.z_steps > 0:
             hidden_states = encoded_layers[-2]
             layers = [self.encoder.layer[-1] for _ in range(self.z_steps)]
@@ -1100,8 +1098,6 @@ class DebertaV2ForMaskedLM(DebertaV2PreTrainedModel):
         self.deberta = DebertaV2Model(config)
         self.cls = DebertaV2OnlyMLMHead(config)
 
-        self.verbose = False
-
         # Initialize weights and apply final processing
         self.post_init()
 
@@ -1132,19 +1128,6 @@ class DebertaV2ForMaskedLM(DebertaV2PreTrainedModel):
 
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        if self.verbose:
-            for i in input_ids[0, :].tolist():
-                print(i, end=", ")
-            print()
-            if attention_mask is not None:
-                for i in attention_mask[0, :].tolist():
-                    print(i, end=", ")
-                print()
-            if position_ids is not None:
-                for i in position_ids[0, :].tolist():
-                    print(i, end=", ")
-                print()
-
         outputs = self.deberta(
             input_ids,
             attention_mask=attention_mask,
@@ -1183,6 +1166,7 @@ class DebertaV2ForCausalLM(DebertaV2ForMaskedLM):
         super().__init__(config)
         config.is_decoder = True
         self.mask_token_id = config.mask_token_id
+        self.cls_token_id = config.cls_token_id
         self.sep_token_id = config.sep_token_id
         self.n_masks = 3
 
@@ -1200,12 +1184,39 @@ class DebertaV2ForCausalLM(DebertaV2ForMaskedLM):
     ):
         position_ids = kwargs.get("position_ids", None)
 
-        if input_ids[0, -1] == 2:
-            input_ids = input_ids[:, :-1]
-            if attention_mask is not None:
-                attention_mask = attention_mask[:, :-1]
-            if position_ids is not None:
-                position_ids = position_ids[:, :-1]
+        assert input_ids[0, 0] != self.cls_token_id, "`add_special_tokens` should be set to `False`, but `[CLS]` token was detected"
+        assert input_ids[0, -1] != self.sep_token_id, "`add_special_tokens` should be set to `False`, but `[SEP]` token was detected"
+
+        batch_size, seq_length = input_ids.shape
+        input_ids = torch.cat(
+            [
+                torch.full((batch_size, 1), self.cls_token_id, device=input_ids.device)
+                input_ids,
+                torch.full((batch_size, self.n_masks), self.mask_token_id, device=input_ids.device),
+                torch.full((batch_size, 1), self.sep_token_id, device=input_ids.device)
+            ],
+            dim=-1
+        )
+
+        if attention_mask is not None:
+            attention_mask = torch.cat(
+                [
+                    torch.full((batch_size, 1), attention_mask[0, 0], device=attention_mask.device),
+                    attention_mask,
+                    torch.full((batch_size, self.n_masks + 1), attention_mask[0, -1], device=attention_mask.device),
+                ],
+                dim=-1
+            )
+
+        if position_ids is not None:
+            position_ids = torch.cat(
+                [
+                    torch.zeros(batch_size, 1, device=position_ids.device),
+                    position_ids + 1,
+                    torch.arange(0, self.n_masks + 1, device=position_ids.device).unsqueeze(0) + position_ids[:, -1:] + 1,
+                ],
+                dim=-1
+            )
 
         # Omit tokens covered by past_key_values
         if past_key_values is not None:
@@ -1228,7 +1239,7 @@ class DebertaV2ForCausalLM(DebertaV2ForMaskedLM):
             {
                 "position_ids": position_ids,
                 "past_key_values": past_key_values,
-                "use_cache": kwargs.get("use_cache"),
+                "use_cache": None,
                 "attention_mask": attention_mask,
             }
         )
@@ -1255,36 +1266,6 @@ class DebertaV2ForCausalLM(DebertaV2ForMaskedLM):
         assert past_key_values is None, "past_key_values is not supported for now"
         assert use_cache is None, "use_cache is not supported for now"
 
-        assert input_ids[0, -1] != self.sep_token_id, "remove the last token if it is a sep token"
-         
-        batch_size, seq_length = input_ids.shape
-        input_ids = torch.cat(
-            [
-                input_ids,
-                torch.full((batch_size, self.n_masks), self.mask_token_id, device=input_ids.device),
-                torch.full((batch_size, 1), self.sep_token_id, device=input_ids.device)
-            ],
-            dim=-1
-        )
-
-        if attention_mask is not None:
-            attention_mask = torch.cat(
-                [
-                    attention_mask,
-                    torch.full((batch_size, self.n_masks + 1), attention_mask[0, -1], device=attention_mask.device),
-                ],
-                dim=-1
-            )
-
-        if position_ids is not None:
-            position_ids = torch.cat(
-                [
-                    position_ids,
-                    torch.arange(0, self.n_masks + 1, device=position_ids.device).unsqueeze(0) + position_ids[:, -1:],
-                ],
-                dim=-1
-            )
-
         outputs = super().forward(
             input_ids,
             attention_mask=attention_mask,
@@ -1297,7 +1278,7 @@ class DebertaV2ForCausalLM(DebertaV2ForMaskedLM):
         )
 
         # shift the outputs and skip excess masks
-        logits = outputs.logits[:, 1:-(self.n_masks), :].contiguous()
+        logits = outputs.logits[:, 2:-self.n_masks, :].contiguous()
 
         loss = None
         if labels is not None: