add new datacollator upgrade examples (#1816)

* upgrade predict_glue.py

* test doc api

* revert

* update collator

* revert dureader_robus

* minor fix

* add glue sample

* stash

* stash

* upgrade examples for new datacollator

* fix squad sample

* add token classification collator

* fix doc

* fix collator for squad

* minor fix

* upgrade msra_ner predict

* add some doc

* add more check

Author: smallv0221

examples/information_extraction/msra_ner/predict.py

@@ -26,7 +26,7 @@
 
 import paddlenlp as ppnlp
 from datasets import load_dataset
-from paddlenlp.data import Stack, Tuple, Pad, Dict
+from paddlenlp.data import DataCollatorForTokenClassification
 from paddlenlp.transformers import BertForTokenClassification, BertTokenizer
 
 parser = argparse.ArgumentParser()
@@ -75,6 +75,7 @@ def do_predict(args):
     # Create dataset, tokenizer and dataloader.
     train_examples, predict_examples = load_dataset(
         'msra_ner', split=('train', 'test'))
+    column_names = train_examples.column_names
     tokenizer = BertTokenizer.from_pretrained(args.model_name_or_path)
 
     label_list = train_examples.features['ner_tags'].feature.names
@@ -104,17 +105,14 @@ def tokenize_and_align_labels(examples):
         return tokenized_inputs
 
     ignore_label = -100
-    batchify_fn = lambda samples, fn=Dict({
-        'input_ids': Pad(axis=0, pad_val=tokenizer.pad_token_id),  # input
-        'token_type_ids': Pad(axis=0, pad_val=tokenizer.pad_token_type_id),  # segment
-        'seq_len': Stack(),
-        'labels': Pad(axis=0, pad_val=ignore_label)  # label
-    }): fn(samples)
+    batchify_fn = DataCollatorForTokenClassification(tokenizer)
 
     id2label = dict(enumerate(label_list))
 
     predict_examples = predict_examples.select(range(len(predict_examples) - 1))
-    predict_ds = predict_examples.map(tokenize_and_align_labels, batched=True)
+    predict_ds = predict_examples.map(tokenize_and_align_labels,
+                                      batched=True,
+                                      remove_columns=column_names)
     predict_data_loader = DataLoader(
         dataset=predict_ds,
         collate_fn=batchify_fn,
@@ -133,11 +131,10 @@ def tokenize_and_align_labels(examples):
     pred_list = []
     len_list = []
     for step, batch in enumerate(predict_data_loader):
-        input_ids, token_type_ids, length, labels = batch
-        logits = model(input_ids, token_type_ids)
+        logits = model(batch['input_ids'], batch['token_type_ids'])
         pred = paddle.argmax(logits, axis=-1)
         pred_list.append(pred.numpy())
-        len_list.append(length.numpy())
+        len_list.append(batch['seq_len'].numpy())
 
     preds = parse_decodes(predict_examples, id2label, pred_list, len_list)
 

examples/information_extraction/msra_ner/train.py

@@ -30,7 +30,7 @@
 from paddlenlp.transformers import BertForTokenClassification, BertTokenizer
 from paddlenlp.transformers import ErnieForTokenClassification, ErnieTokenizer
 from paddlenlp.transformers import ErnieCtmForTokenClassification, ErnieCtmTokenizer
-from paddlenlp.data import Stack, Tuple, Pad, Dict
+from paddlenlp.data import DataCollatorForTokenClassification
 from paddlenlp.utils.log import logger
 
 MODEL_CLASSES = {
@@ -68,13 +68,12 @@ def evaluate(model, loss_fct, metric, data_loader, label_num, mode="valid"):
     metric.reset()
     avg_loss, precision, recall, f1_score = 0, 0, 0, 0
     for batch in data_loader:
-        input_ids, token_type_ids, length, labels = batch
-        logits = model(input_ids, token_type_ids)
-        loss = loss_fct(logits, labels)
+        logits = model(batch['input_ids'], batch['token_type_ids'])
+        loss = loss_fct(logits, batch['labels'])
         avg_loss = paddle.mean(loss)
         preds = logits.argmax(axis=2)
         num_infer_chunks, num_label_chunks, num_correct_chunks = metric.compute(
-            length, preds, labels)
+            batch['seq_len'], preds, batch['labels'])
         metric.update(num_infer_chunks.numpy(),
                       num_label_chunks.numpy(), num_correct_chunks.numpy())
         precision, recall, f1_score = metric.accumulate()
@@ -125,16 +124,14 @@ def tokenize_and_align_labels(examples):
         return tokenized_inputs
 
     train_ds = train_ds.select(range(len(train_ds) - 1))
-    train_ds = train_ds.map(tokenize_and_align_labels, batched=True)
+    column_names = train_ds.column_names
+    train_ds = train_ds.map(tokenize_and_align_labels,
+                            batched=True,
+                            remove_columns=column_names)
 
     ignore_label = -100
 
-    batchify_fn = lambda samples, fn=Dict({
-        'input_ids': Pad(axis=0, pad_val=tokenizer.pad_token_id, dtype='int32'),  # input
-        'token_type_ids': Pad(axis=0, pad_val=tokenizer.pad_token_type_id, dtype='int32'),  # segment
-        'seq_len': Stack(dtype='int64'),  # seq_len
-        'labels': Pad(axis=0, pad_val=ignore_label, dtype='int64')  # label
-    }): fn(samples)
+    batchify_fn = DataCollatorForTokenClassification(tokenizer, ignore_label)
 
     train_batch_sampler = paddle.io.DistributedBatchSampler(
         train_ds, batch_size=args.batch_size, shuffle=True, drop_last=True)
@@ -148,7 +145,9 @@ def tokenize_and_align_labels(examples):
 
     test_ds = raw_datasets['test']
     test_ds = test_ds.select(range(len(test_ds) - 1))
-    test_ds = test_ds.map(tokenize_and_align_labels, batched=True)
+    test_ds = test_ds.map(tokenize_and_align_labels,
+                          batched=True,
+                          remove_columns=column_names)
 
     test_data_loader = DataLoader(
         dataset=test_ds,
@@ -160,7 +159,9 @@ def tokenize_and_align_labels(examples):
     if args.dataset == "peoples_daily_ner":
         dev_ds = raw_datasets['validation']
         dev_ds = dev_ds.select(range(len(dev_ds) - 1))
-        dev_ds = dev_ds.map(tokenize_and_align_labels, batched=True)
+        dev_ds = dev_ds.map(tokenize_and_align_labels,
+                            batched=True,
+                            remove_columns=column_names)
 
         dev_data_loader = DataLoader(
             dataset=dev_ds,
@@ -205,9 +206,8 @@ def tokenize_and_align_labels(examples):
     for epoch in range(args.num_train_epochs):
         for step, batch in enumerate(train_data_loader):
             global_step += 1
-            input_ids, token_type_ids, _, labels = batch
-            logits = model(input_ids, token_type_ids)
-            loss = loss_fct(logits, labels)
+            logits = model(batch['input_ids'], batch['token_type_ids'])
+            loss = loss_fct(logits, batch['labels'])
             avg_loss = paddle.mean(loss)
             if global_step % args.logging_steps == 0:
                 print(

examples/language_model/bert/run_glue.py

@@ -27,7 +27,7 @@
 from paddle.metric import Metric, Accuracy, Precision, Recall
 
 from datasets import load_dataset
-from paddlenlp.data import Stack, Tuple, Pad, Dict
+from paddlenlp.data import default_data_collator, DataCollatorWithPadding
 from paddlenlp.data.sampler import SamplerHelper
 from paddlenlp.transformers import BertForSequenceClassification, BertTokenizer
 from paddlenlp.transformers import ElectraForSequenceClassification, ElectraTokenizer
@@ -196,10 +196,9 @@ def evaluate(model, loss_fct, metric, data_loader):
     model.eval()
     metric.reset()
     for batch in data_loader:
-        input_ids, segment_ids, labels = batch
-        logits = model(input_ids, segment_ids)
-        loss = loss_fct(logits, labels)
-        correct = metric.compute(logits, labels)
+        logits = model(batch['input_ids'], batch['token_type_ids'])
+        loss = loss_fct(logits, batch['labels'])
+        correct = metric.compute(logits, batch['labels'])
         metric.update(correct)
     res = metric.accumulate()
     if isinstance(metric, AccuracyAndF1):
@@ -266,11 +265,7 @@ def preprocess_function(examples):
                             remove_columns=columns)
     train_batch_sampler = paddle.io.DistributedBatchSampler(
         train_ds, batch_size=args.batch_size, shuffle=True)
-    batchify_fn = lambda samples, fn=Dict({
-        'input_ids': Pad(axis=0, pad_val=tokenizer.pad_token_id),  # input
-        'token_type_ids': Pad(axis=0, pad_val=tokenizer.pad_token_type_id),  # segment
-        'labels': Stack(dtype="int64" if label_list else "float32")  # label
-    }): fn(samples)
+    batchify_fn = DataCollatorWithPadding(tokenizer)
     train_data_loader = DataLoader(
         dataset=train_ds,
         batch_sampler=train_batch_sampler,
@@ -358,13 +353,11 @@ def preprocess_function(examples):
     for epoch in range(args.num_train_epochs):
         for step, batch in enumerate(train_data_loader):
             global_step += 1
-
-            input_ids, segment_ids, labels = batch
             with paddle.amp.auto_cast(
                     args.use_amp,
                     custom_white_list=["layer_norm", "softmax", "gelu"]):
-                logits = model(input_ids, segment_ids)
-                loss = loss_fct(logits, labels)
+                logits = model(batch['input_ids'], batch['token_type_ids'])
+                loss = loss_fct(logits, batch['labels'])
             if args.use_amp:
                 scaler.scale(loss).backward()
                 scaler.minimize(optimizer, loss)

examples/machine_reading_comprehension/SQuAD/run_squad.py

@@ -28,7 +28,7 @@
 
 import paddlenlp as ppnlp
 
-from paddlenlp.data import Pad, Stack, Tuple, Dict
+from paddlenlp.data import default_data_collator, DataCollatorWithPadding
 from paddlenlp.transformers import BertForQuestionAnswering, BertTokenizer, ErnieForQuestionAnswering, ErnieTokenizer, FunnelForQuestionAnswering, FunnelTokenizer
 from paddlenlp.transformers import LinearDecayWithWarmup
 from paddlenlp.metrics.squad import squad_evaluate, compute_prediction
@@ -170,19 +170,18 @@ def set_seed(args):
 
 
 @paddle.no_grad()
-def evaluate(model, data_loader, raw_dataset, args):
+def evaluate(model, data_loader, raw_dataset, features, args):
     model.eval()
 
     all_start_logits = []
     all_end_logits = []
     tic_eval = time.time()
 
     for batch in data_loader:
-        input_ids, token_type_ids, attention_mask = batch
         start_logits_tensor, end_logits_tensor = model(
-            input_ids,
-            token_type_ids=token_type_ids,
-            attention_mask=attention_mask)
+            batch['input_ids'],
+            token_type_ids=batch['token_type_ids'],
+            attention_mask=batch['attention_mask'])
 
         for idx in range(start_logits_tensor.shape[0]):
             if len(all_start_logits) % 1000 == 0 and len(all_start_logits):
@@ -194,7 +193,7 @@ def evaluate(model, data_loader, raw_dataset, args):
             all_end_logits.append(end_logits_tensor.numpy()[idx])
 
     all_predictions, all_nbest_json, scores_diff_json = compute_prediction(
-        raw_dataset, data_loader.dataset, (all_start_logits, all_end_logits),
+        raw_dataset, features, (all_start_logits, all_end_logits),
         args.version_2_with_negative, args.n_best_size, args.max_answer_length,
         args.null_score_diff_threshold)
 
@@ -262,13 +261,7 @@ def run(args):
                                       num_proc=4)
         train_batch_sampler = paddle.io.DistributedBatchSampler(
             train_ds, batch_size=args.batch_size, shuffle=True)
-        train_batchify_fn = lambda samples, fn=Dict({
-            "input_ids": Pad(axis=0, pad_val=tokenizer.pad_token_id),
-            "token_type_ids": Pad(axis=0, pad_val=tokenizer.pad_token_type_id),
-            'attention_mask': Pad(axis=0, pad_val=tokenizer.pad_token_type_id),
-            "start_positions": Stack(dtype="int64"),
-            "end_positions": Stack(dtype="int64")
-        }): fn(samples)
+        train_batchify_fn = DataCollatorWithPadding(tokenizer)
 
         train_data_loader = DataLoader(
             dataset=train_ds,
@@ -304,12 +297,12 @@ def run(args):
         for epoch in range(num_train_epochs):
             for step, batch in enumerate(train_data_loader):
                 global_step += 1
-                input_ids, token_type_ids, attention_mask, start_positions, end_positions = batch
                 logits = model(
-                    input_ids=input_ids,
-                    token_type_ids=token_type_ids,
-                    attention_mask=attention_mask)
-                loss = criterion(logits, (start_positions, end_positions))
+                    input_ids=batch['input_ids'],
+                    token_type_ids=batch['token_type_ids'],
+                    attention_mask=batch['attention_mask'])
+                loss = criterion(logits, (batch['start_positions'],
+                                          batch['end_positions']))
                 if global_step % args.logging_steps == 0:
                     print(
                         "global step %d, epoch: %d, batch: %d, loss: %f, speed: %.2f step/s"
@@ -344,20 +337,17 @@ def run(args):
                                   num_proc=4)
         dev_batch_sampler = paddle.io.BatchSampler(
             dev_ds, batch_size=args.batch_size, shuffle=False)
-
-        dev_batchify_fn = lambda samples, fn=Dict({
-            "input_ids": Pad(axis=0, pad_val=tokenizer.pad_token_id),
-            "token_type_ids": Pad(axis=0, pad_val=tokenizer.pad_token_type_id),
-            "attention_mask": Pad(axis=0, pad_val=tokenizer.pad_token_type_id)
-        }): fn(samples)
+        dev_ds_for_model = dev_ds.remove_columns(
+            ["example_id", "offset_mapping"])
+        dev_batchify_fn = DataCollatorWithPadding(tokenizer)
 
         dev_data_loader = DataLoader(
-            dataset=dev_ds,
+            dataset=dev_ds_for_model,
             batch_sampler=dev_batch_sampler,
             collate_fn=dev_batchify_fn,
             return_list=True)
 
-        evaluate(model, dev_data_loader, dev_examples, args)
+        evaluate(model, dev_data_loader, dev_examples, dev_ds, args)
 
 
 if __name__ == "__main__":