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)