Improve CodeGen

examples/code_generation/codegen/run_clm.py

@@ -252,7 +252,7 @@ def do_train(args):
                            block_size)
     dev_set = process_ds(dev_set, tokenizer, args.overwrite_cache, block_size)
 
-    batchify_fn = DataCollatorWithPadding(tokenizer)
+    batchify_fn = DataCollatorWithPadding(tokenizer, return_attention_mask=True)
 
     train_batch_sampler = DistributedBatchSampler(
         train_set, batch_size=args.train_batch_size, shuffle=True)

paddlenlp/data/data_collator.py

@@ -192,6 +192,7 @@ class DataCollatorWithPadding:
     max_length: Optional[int] = None
     pad_to_multiple_of: Optional[int] = None
     return_tensors: str = "pd"
+    return_attention_mask: Optional[bool] = None
 
     def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
         batch = self.tokenizer.pad(
@@ -200,7 +201,7 @@ def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
             max_length=self.max_length,
             pad_to_multiple_of=self.pad_to_multiple_of,
             return_tensors=self.return_tensors,
-        )
+            return_attention_mask=self.return_attention_mask)
         if "label" in batch:
             batch["labels"] = batch["label"]
             del batch["label"]

paddlenlp/transformers/codegen/modeling.py

@@ -412,6 +412,7 @@ def forward(
         self,
         input_ids=None,
         attention_mask=None,
+        token_type_ids=None,
         use_cache=False,
         cache=None,
     ):
@@ -472,9 +473,15 @@ def forward(
         if attention_mask is None:
             assert input_ids is not None, "input_ids should be " \
                                           "specified when generating attention_mask"
-            attention_mask = paddle.cast(
-                input_ids == self.pad_token_id,
-                dtype=paddle.get_default_dtype()).unsqueeze([1, 2]) * -1e4
+            if batch_size == 1 and past_length != 0:
+                batch_size, seq_len = input_shape
+                attention_mask = paddle.ones(
+                    [batch_size, 1, 1, seq_len + past_length],
+                    dtype=paddle.get_default_dtype())
+            else:
+                attention_mask = paddle.cast(
+                    input_ids == self.pad_token_id,
+                    dtype=paddle.get_default_dtype()).unsqueeze([1, 2]) * -1e4
         # For 2D attention_mask from tokenizer
         elif attention_mask.ndim == 2:
             attention_mask = paddle.unsqueeze(
@@ -483,6 +490,10 @@ def forward(
             attention_mask.stop_gradient = True
 
         inputs_embeds = self.wte(input_ids)
+        if token_type_ids is not None:
+            token_type_embeds = self.wte(token_type_ids)
+            inputs_embeds = inputs_embeds + token_type_embeds
+
         hidden_states = self.drop(inputs_embeds)
         output_shape = input_shape[:] + [hidden_states.shape[-1]]
 
@@ -579,6 +590,7 @@ def prepare_inputs_for_generation(self, input_ids, cache=None, **kwargs):
     def forward(self,
                 input_ids=None,
                 attention_mask=None,
+                token_type_ids=None,
                 use_cache=False,
                 cache=None):
         r"""
@@ -613,6 +625,7 @@ def forward(self,
 
         transformer_outputs = self.transformer(input_ids,
                                                attention_mask=attention_mask,
+                                               token_type_ids=token_type_ids,
                                                use_cache=use_cache,
                                                cache=cache)