create

游雁

2022-11-26 c087854f71960341933a71442583dbc53d9b4e14

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
import random
 
from itertools import count
from functools import partial
from torch.utils.data import IterableDataset
from funasr.datasets.large_datasets.datapipes.map import MapperIterDataPipe
 
tiebreaker = count()
 
 
def _default_len_fn(token):
    return len(token), next(tiebreaker)
 
 
def _token_len_fn(token, len_fn):
    return len_fn(token), next(tiebreaker), token
 
 
class MaxTokenBucketizerIterDataPipe(IterableDataset):
 
    def __init__(
            self,
            datapipe,
            batch_size=8000,
            len_fn=_default_len_fn,
            buffer_size=10240,
            sort_size=500
    ):
        assert batch_size > 0, "Batch size is required to be larger than 0!"
        assert buffer_size >= -1, "Buffer size is required to be larger than -1!"
        assert sort_size > 0, "Sort size is required to be larger than 0!"
 
        datapipe = MapperIterDataPipe(datapipe, fn=partial(_token_len_fn, len_fn=len_fn))
        self.datapipe = datapipe
        self.batch_size = batch_size
        self.buffer_size = buffer_size
        self.sort_size = sort_size
 
    def set_epoch(self, epoch):
        self.epoch = epoch
 
    def __iter__(self):
        buffer = []
        batch = []
        bucket = []
        max_lengths = 0
        batch_lengths = 0
 
        if self.buffer_size == -1:
            for d in self.datapipe:
                if d[0] > self.batch_size:
                    continue
                buffer.append(d)
            buffer.sort()
            for sample in buffer:
                length, _, token = sample
                if length > max_lengths:
                    max_lengths = length
                batch_lengths = max_lengths * (len(batch) + 1)
                if batch_lengths > self.batch_size:
                    bucket.append(batch)
                    batch = []
                    max_lengths = length
                batch.append(token)
            random.shuffle(bucket)
            if bucket:
                for batch_sample in bucket:
                    yield batch_sample
            if batch:
                yield batch
 
        elif self.buffer_size == 0:
            for d in self.datapipe:
                if d[0] > self.batch_size:
                    continue
                length, _, token = d
                if length > self.batch_size:
                    continue
                if length > max_lengths:
                    max_lengths = length
                batch_lengths = max_lengths * (len(batch) + 1)
                if batch_lengths > self.batch_size:
                    yield batch
                    batch = []
                    max_lengths = length
                batch.append(token)
            if batch:
                yield batch
 
        else:
            for d in self.datapipe:
                if d[0] > self.batch_size:
                    continue
                buffer.append(d)
                if len(buffer) == self.buffer_size:
                    random.shuffle(buffer)
                    for sample in buffer:
                        bucket.append(sample)
                        if len(bucket) == self.sort_size:
                            bucket.sort()
                            for x in bucket:
                                length, _, token = x
                                if length > max_lengths:
                                    max_lengths = length
                                batch_lengths = max_lengths * (len(batch) + 1)
                                if batch_lengths > self.batch_size:
                                    yield batch
                                    batch = []
                                    max_lengths = length
                                batch.append(token)
                            bucket = []
                    buffer = []
 
            if buffer:
                random.shuffle(buffer)
                for sample in buffer:
                    bucket.append(sample)
                    if len(bucket) == self.sort_size:
                        bucket.sort()
                        for x in bucket:
                            length, _, token = x
                            if length > max_lengths:
                                max_lengths = length
                            batch_lengths = max_lengths * (len(batch) + 1)
                            if batch_lengths > self.batch_size:
                                yield batch
                                batch = []
                                max_lengths = length
                            batch.append(token)
                        bucket = []
                buffer = []
 
            if bucket:
                bucket.sort()
                for x in bucket:
                    length, _, token = x
                    if length > max_lengths:
                        max_lengths = length
                    batch_lengths = max_lengths * (len(batch) + 1)
                    if batch_lengths > self.batch_size:
                        yield batch
                        batch = []
                        max_lengths = length
                    batch.append(token)
                bucket = []
 
            if batch:
                yield batch