def __getitem__(self, index): with data_utils.numpy_seed(self.seed, self.epoch, index): tokens = self.dataset[index] assert tokens[-1] == self.eos source, target = tokens, tokens.clone() # cannot permute code if self.permute_sentence_ratio > 0.0: source = self.permute_sentences(source, self.permute_sentence_ratio) if self.mask_ratio > 0: source = self.add_whole_word_mask(source, self.mask_ratio) if self.insert_ratio > 0: source = self.add_insertion_noise(source, self.insert_ratio) if self.rotate_ratio > 0.0 and np.random.random( ) < self.rotate_ratio: source = self.add_rolling_noise(source) assert (source >= 0).all() assert (source[1:-1] >= 1).all() assert (source <= len(self.vocab)).all() assert source[0] == self.vocab.bos() assert source[-1] == self.eos return { 'id': index, 'source': source, 'target': target, }
def load_dataset(self, split, epoch=1, combine=False, **kwargs): """Load a given dataset split. Args: split (str): name of the split (e.g., train, valid, test) """ paths = utils.split_paths(self.args['task']['data']) assert len(paths) > 0 data_path = paths[(epoch - 1) % len(paths)] split_path = os.path.join(data_path, split) dataset = data_utils.load_indexed_dataset( path=split_path, dictionary=self.source_dictionary, dataset_impl=self.args['dataset']['dataset_impl'], combine=combine, ) if dataset is None: raise FileNotFoundError('Dataset not found: {} ({})'.format( split, split_path)) # create continuous blocks of tokens dataset = TokenBlockDataset( dataset, dataset.sizes, self.args['task']['tokens_per_sample'] - 1, # one less for <s> pad=self.source_dictionary.pad(), eos=self.source_dictionary.eos(), break_mode=self.args['task']['sample_break_mode'], ) LOGGER.info('loaded {} blocks from: {}'.format(len(dataset), split_path)) # prepend beginning-of-sentence token (<s>, equiv. to [CLS] in BERT) dataset = PrependTokenDataset(dataset, self.source_dictionary.bos()) # create masked input and targets mask_whole_words = get_whole_word_mask(self.args, self.source_dictionary) \ if self.args['task']['mask_whole_words'] else None src_dataset, tgt_dataset = MaskTokensDataset.apply_mask( dataset, self.source_dictionary, pad_idx=self.source_dictionary.pad(), mask_idx=self.mask_idx, seed=self.args['common']['seed'], mask_prob=self.args['task']['mask_prob'], leave_unmasked_prob=self.args['task']['leave_unmasked_prob'], random_token_prob=self.args['task']['random_token_prob'], freq_weighted_replacement=self.args['task'] ['freq_weighted_replacement'], mask_whole_words=mask_whole_words, ) with data_utils.numpy_seed(self.args['common']['seed'] + epoch): shuffle = np.random.permutation(len(src_dataset)) self.datasets[split] = SortDataset( NestedDictionaryDataset( { 'id': IdDataset(), 'net_input': { 'src_tokens': PadDataset( src_dataset, pad_idx=self.source_dictionary.pad(), left_pad=False, ), 'src_lengths': NumelDataset(src_dataset, reduce=False), }, 'target': PadDataset( tgt_dataset, pad_idx=self.source_dictionary.pad(), left_pad=False, ), 'nsentences': NumSamplesDataset(), 'ntokens': NumelDataset(src_dataset, reduce=True), }, sizes=[src_dataset.sizes], ), sort_order=[ shuffle, src_dataset.sizes, ], )
def load_dataset(self, split, epoch=1, combine=False, **kwargs): """Load a given dataset split. Args: split (str): name of the split (e.g., train, valid, test) """ # paths = self.args.data.split(':') paths = utils.split_paths(self.args['task']['data']) assert len(paths) > 0 data_path = paths[(epoch - 1) % len(paths)] # split_path = os.path.join(data_path, split) if self.langs is None: languages = sorted([ name for name in os.listdir(data_path) if os.path.isdir(os.path.join(data_path, name)) ]) else: languages = self.langs # .split(',') # for name in languages: # assert os.path.exists(os.path.join(data_path, name)), FileNotFoundError(os.path.join(data_path, name)) LOGGER.info("| Training on {0} languages: {1}".format( len(languages), languages)) LOGGER.info("| Language to id mapping: ", {lang: id for id, lang in enumerate(languages)}) mask_whole_words = get_whole_word_mask(self.args, self.dictionary) lang_datasets = [] for language in languages: # split_path = os.path.join(data_path, language, split) if language == 'docstring': split_path = os.path.join(data_path, language, f"{split}.docstring.spm") else: split_path = os.path.join(data_path, language, f"{split}.code.spm") # split_path = os.path.join(data_path, language, f"{split}.spm.{language}") # dataset = data_utils.load_indexed_dataset( # split_path, # self.source_dictionary, # self.args['dataset']['dataset_impl'], # combine=combine, # ) dataset = load_lang_dataset_denoising( path=split_path, impl=self.args['dataset']['dataset_impl'], dict=self.source_dictionary) if dataset is None: raise FileNotFoundError('Dataset not found: {} ({})'.format( split, split_path)) dataset = AppendTokenDataset( TruncateDataset( StripTokenDataset(dataset, self.source_dictionary.eos()), self.args['task']['max_source_positions'] - 3), # <lang>, <bos>, <eos> token=self.source_dictionary.eos(), ) end_token = self.source_dictionary.index('[{}]'.format(language)) \ if self.args['task']['add_lang_token'] else self.source_dictionary.eos() # create continuous blocks of tokens dataset = TokenBlockDataset( dataset, dataset.sizes, self.args['task']['tokens_per_sample'] - 2, # one less for <s> and one for </s> pad=self.source_dictionary.pad(), eos=end_token, break_mode=self.args['task']['sample_break_mode'], document_sep_len=0, ) LOGGER.info('| loaded {} blocks from: {}'.format( len(dataset), split_path)) # prepend beginning-of-sentence token (<s>, equiv. to [CLS] in BERT) dataset = PrependTokenDataset(dataset, self.source_dictionary.bos()) dataset = AppendTokenDataset(dataset, end_token) lang_dataset = DenoisingDataset( dataset, dataset.sizes, self.dictionary, self.mask_idx, mask_whole_words, shuffle=self.args['dataset']['shuffle_instance'], seed=self.seed, args=self.args, eos=None if not self.args['task']['add_lang_token'] else self.source_dictionary.index('[{}]'.format(language)), ) lang_datasets.append(lang_dataset) dataset_lengths = np.array( [len(d) for d in lang_datasets], dtype=float, ) LOGGER.info('| loaded total {} blocks for all languages'.format( dataset_lengths.sum(), )) if split == self.args['dataset']['train_subset']: # For train subset, additionally up or down sample languages. sample_probs = self._get_sample_prob(dataset_lengths) LOGGER.info( "| Sample probability by language: ", { lang: "{0:.4f}".format(sample_probs[id]) for id, lang in enumerate(languages) }) size_ratio = (sample_probs * dataset_lengths.sum()) / dataset_lengths LOGGER.info( "| Up/Down Sampling ratio by language: ", { lang: "{0:.2f}".format(size_ratio[id]) for id, lang in enumerate(languages) }) resampled_lang_datasets = [ ResamplingDataset( lang_datasets[i], size_ratio=size_ratio[i], seed=self.args['common']['seed'], epoch=epoch, replace=size_ratio[i] >= 1.0, ) for i, d in enumerate(lang_datasets) ] dataset = ConcatDataset(resampled_lang_datasets, ) else: dataset = ConcatDataset(lang_datasets) lang_splits = [split] # for lang_id, lang_dataset in enumerate(lang_datasets): # split_name = split + '_' + languages[lang_id] # lang_splits.append(split_name) # self.datasets[split_name] = lang_dataset if split in self.args['dataset']['valid_subset']: self.args['dataset']['valid_subset'] = self.args['dataset'][ 'valid_subset'].replace(split, ','.join(lang_splits)) with data_utils.numpy_seed(self.args['common']['seed'] + epoch): shuffle = np.random.permutation(len(dataset)) self.datasets[split] = SortDataset( dataset, sort_order=[ shuffle, dataset.sizes, ], )
def load_masked_code_docstring_dataset_roberta(args, epoch, data_path, split, src, src_dict, tgt, tgt_dict, combine, dataset_impl, upsample_primary, left_pad_source, left_pad_target, max_source_positions, max_target_positions, prepend_bos=False, load_alignments=False, truncate_source=False, append_source_id=False): source_path = os.path.join(data_path, '{}.code'.format(split)) target_path = os.path.join(data_path, '{}.docstring'.format(split)) # source_dataset source_dataset = data_utils.load_indexed_dataset(source_path, 'text', src_dict, tokenizer=None, dataset_impl=dataset_impl) if source_dataset is None: raise FileNotFoundError('Dataset not found: {} ({})'.format( split, source_path)) # target_dataset target_dataset = data_utils.load_indexed_dataset(target_path, 'text', tgt_dict, tokenizer=None, dataset_impl=dataset_impl) if target_dataset is None: raise FileNotFoundError('Dataset not found: {} ({})'.format( split, target_path)) # concate dataset dataset = ConcatSentencesDataset([source_dataset, target_dataset]) # create continuous blocks of tokens dataset = TokenBlockDataset( dataset, dataset.sizes, args['task']['tokens_per_sample'] - 1, # one less for <s> pad=src_dict.pad(), eos=src_dict.eos(), break_mode=args['task']['sample_break_mode'], ) # LOGGER.info('loaded {} blocks from: {}'.format(len(dataset), split_path)) # # prepend beginning-of-sentence token (<s>, equiv. to [CLS] in BERT) dataset = PrependTokenDataset(dataset, src_dict.bos()) # .source_dictionary.bos() # # # create masked input and targets mask_whole_words = get_whole_word_mask(args, src_dict) \ if args['task']['mask_whole_words'] else None src_dataset, tgt_dataset = MaskTokensDataset.apply_mask( dataset, src_dict, pad_idx=src_dict.pad(), mask_idx=src_dict.index(constants.T_MASK), # self.mask_idx, seed=args['common']['seed'], mask_prob=args['task']['mask_prob'], leave_unmasked_prob=args['task']['leave_unmasked_prob'], random_token_prob=args['task']['random_token_prob'], freq_weighted_replacement=args['task']['freq_weighted_replacement'], mask_whole_words=mask_whole_words, ) with data_utils.numpy_seed(args['common']['seed'] + epoch): shuffle = np.random.permutation(len(src_dataset)) return SortDataset( NestedDictionaryDataset( { 'id': IdDataset(), 'net_input': { 'src_tokens': PadDataset( src_dataset, pad_idx=src_dict.pad(), left_pad=False, ), 'src_lengths': NumelDataset(src_dataset, reduce=False), }, 'target': PadDataset( tgt_dataset, pad_idx=src_dict.pad(), left_pad=False, ), 'nsentences': NumSamplesDataset(), 'ntokens': NumelDataset(src_dataset, reduce=True), }, sizes=[src_dataset.sizes], ), sort_order=[ shuffle, src_dataset.sizes, ], )
def shuffle_batches(batches, seed): with data_utils.numpy_seed(seed): np.random.shuffle(batches) return batches
def load_masked_traverse_dataset_roberta( args, epoch, data_path, split, source_dictionary, combine, ): split_path = os.path.join(data_path, '{}.ast_trav_df'.format(split)) dataset = data_utils.load_indexed_dataset( path=split_path, dictionary=source_dictionary, dataset_impl=args['dataset']['dataset_impl'], combine=combine, ) if dataset is None: raise FileNotFoundError('Dataset not found: {} ({})'.format( split, split_path)) # # create continuous blocks of tokens # dataset = TokenBlockDataset( # dataset, # dataset.sizes, # args['task']['tokens_per_sample'] - 1, # one less for <s> # pad=source_dictionary.pad(), # eos=source_dictionary.eos(), # break_mode=args['task']['sample_break_mode'], # ) # LOGGER.info('loaded {} blocks from: {}'.format(len(dataset), split_path)) # prepend beginning-of-sentence token (<s>, equiv. to [CLS] in BERT) dataset = PrependTokenDataset( dataset, source_dictionary.bos()) # .source_dictionary.bos() # create masked input and targets mask_whole_words = get_whole_word_mask(args, source_dictionary) \ if args['task']['mask_whole_words'] else None src_dataset, tgt_dataset = MaskTokensDataset.apply_mask( dataset, source_dictionary, pad_idx=source_dictionary.pad(), mask_idx=source_dictionary.index(constants.MASK), # self.mask_idx, seed=args['common']['seed'], mask_prob=args['task']['mask_prob'], leave_unmasked_prob=args['task']['leave_unmasked_prob'], random_token_prob=args['task']['random_token_prob'], freq_weighted_replacement=args['task']['freq_weighted_replacement'], mask_whole_words=mask_whole_words, ) with data_utils.numpy_seed(args['common']['seed'] + epoch): shuffle = np.random.permutation(len(src_dataset)) return SortDataset( NestedDictionaryDataset( { 'id': IdDataset(), 'net_input': { 'src_tokens': PadDataset( src_dataset, pad_idx=source_dictionary.pad(), left_pad=False, ), 'src_lengths': NumelDataset(src_dataset, reduce=False), }, 'target': PadDataset( tgt_dataset, pad_idx=source_dictionary.pad(), left_pad=False, ), 'nsentences': NumSamplesDataset(), 'ntokens': NumelDataset(src_dataset, reduce=True), }, sizes=[src_dataset.sizes], ), sort_order=[ shuffle, src_dataset.sizes, ], )
def get_batch_iterator( self, dataset, max_tokens=None, max_sentences=None, max_positions=None, ignore_invalid_inputs=False, required_batch_size_multiple=1, seed=1, num_shards=1, shard_id=0, num_workers=0, epoch=1, ): """ Get an iterator that yields batches of data from the given dataset. Args: dataset (~fairseq.data.NccDataset): dataset to batch max_tokens (int, optional): max number of tokens in each batch (default: None). max_sentences (int, optional): max number of sentences in each batch (default: None). max_positions (optional): max sentence length supported by the model (default: None). ignore_invalid_inputs (bool, optional): don't raise Exception for sentences that are too long (default: False). required_batch_size_multiple (int, optional): require batch size to be a multiple of N (default: 1). seed (int, optional): seed for random number generator for reproducibility (default: 1). num_shards (int, optional): shard the data iterator into N shards (default: 1). shard_id (int, optional): which shard of the data iterator to return (default: 0). num_workers (int, optional): how many subprocesses to use for data loading. 0 means the data will be loaded in the main process (default: 0). epoch (int, optional): the epoch to start the iterator from (default: 1). Returns: ~fairseq.iterators.EpochBatchIterator: a batched iterator over the given dataset split """ # For default fairseq task, return same iterator across epochs # as datasets are not dynamic, can be overridden in task specific # setting. if dataset in self.dataset_to_epoch_iter: return self.dataset_to_epoch_iter[dataset] assert isinstance(dataset, NccDataset) # initialize the dataset with the correct starting epoch dataset.set_epoch(epoch) # get indices ordered by example size with data_utils.numpy_seed(seed): indices = dataset.ordered_indices() # filter examples that are too large if max_positions is not None: indices = data_utils.filter_by_size( indices, dataset, max_positions, raise_exception=(not ignore_invalid_inputs), ) # create mini-batches with given size constraints batch_sampler = data_utils.batch_by_size( indices, dataset.num_tokens, max_tokens=max_tokens, max_sentences=max_sentences, required_batch_size_multiple=required_batch_size_multiple, ) # return a reusable, sharded iterator epoch_iter = iterators.EpochBatchIterator( dataset=dataset, collate_fn=dataset.collater, batch_sampler=batch_sampler, seed=seed, num_shards=num_shards, shard_id=shard_id, num_workers=num_workers, epoch=epoch, ) self.dataset_to_epoch_iter[dataset] = epoch_iter return epoch_iter
def __getitem__(self, index: int): with data_utils.numpy_seed(self.seed, self.epoch, index): item = self.dataset[index]['code_q'] sz = len(item) assert self.mask_idx not in item, \ 'Dataset contains mask_idx (={}), this is not expected!'.format( self.mask_idx, ) if self.mask_whole_words is not None: word_begins_mask = self.mask_whole_words.gather(0, item) word_begins_idx = word_begins_mask.nonzero().view(-1) sz = len(word_begins_idx) words = np.split(word_begins_mask, word_begins_idx)[1:] assert len(words) == sz word_lens = list(map(len, words)) # decide elements to mask mask = np.full(sz, False) num_mask = int( # add a random number for probabilistic rounding self.mask_prob * sz + np.random.rand()) mask[np.random.choice(sz, num_mask, replace=False)] = True if self.return_masked_tokens: # exit early if we're just returning the masked tokens # (i.e., the targets for masked LM training) if self.mask_whole_words is not None: mask = np.repeat(mask, word_lens) new_item = np.full(len(mask), self.pad_idx) new_item[mask] = item[torch.from_numpy(mask.astype(np.uint8)) == 1] return torch.from_numpy(new_item) # decide unmasking and random replacement rand_or_unmask_prob = self.random_token_prob + self.leave_unmasked_prob if rand_or_unmask_prob > 0.0: rand_or_unmask = mask & (np.random.rand(sz) < rand_or_unmask_prob) if self.random_token_prob == 0.0: unmask = rand_or_unmask rand_mask = None elif self.leave_unmasked_prob == 0.0: unmask = None rand_mask = rand_or_unmask else: unmask_prob = self.leave_unmasked_prob / rand_or_unmask_prob decision = np.random.rand(sz) < unmask_prob unmask = rand_or_unmask & decision rand_mask = rand_or_unmask & (~decision) else: unmask = rand_mask = None if unmask is not None: mask = mask ^ unmask if self.mask_whole_words is not None: mask = np.repeat(mask, word_lens) new_item = np.copy(item) new_item[mask] = self.mask_idx if rand_mask is not None: num_rand = rand_mask.sum() if num_rand > 0: if self.mask_whole_words is not None: rand_mask = np.repeat(rand_mask, word_lens) num_rand = rand_mask.sum() new_item[rand_mask] = np.random.choice( len(self.vocab), num_rand, p=self.weights, ) return torch.from_numpy(new_item)