import torch
|
|
class BatchSampler(torch.utils.data.BatchSampler):
|
|
def __init__(self, dataset=None, args=None, drop_last=True, ):
|
|
self.drop_last = drop_last
|
self.pre_idx = -1
|
self.dataset = dataset
|
self.batch_size_type = args.batch_size_type
|
self.batch_size = args.batch_size
|
self.sort_size = args.sort_size
|
self.max_length_token = args.max_length_token
|
self.total_samples = len(dataset)
|
|
|
def __len__(self):
|
return self.total_samples
|
|
|
def __iter__(self):
|
batch = []
|
max_token = 0
|
num_sample = 0
|
|
iter_num = (self.total_samples-1) // self.sort_size + 1
|
for iter in range(self.pre_idx + 1, iter_num):
|
datalen_with_index = []
|
for i in range(self.sort_size):
|
idx = iter * self.sort_size + i
|
if idx >= self.total_samples:
|
continue
|
|
if self.batch_size_type == "example":
|
sample_len_cur = 1
|
else:
|
idx_map = self.dataset.shuffle_idx[idx]
|
# prompt = self.dataset.indexed_dataset[idx_map]["prompt"]
|
sample_len_cur = self.dataset.indexed_dataset[idx_map]["source_len"] + \
|
self.dataset.indexed_dataset[idx_map]["target_len"]
|
|
datalen_with_index.append([idx, sample_len_cur])
|
|
datalen_with_index_sort = sorted(datalen_with_index, key=lambda x: x[1])
|
for item in datalen_with_index_sort:
|
idx, sample_len_cur = item
|
if sample_len_cur > self.max_length_token:
|
continue
|
max_token_cur = max(max_token, sample_len_cur)
|
max_token_padding = (1 + num_sample) * max_token_cur
|
if max_token_padding <= self.batch_size:
|
batch.append(idx)
|
max_token = max_token_cur
|
num_sample += 1
|
else:
|
yield batch
|
max_token = sample_len_cur
|
num_sample = 1
|
batch = [idx]
|
|
