paddlenlp

1
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14

15
import numbers
16

17
import numpy as np
18
import paddle
19

20
try:
21
    from collections.abc import Mapping, Sequence
22
except:
23
    from collections import Sequence, Mapping
24

25
from ppfleetx.data.sampler import Stack, Tuple
26

27

28
def collate_fn(batch):
29
    """
30
    Default batch collating function for :code:`paddle.io.DataLoader`,
31
    get input data as a list of sample datas, each element in list
32
    if the data of a sample, and sample data should composed of list,
33
    dictionary, string, number, numpy array and paddle.Tensor, this
34
    function will parse input data recursively and stack number,
35
    numpy array and paddle.Tensor datas as batch datas. e.g. for
36
    following input data:
37
    [{'image': np.array(shape=[3, 224, 224]), 'label': 1},
38
     {'image': np.array(shape=[3, 224, 224]), 'label': 3},
39
     {'image': np.array(shape=[3, 224, 224]), 'label': 4},
40
     {'image': np.array(shape=[3, 224, 224]), 'label': 5},]
41

42

43
    This default collate function zipped each number and numpy array
44
    field together and stack each field as the batch field as follows:
45
    {'image': np.array(shape=[4, 3, 224, 224]), 'label': np.array([1, 3, 4, 5])}
46
    Args:
47
        batch(list of sample data): batch should be a list of sample data.
48

49
    Returns:
50
        Batched data: batched each number, numpy array and paddle.Tensor
51
                      in input data.
52
    """
53
    sample = batch[0]
54
    if isinstance(sample, np.ndarray):
55
        batch = np.stack(batch, axis=0)
56
        return batch
57
    elif isinstance(sample, paddle.Tensor):
58
        return paddle.stack(batch, axis=0)
59
    elif isinstance(sample, numbers.Number):
60
        batch = np.array(batch)
61
        return batch
62
    elif isinstance(sample, (str, bytes)):
63
        return batch
64
    elif isinstance(sample, Mapping):
65
        return {key: collate_fn([d[key] for d in batch]) for key in sample}
66
    elif isinstance(sample, Sequence):
67
        sample_fields_num = len(sample)
68
        if not all(len(sample) == sample_fields_num for sample in iter(batch)):
69
            raise RuntimeError("fileds number not same among samples in a batch")
70
        return [collate_fn(fields) for fields in zip(*batch)]
71

72
    raise TypeError(
73
        "batch data con only contains: tensor, numpy.ndarray, " "dict, list, number, but got {}".format(type(sample))
74
    )
75

76

77
def default_collate_fn(batch_transform=None):
78
    if batch_transform is not None:
79
        # batch_ops = create_preprocess_operators(batch_transform)
80

81
        # def inner_collate_fn(batch):
82
        #     batch = transform(batch, batch_ops)
83
        #     batch = collate_fn(batch)
84
        #     return batch
85

86
        # return inner_collate_fn
87
        pass
88
    else:
89
        return collate_fn
90

91

92
def gpt_collate_fn(batch):
93
    return Tuple([Stack() for raw in zip(*batch)])(batch)
94