wehub-resource-sync
348 lines
14 KiB
Python
348 lines
14 KiB
Python
import glob
|
|
import os
|
|
import torch
|
|
import numpy as np
|
|
import time
|
|
import json
|
|
import random
|
|
import itertools
|
|
import hydra
|
|
import copy
|
|
from omegaconf import DictConfig, OmegaConf
|
|
|
|
from infinibatch import iterators
|
|
from .basic_loader import BaseBatchGen
|
|
from .utils import NativeCheckpointableIterator, WeightIterator, EOL_SYMBOL
|
|
from .utils import safe_getattr, safe_hasattr
|
|
|
|
|
|
class LMLoader(BaseBatchGen):
|
|
|
|
def __init__(
|
|
self,
|
|
args,
|
|
dataset,
|
|
dictionary,
|
|
tokenizer,
|
|
max_tokens=None,
|
|
max_sentences=None,
|
|
max_positions=None,
|
|
ignore_invalid_inputs=False,
|
|
required_batch_size_multiple=1,
|
|
seed=1,
|
|
epoch=1,
|
|
num_shards=1,
|
|
shard_id=0,
|
|
disable_prefetching=False,
|
|
data_name='gpt',
|
|
):
|
|
super().__init__()
|
|
self.args = args
|
|
self.data = dataset.data
|
|
self.data_dir = dataset.data_dir
|
|
self.shuffle = dataset.shuffle
|
|
self.dictionary = dictionary
|
|
self.tokenizer = tokenizer
|
|
|
|
self.max_tokens = max_tokens
|
|
self.max_sentences = max_sentences
|
|
self.max_positions = max_positions
|
|
self.tokens_per_sample = args.tokens_per_sample
|
|
self.mlm_cut_length = safe_getattr(args, "mlm_cut_length", 0)
|
|
self.mlm_tokens_proportion = safe_getattr(args, "mlm_tokens_proportion", 0)
|
|
self.pad_to_max_len = safe_getattr(args, "pad_to_max_len", False)
|
|
self.ignore_invalid_inputs = ignore_invalid_inputs
|
|
self.required_batch_size_multiple = required_batch_size_multiple
|
|
self.seed = str(seed)
|
|
self.epoch = epoch
|
|
self.num_shards = num_shards
|
|
self.shard_id = shard_id
|
|
|
|
self.batch_read_ahead = args.batch_read_ahead
|
|
self.disable_prefetching = disable_prefetching
|
|
self.data_name = data_name
|
|
self._setup()
|
|
|
|
self._build_iter()
|
|
|
|
def _setup(self):
|
|
pass
|
|
|
|
def _build_iter(self):
|
|
tokenized_lines = self._tokenize()
|
|
self.padded_batches = self._batchify(tokenized_lines)
|
|
|
|
if self.disable_prefetching:
|
|
prefetch_batches = self.padded_batches
|
|
else:
|
|
prefetch_batches = iterators.PrefetchIterator(
|
|
self.padded_batches,
|
|
buffer_size=10000,
|
|
buffer_in_main_process=True,
|
|
log_empty_buffer_warning=True and self.shard_id == 0,
|
|
)
|
|
|
|
prefetch_batches = iterators.MapIterator(
|
|
prefetch_batches, self._move_to_tensor
|
|
)
|
|
|
|
self._iter = prefetch_batches
|
|
|
|
def _tokenize(self):
|
|
'''
|
|
data:
|
|
{
|
|
'source': list[Path],
|
|
}
|
|
'''
|
|
dataset = list(zip(self.data['source']))
|
|
|
|
if self.shuffle:
|
|
chunk_files = \
|
|
iterators.InfinitePermutationSourceIterator(
|
|
dataset,
|
|
seed=self.seed,
|
|
shuffle=self.shuffle,
|
|
num_instances=self.num_shards,
|
|
instance_rank=self.shard_id,
|
|
)
|
|
else:
|
|
chunk_files = \
|
|
iterators.ChunkedSourceIterator(
|
|
dataset,
|
|
num_instances=self.num_shards,
|
|
instance_rank=self.shard_id,
|
|
)
|
|
|
|
tokenized_lines = iterators.SelectManyIterator(chunk_files, lambda files: self._read_from_files(*files))
|
|
tokenized_lines = iterators.SamplingRandomMapIterator(tokenized_lines, self._prepare, self.seed)
|
|
|
|
return tokenized_lines
|
|
|
|
def getstate(self):
|
|
state = super().getstate()
|
|
state["epoch"] = self.epoch
|
|
state["iterations_in_epoch"] = None
|
|
return state
|
|
|
|
def _batchify(self, lines):
|
|
|
|
if self.max_sentences is not None:
|
|
if self.batch_read_ahead > 0:
|
|
lines = iterators.BlockwiseShuffleIterator(lines, self.batch_read_ahead, self.seed)
|
|
batches = iterators.FixedBatchIterator(lines, self.max_sentences)
|
|
else:
|
|
# -
|
|
def dynamic_batch_size(sample):
|
|
lengths = [len(x) for x in sample]
|
|
batch_size = self.max_tokens // max(lengths) // self.required_batch_size_multiple * self.required_batch_size_multiple
|
|
return max(1, batch_size)
|
|
|
|
batches = iterators.BucketedReadaheadBatchIterator(
|
|
lines,
|
|
read_ahead=self.batch_read_ahead,
|
|
key=(lambda x: max(len(x[0]), len(x[1]))) if self.shuffle else None,
|
|
batch_size=dynamic_batch_size,
|
|
shuffle=self.shuffle,
|
|
seed=self.seed,
|
|
)
|
|
|
|
def collate(batch):
|
|
batch_size = len(batch)
|
|
mlm_batch_size = sum([len(x[2]) for x in batch])
|
|
|
|
gpt_max_length = max([len(x[0]) for x in batch])
|
|
if self.pad_to_max_len:
|
|
gpt_max_length = self.tokens_per_sample
|
|
|
|
mlm_max_length = 0
|
|
mlm_ntokens = 0
|
|
for x in batch:
|
|
for y in x[2]:
|
|
mlm_max_length = max(mlm_max_length, len(y))
|
|
mlm_ntokens += len(y)
|
|
|
|
gpt_source_ids = np.full(shape=(batch_size, gpt_max_length-1), dtype=np.int32,
|
|
fill_value=self.dictionary.pad())
|
|
gpt_target_ids = np.full(shape=(batch_size, gpt_max_length-1), dtype=np.int32,
|
|
fill_value=self.dictionary.pad())
|
|
mlm_source_ids = np.full(shape=(mlm_batch_size, mlm_max_length), dtype=np.int32,
|
|
fill_value=self.dictionary.pad())
|
|
gpt_input_mask_all = np.full(shape=(batch_size, gpt_max_length-1), dtype=np.int32, fill_value=0)
|
|
gpt_loss_mask_all = np.full(shape=(batch_size, gpt_max_length-1), dtype=np.int32, fill_value=1)
|
|
mlm_mask_all = np.full(shape=(mlm_batch_size, mlm_max_length), dtype=np.int32, fill_value=0)
|
|
|
|
mlm_index = 0
|
|
for i, (gpt_ids, gpt_input_mask, mlm_ids_list, mlm_mask_list, gpt_loss_mask) in enumerate(batch):
|
|
gpt_source_ids[i, :len(gpt_ids)-1] = gpt_ids[:-1]
|
|
gpt_target_ids[i, :len(gpt_ids)-1] = gpt_ids[1:]
|
|
gpt_input_mask_all[i, :len(gpt_ids)-1] = gpt_input_mask[:-1]
|
|
gpt_loss_mask_all[i, :len(gpt_ids)-1] = gpt_loss_mask[1:]
|
|
|
|
for j, (mlm_ids, mlm_mask) in enumerate(zip(mlm_ids_list, mlm_mask_list)):
|
|
mlm_source_ids[mlm_index, :len(mlm_ids)] = mlm_ids
|
|
mlm_mask_all[mlm_index, :len(mlm_mask)] = mlm_mask
|
|
mlm_index += 1
|
|
|
|
ret_batch = {
|
|
'text':{
|
|
'net_input': {
|
|
'src_tokens': gpt_source_ids.astype(np.int64),
|
|
'mlm_src_tokens': mlm_source_ids.astype(np.int64) if mlm_batch_size !=0 else None,
|
|
'gpt_input_mask': gpt_input_mask_all.astype(np.bool_),
|
|
'gpt_loss_mask': gpt_loss_mask_all.astype(np.bool_),
|
|
'mlm_mask': mlm_mask_all.astype(np.bool_) if mlm_batch_size !=0 else None
|
|
},
|
|
'target': gpt_target_ids.astype(np.int64),
|
|
'nsentences': batch_size,
|
|
'ntokens': sum([len(x[0]) for x in batch]),
|
|
'mlm_ntokens': mlm_ntokens
|
|
}
|
|
}
|
|
|
|
return ret_batch
|
|
|
|
def collate_for_gpt(batch):
|
|
batch_size = len(batch)
|
|
gpt_max_length = max([len(x[0]) for x in batch])
|
|
if self.pad_to_max_len:
|
|
gpt_max_length = self.tokens_per_sample
|
|
|
|
gpt_source_ids = np.full(shape=(batch_size, gpt_max_length-1), dtype=np.int32,
|
|
fill_value=self.dictionary.pad())
|
|
gpt_target_ids = np.full(shape=(batch_size, gpt_max_length-1), dtype=np.int32,
|
|
fill_value=self.dictionary.pad())
|
|
gpt_input_mask_all = np.full(shape=(batch_size, gpt_max_length-1), dtype=np.int32, fill_value=0)
|
|
gpt_loss_mask_all = np.full(shape=(batch_size, gpt_max_length-1), dtype=np.int32, fill_value=1)
|
|
|
|
for i, (gpt_ids, gpt_input_mask, mlm_ids_list, mlm_mask_list, gpt_loss_mask) in enumerate(batch):
|
|
gpt_source_ids[i, :len(gpt_ids)-1] = gpt_ids[:-1]
|
|
gpt_target_ids[i, :len(gpt_ids)-1] = gpt_ids[1:]
|
|
gpt_input_mask_all[i, :len(gpt_ids)-1] = gpt_input_mask[:-1]
|
|
gpt_loss_mask_all[i, :len(gpt_ids)-1] = gpt_loss_mask[1:]
|
|
|
|
ret_batch = {
|
|
self.data_name:{
|
|
'net_input': {
|
|
'src_tokens': gpt_source_ids.astype(np.int64),
|
|
},
|
|
'target': gpt_target_ids.astype(np.int64),
|
|
'nsentences': batch_size,
|
|
'ntokens': sum([len(x[0]) for x in batch]),
|
|
'mlm_ntokens': 0
|
|
}
|
|
}
|
|
|
|
return ret_batch
|
|
|
|
if self.mlm_tokens_proportion == 0:
|
|
padded_batches = iterators.MapIterator(
|
|
batches, collate_for_gpt
|
|
)
|
|
else:
|
|
padded_batches = iterators.MapIterator(
|
|
batches, collate
|
|
)
|
|
|
|
return padded_batches
|
|
|
|
def _prepare(self, _random, doc):
|
|
mlm_tokens, mlm_mask, gpt_input_mask, gpt_loss_mask = self._mlm_cut(_random, doc)
|
|
full_tokens = self._gpt(doc)
|
|
return full_tokens, gpt_input_mask, mlm_tokens, mlm_mask, gpt_loss_mask
|
|
|
|
def _mlm_cut(self, _random, doc):
|
|
eod_index = self.dictionary.indices[EOL_SYMBOL]
|
|
|
|
if self.mlm_tokens_proportion == 0:
|
|
mlm_tokens = []
|
|
mlm_mask = []
|
|
gpt_input_mask = [0] * len(doc)
|
|
gpt_loss_mask = [1] * len(doc)
|
|
return mlm_tokens, mlm_mask, gpt_input_mask, gpt_loss_mask
|
|
|
|
cut_start = np.arange(1, len(doc)-3/2*self.mlm_cut_length, self.mlm_cut_length, dtype=int)
|
|
|
|
_random.shuffle(cut_start)
|
|
mlm_tokens = []
|
|
mlm_mask = []
|
|
start_list = []
|
|
gpt_input_mask = np.zeros(len(doc), dtype=int)
|
|
gpt_loss_mask = np.ones(len(doc), dtype=int)
|
|
mlm_tokens_total_num = (len(doc)-1) * self.mlm_tokens_proportion
|
|
|
|
mlm_tokens_cur_num = 0
|
|
|
|
for start in cut_start:
|
|
eod_num = doc[start:start+self.mlm_cut_length].count(eod_index)
|
|
if eod_num >= 2:
|
|
continue
|
|
elif eod_num == 1:
|
|
eod_pos = doc[start:start+self.mlm_cut_length].index(eod_index)
|
|
if self.mlm_cut_length - eod_pos < 20:
|
|
continue
|
|
start_ind, end_ind = start+eod_pos+1, start + self.mlm_cut_length
|
|
else:
|
|
cut_pos = _random.randint(0, self.mlm_cut_length-1)
|
|
if cut_pos >= self.mlm_cut_length/2:
|
|
start_ind, end_ind = start, start + cut_pos + 1
|
|
else:
|
|
start_ind, end_ind = start + cut_pos, start + self.mlm_cut_length
|
|
|
|
assert eod_index not in doc[start_ind:end_ind]
|
|
|
|
start_list.append(start)
|
|
mlm_tokens.append([self.dictionary.bos()] + doc[start_ind:end_ind])
|
|
mlm_tokens_cur_num += end_ind - start_ind
|
|
mlm_mask.append([0] + [1]*(end_ind - start_ind))
|
|
gpt_input_mask[start_ind:end_ind] = 1
|
|
gpt_loss_mask[start_ind:end_ind-1] = 0
|
|
|
|
if mlm_tokens_cur_num > mlm_tokens_total_num:
|
|
break
|
|
|
|
ind = np.array(start_list).argsort()
|
|
start_list = np.array(start_list)[ind]
|
|
mlm_tokens = np.array(mlm_tokens, dtype=object)[ind]
|
|
mlm_mask = np.array(mlm_mask, dtype=object)[ind]
|
|
|
|
return mlm_tokens, mlm_mask, gpt_input_mask, gpt_loss_mask
|
|
|
|
def _gpt(self, doc):
|
|
return doc
|
|
|
|
def _read_from_files(self, source_file):
|
|
data = []
|
|
file_path = os.path.join(self.data_dir, source_file)
|
|
|
|
if not os.path.exists(file_path):
|
|
print('| file {} not exists'.format(file_path), flush=True)
|
|
return iter([]) # skip bad file
|
|
|
|
with open(file_path, 'r', encoding='utf8') as f:
|
|
lines = f.read().strip().split('\n')
|
|
|
|
gpt_format_text = []
|
|
for line in lines:
|
|
gpt_format_text.extend(list(filter(None, json.loads(line)["text"].split("\n"))))
|
|
gpt_format_text.append('')
|
|
|
|
tokenized_lines = [self.tokenizer.encode(line) for line in gpt_format_text]
|
|
tokenized_ids = [self.dictionary.encode_line(line, add_if_not_exist=False) for line in tokenized_lines]
|
|
|
|
doc = [self.dictionary.bos()]
|
|
for ids in tokenized_ids:
|
|
if len(ids) > self.tokens_per_sample: # drop too long sentence
|
|
continue
|
|
|
|
if len(doc) + len(ids) > self.tokens_per_sample:
|
|
if len(doc) > 5/2*self.mlm_cut_length + 1:
|
|
data.append(doc)
|
|
doc = [self.dictionary.bos()]
|
|
doc.extend(ids)
|
|
|
|
if len(doc) > 1 and len(doc) <= self.tokens_per_sample:
|
|
if len(doc) > 5/2*self.mlm_cut_length + 1:
|
|
data.append(doc)
|
|
|
|
return data |