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