def main(args):
tokenizer = build_tokenizer(args)
src_indices = get_indices(args.input_src, tokenizer)
trg_indices = get_indices(args.input_trg, tokenizer)
src_dataset = IndexDataset(src_indices)
trg_dataset = IndexDataset(trg_indices)
eos = tokenizer.sep_token_id
bos = tokenizer.cls_token_id
max_pos = args.max_pos
datasets = []
src_dataset = TruncateDataset(
StripTokenDataset(src_dataset, eos), max_pos - 2,)
trg_dataset = TruncateDataset(
StripTokenDataset(trg_dataset, eos), max_pos - 2,)
datasets.append(
TLMDataset(src_dataset, trg_dataset, bos, eos))
datasets.append(
TLMDataset(trg_dataset, src_dataset, bos, eos))
dataset = ConcatDataset(datasets)
print("| get all items ...")
items = [i for i in tqdm(dataset)]
print("| writing binary file ...")
prefix = os.path.join(args.output, "train.0")
save_items(items, prefix, len(tokenizer))
def build_dataset_for_inference(self, src_tokens, src_lengths, **kwargs):
"""
Generate batches for inference. We assume that the input begins with a
bos symbol (`<s>`) and ends with an eos symbol (`</s>`).
"""
pad = self.source_dictionary.pad()
eos = self.source_dictionary.eos()
src_dataset = TokenBlockDataset(
src_tokens,
src_lengths,
block_size=self.args.tokens_per_sample - 2, # for <s> and </s>
pad=pad,
eos=eos,
break_mode=self.args.sample_break_mode,
document_sep_len=0,
)
prev_output_tokens = PrependTokenDataset(
StripTokenDataset(src_dataset, eos), eos)
src_dataset = PadDataset(src_dataset, pad_idx=pad, left_pad=False)
return NestedDictionaryDataset(
{
"id": IdDataset(),
"net_input": {
"src_tokens":
src_dataset,
"src_lengths":
NumelDataset(src_dataset, reduce=False),
"prev_output_tokens":
PadDataset(prev_output_tokens, pad_idx=pad,
left_pad=False),
},
"target": src_dataset,
},
sizes=[np.array(src_lengths)],
)
def load_dataset(self,
split: str,
task_cfg: FairseqDataclass = None,
**kwargs):
data_path = self.cfg.data
task_cfg = task_cfg or self.cfg
has_unpaired_text = os.path.exists(
os.path.join(self.cfg.text_data, f"{split}.idx"))
self.datasets[split] = ExtractedFeaturesDataset(
path=data_path,
split=split,
min_length=3,
max_length=task_cfg.max_length,
labels=None if has_unpaired_text else task_cfg.labels,
label_dict=self.target_dictionary,
shuffle=getattr(task_cfg, "shuffle", True),
sort_by_length=task_cfg.sort_by_length,
)
logger.info(f"split {split} has unpaired text? {has_unpaired_text}")
if has_unpaired_text:
text_dataset = data_utils.load_indexed_dataset(
os.path.join(self.cfg.text_data, split),
self.target_dictionary)
text_dataset = StripTokenDataset(text_dataset,
self.target_dictionary.eos())
self.datasets[split] = RandomInputDataset(
self.datasets[split],
text_dataset,
["random_label"],
add_to_input=True,
pad_idx=self.target_dictionary.pad(),
)
def build_dataset_for_inference(self,
src_tokens,
src_lengths,
language="en_XX",
**kwargs):
"""
Generate batches for inference. We prepend an eos token to src_tokens
(or bos if `--add-bos-token` is set) and we append a <pad> to target.
This is convenient both for generation with a prefix and LM scoring.
"""
dataset = StripTokenDataset(
TokenBlockDataset(
src_tokens,
src_lengths,
block_size=None, # ignored for "eos" break mode
pad=self.source_dictionary.pad(),
eos=self.source_dictionary.eos(),
break_mode="eos",
),
# remove eos from (end of) target sequence
self.source_dictionary.eos(),
)
src_lang_idx = self.dictionary.index(lang_token(language))
src_dataset = PrependTokenDataset(
dataset,
token=((src_lang_idx or self.source_dictionary.bos()) if getattr(
self.args, "add_bos_token", False) else
self.source_dictionary.eos()),
)
max_seq_len = max(src_lengths) + 1
tgt_dataset = AppendTokenDataset(dataset,
token=self.source_dictionary.pad())
return NestedDictionaryDataset(
{
"id":
IdDataset(),
"net_input": {
"src_tokens":
PadDataset(
src_dataset,
pad_idx=self.source_dictionary.pad(),
left_pad=False,
pad_length=max_seq_len,
),
"src_lengths":
NumelDataset(src_dataset, reduce=False),
},
"target":
PadDataset(
tgt_dataset,
pad_idx=self.source_dictionary.pad(),
left_pad=False,
pad_length=max_seq_len,
),
},
sizes=[np.array(src_lengths)],
)
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.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(
split_path,
self.dictionary,
self.args.dataset_impl,
combine=combine,
)
if dataset is None:
raise FileNotFoundError("Dataset not found: {} ({})".format(
split, split_path))
dataset = StripTokenDataset(dataset, self.dictionary.eos())
# create continuous blocks of tokens
dataset = TokenBlockDataset(
dataset,
dataset.sizes,
self.args.tokens_per_sample -
2, # one less for <s> and one for </s>
pad=self.dictionary.pad(),
eos=self.dictionary.eos(),
break_mode=self.args.sample_break_mode,
document_sep_len=0,
)
# prepend beginning-of-sentence token (<s>, equiv. to [CLS] in BERT)
dataset = PrependTokenDataset(dataset, self.source_dictionary.bos())
dataset = AppendTokenDataset(dataset, self.source_dictionary.eos())
mask_whole_words = (get_whole_word_mask(self.args,
self.source_dictionary)
if self.args.mask_length != "subword" else None)
self.datasets[split] = DenoisingDataset(
dataset,
dataset.sizes,
self.dictionary,
self.mask_idx,
mask_whole_words,
shuffle=self.args.shuffle_instance,
seed=self.seed,
args=self.args,
)
logger.info(
"Split: {0}, Loaded {1} samples of denoising_dataset".format(
split,
len(self.datasets[split]),
))
def load_dataset(self, split, epoch=0, combine=False, **kwargs):
"""
TODO:
- break_mode="，。"
"""
paths = utils.split_paths(self.args.data)
assert len(paths) > 0
data_path = paths[epoch % len(paths)]
def get_path(type, split):
return os.path.join(data_path, type, split)
def make_dataset(type, dictionary):
split_path = get_path(type, split)
dataset = data_utils.load_indexed_dataset(
split_path,
dictionary,
self.args.dataset_impl,
combine=combine,
)
if dataset is None:
raise FileNotFoundError('Dataset not found: {} ({})'.format(
split, split_path))
return dataset
dataset = make_dataset('input', self.dictionary)
dataset = TruncateDataset(
RStripTokenDataset(dataset, self.dictionary.eos()),
self.args.tokens_per_sample - 2)
# prepend beginning-of-sentence token (<s>, equiv. to [CLS] in BERT)。
# https://github.com/pytorch/fairseq/blob/master/fairseq/tasks/translation.py#L71
# https://github.com/pytorch/fairseq/blob/77983ee1a52c4e011e54cc6bfa5352b7811ec96d/fairseq/tasks/denoising.py#L127
dataset = PrependTokenDataset(dataset, self.source_dictionary.bos())
dataset = AppendTokenDataset(dataset, self.source_dictionary.eos())
meta_dataset = make_dataset('meta', self.meta_dictionary)
meta_dataset = StripTokenDataset(
meta_dataset, id_to_strip=self.meta_dictionary.eos())
s2s_dataset = KnowledgeLanguagePairDataset.apply_mask(
dataset,
dataset.sizes,
self.source_dictionary,
meta=meta_dataset,
meta_sizes=meta_dataset.sizes,
meta_dict=self.meta_dictionary,
shuffle=True,
mask_prob=self.args.mask_prob,
leave_unmasked_prob=self.args.leave_unmasked_prob,
random_token_prob=self.args.random_token_prob,
sub_task=self.args.sub_task,
)
self.datasets[split] = s2s_dataset
def load_dataset(self, split, epoch=0, combine=False, **kwargs):
"""
TODO:
- break_mode="，。"
"""
paths = utils.split_paths(self.cfg.data)
assert len(paths) > 0
data_path = paths[epoch % len(paths)]
def get_path(type, split):
return os.path.join(data_path, type, split)
def make_dataset(type, dictionary):
split_path = get_path(type, split)
dataset = data_utils.load_indexed_dataset(
split_path,
dictionary,
self.cfg.dataset_impl,
combine=combine,
)
if dataset is None:
raise FileNotFoundError('Dataset not found: {} ({})'.format(
split, split_path))
return dataset
dataset = make_dataset('input', self.dictionary)
dataset = TruncateDataset(
RStripTokenDataset(dataset, self.dictionary.eos()),
self.cfg.tokens_per_sample - 2)
# prepend beginning-of-sentence token (<s>, equiv. to [CLS] in BERT)。
dataset = PrependTokenDataset(dataset, self.source_dictionary.bos())
dataset = AppendTokenDataset(dataset, self.source_dictionary.eos())
meta_dataset = make_dataset('meta', self.meta_dictionary)
meta_dataset = StripTokenDataset(
meta_dataset, id_to_strip=self.meta_dictionary.eos())
s2s_dataset = KnowledgeLanguagePairDataset.apply_mask(
dataset,
dataset.sizes,
self.source_dictionary,
meta=meta_dataset,
meta_sizes=meta_dataset.sizes,
meta_dict=self.meta_dictionary,
shuffle=True,
mask_idx=self.mask_idx,
mask_prob=self.cfg.mask_prob,
leave_unmasked_prob=self.cfg.leave_unmasked_prob,
random_token_prob=self.cfg.random_token_prob,
sub_task=self.cfg.sub_task,
)
self.datasets[split] = s2s_dataset
def _load_dataset_split(self, split, epoch, combine):
paths = utils.split_paths(self.cfg.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(
split_path,
self.dictionary,
self.cfg.dataset_impl,
combine=combine,
)
if dataset is None:
raise FileNotFoundError(
"Dataset not found: {} ({})".format(split, split_path)
)
dataset = StripTokenDataset(dataset, self.dictionary.eos())
dataset = maybe_shorten_dataset(
dataset,
split,
self.cfg.shorten_data_split_list,
self.cfg.shorten_method,
self.cfg.tokens_per_sample,
self.cfg.seed,
)
# create continuous blocks of tokens
dataset = TokenBlockDataset(
dataset,
dataset.sizes,
self.cfg.tokens_per_sample - 2,
# one less for <s> and one for </s>
pad=self.dictionary.pad(),
eos=self.dictionary.eos(),
break_mode=self.cfg.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, self.source_dictionary.eos())
return dataset
def build_dataset_for_inference(self, src_tokens, src_lengths, **kwargs):
dataset = StripTokenDataset(
TokenBlockDataset(
src_tokens,
src_lengths,
block_size=None, # ignored for "eos" break mode
pad=self.source_dictionary.pad(),
eos=self.source_dictionary.eos(),
break_mode="eos",
),
# remove eos from (end of) target sequence
self.source_dictionary.eos(),
)
src_dataset = PrependTokenDataset(
dataset,
token=(self.source_dictionary.bos() if getattr(
self.args, "add_bos_token", False) else
self.source_dictionary.eos()),
)
tgt_dataset = AppendTokenDataset(dataset,
token=self.source_dictionary.pad())
return 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),
},
sizes=[np.array(src_lengths)],
)
def load_langpair_dataset(
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,
num_buckets=0,
shuffle=True,
pad_to_multiple=1,
add_lang_token=False,
):
def split_exists(split, src, tgt, lang, data_path):
logger.info(
os.path.join(data_path,
"{}.{}-{}.{}".format(split, src, tgt, lang)))
filename = os.path.join(data_path,
"{}.{}-{}.{}".format(split, src, tgt, lang))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
def split_exists_self(split, src, data_path):
logger.info(
os.path.join(data_path, "{}.{}-{}.{}".format(split, src, src,
src)))
filename = os.path.join(data_path,
"{}.{}-{}.{}".format(split, src, src, src))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
def split_exists_valid(split, lang, data_path):
logger.info(os.path.join(data_path, "{}.{}".format(split, lang)))
filename = os.path.join(data_path, "{}.{}".format(split, lang))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
src_datasets = []
tgt_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else "")
# print(split_k, src, tgt, src, data_path)
prefix_src = None
prefix_tgt = None
if not "-" in split_k:
# infer langcode
if split_exists(split_k, src, tgt, src, data_path):
prefix = os.path.join(data_path,
"{}.{}-{}.".format(split_k, src, tgt))
elif split_exists(split_k, tgt, src, src, data_path):
prefix = os.path.join(data_path,
"{}.{}-{}.".format(split_k, tgt, src))
else:
if k > 0:
break
else:
raise FileNotFoundError(
"Dataset not found: {} ({}) {} {}".format(
split, data_path, src, tgt))
else:
# infer langcode
if split_exists_valid(split_k, src, data_path):
prefix = os.path.join(data_path, split_k + ".")
else:
if k > 0:
break
else:
raise FileNotFoundError(
"Dataset not found: {} ({}) ".format(split, data_path))
if prefix_src != None:
prefix = prefix_src
src_dataset = data_utils.load_indexed_dataset(prefix + src, src_dict,
dataset_impl)
if truncate_source:
src_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
)
src_datasets.append(src_dataset)
if prefix_tgt != None:
prefix = prefix_tgt
tgt_dataset = data_utils.load_indexed_dataset(prefix + tgt, tgt_dict,
dataset_impl)
if tgt_dataset is not None:
tgt_datasets.append(tgt_dataset)
logger.info("{} {} {}-{} {} examples".format(data_path, split_k,
src, tgt,
len(src_datasets[-1])))
if not combine:
break
assert len(src_datasets) == len(tgt_datasets) or len(tgt_datasets) == 0
if len(src_datasets) == 1:
src_dataset = src_datasets[0]
tgt_dataset = tgt_datasets[0] if len(tgt_datasets) > 0 else None
else:
sample_ratios = [1] * len(src_datasets)
sample_ratios[0] = upsample_primary
logger.info("::::data sample_ratios:{}".format(sample_ratios))
src_dataset = ConcatDataset(src_datasets, sample_ratios)
if len(tgt_datasets) > 0:
tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
else:
tgt_dataset = None
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(
tgt_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
if tgt_dataset is not None:
tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
eos = None
if append_source_id:
src_dataset = AppendTokenDataset(src_dataset,
src_dict.index("[{}]".format(src)))
if tgt_dataset is not None:
tgt_dataset = AppendTokenDataset(
tgt_dataset, tgt_dict.index("[{}]".format(tgt)))
eos = tgt_dict.index("[{}]".format(tgt))
eos = None
if add_lang_token:
src_dataset = PrependTokenDataset(src_dataset,
src_dict.index("[{}]".format(src)))
if tgt_dataset is not None:
tgt_dataset = PrependTokenDataset(
tgt_dataset, tgt_dict.index("[{}]".format(tgt)))
align_dataset = None
if load_alignments:
align_path = os.path.join(data_path,
"{}.align.{}-{}".format(split, src, tgt))
if indexed_dataset.dataset_exists(align_path, impl=dataset_impl):
align_dataset = data_utils.load_indexed_dataset(
align_path, None, dataset_impl)
tgt_dataset_sizes = tgt_dataset.sizes if tgt_dataset is not None else None
return LanguagePairDataset(
src_dataset,
src_dataset.sizes,
src_dict,
tgt_dataset,
tgt_dataset_sizes,
tgt_dict,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
align_dataset=align_dataset,
eos=eos,
num_buckets=num_buckets,
shuffle=shuffle,
pad_to_multiple=pad_to_multiple,
)
def load_langpair_dataset(
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,
num_buckets=0,
shuffle=True,
):
def split_exists(split, src, tgt, lang, data_path):
filename = os.path.join(data_path,
'{}.{}-{}.{}'.format(split, src, tgt, lang))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
src_datasets = []
tgt_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else '')
# infer langcode
if split_exists(split_k, src, tgt, src, data_path):
prefix = os.path.join(data_path,
'{}.{}-{}.'.format(split_k, src, tgt))
elif split_exists(split_k, tgt, src, src, data_path):
prefix = os.path.join(data_path,
'{}.{}-{}.'.format(split_k, tgt, src))
else:
if k > 0:
break
else:
raise FileNotFoundError('Dataset not found: {} ({})'.format(
split, data_path))
src_dataset = data_utils.load_indexed_dataset(prefix + src, src_dict,
dataset_impl)
if truncate_source:
src_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
)
src_datasets.append(src_dataset)
tgt_dataset = data_utils.load_indexed_dataset(prefix + tgt, tgt_dict,
dataset_impl)
if tgt_dataset is not None:
tgt_datasets.append(tgt_dataset)
logger.info('{} {} {}-{} {} examples'.format(data_path, split_k,
src, tgt,
len(src_datasets[-1])))
if not combine:
break
assert len(src_datasets) == len(tgt_datasets) or len(tgt_datasets) == 0
if len(src_datasets) == 1:
src_dataset = src_datasets[0]
tgt_dataset = tgt_datasets[0] if len(tgt_datasets) > 0 else None
# these features are not yet implemented for the cluster code
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(
tgt_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
if tgt_dataset is not None:
tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
eos = None
if append_source_id:
src_dataset = AppendTokenDataset(
src_dataset, src_dict.index('[{}]'.format(src)))
if tgt_dataset is not None:
tgt_dataset = AppendTokenDataset(
tgt_dataset, tgt_dict.index('[{}]'.format(tgt)))
eos = tgt_dict.index('[{}]'.format(tgt))
align_dataset = None
if load_alignments:
align_path = os.path.join(data_path,
'{}.align.{}-{}'.format(split, src, tgt))
if indexed_dataset.dataset_exists(align_path, impl=dataset_impl):
align_dataset = data_utils.load_indexed_dataset(
align_path, None, dataset_impl)
tgt_dataset_sizes = tgt_dataset.sizes if tgt_dataset is not None else None
return LanguagePairDataset(
src_dataset,
src_dataset.sizes,
src_dict,
tgt_dataset,
tgt_dataset_sizes,
tgt_dict,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
align_dataset=align_dataset,
eos=eos,
num_buckets=num_buckets,
shuffle=shuffle,
)
else:
# sample_ratios = [1] * len(src_datasets)
# sample_ratios[0] = upsample_primary
# src_dataset = ConcatDataset(src_datasets, sample_ratios)
# if len(tgt_datasets) > 0:
# tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
# else:
# tgt_dataset = None
datasets = []
eos = None
align_dataset = None
for i in range(0, len(src_datasets)):
src_dataset = src_datasets[i]
tgt_dataset = tgt_datasets[i]
tgt_dataset_sizes = tgt_dataset.sizes if tgt_dataset is not None else None
datasets.append(
LanguagePairDataset(
src_dataset,
src_dataset.sizes,
src_dict,
tgt_dataset,
tgt_dataset_sizes,
tgt_dict,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
align_dataset=align_dataset,
eos=eos,
num_buckets=num_buckets,
shuffle=shuffle,
))
return datasets
shorten_data_split_list = ''
paths = utils.split_paths(bin_dir)
assert len(paths) > 0
split = 'valid'
data_path = paths[0]
split_path = os.path.join(data_path, split)
dataset = data_utils.load_indexed_dataset(
split_path,
dictionary,
dataset_impl,
combine=False,
)
dataset = StripTokenDataset(dataset, dictionary.eos())
dataset = maybe_shorten_dataset(
dataset,
split,
shorten_data_split_list,
shorten_method,
tokens_per_sample,
seed,
)
prev_size = len(dataset)
# create continuous blocks of tokens
dataset = TokenBlockDataset(
dataset,
def load_dataset(self, split, combine=False, **kwargs):
"""Load a given dataset split (e.g., train, valid, test)."""
def get_path(type, split):
return os.path.join(self.args.data, type, split)
def make_dataset(type, dictionary):
split_path = get_path(type, split)
dataset = data_utils.load_indexed_dataset(
split_path,
self.source_dictionary,
self.args.dataset_impl,
combine=combine,
)
return dataset
input0 = make_dataset('input0', self.source_dictionary)
assert input0 is not None, 'could not find dataset: {}'.format(
get_path(type, split))
input1 = make_dataset('input1', self.source_dictionary)
if self.args.init_token is not None:
input0 = PrependTokenDataset(input0, self.args.init_token)
if input1 is None:
src_tokens = input0
else:
if self.args.separator_token is not None:
input1 = PrependTokenDataset(input1, self.args.separator_token)
src_tokens = ConcatSentencesDataset(input0, input1)
with data_utils.numpy_seed(self.args.seed):
shuffle = np.random.permutation(len(src_tokens))
if self.args.truncate_sequence:
src_tokens = TruncateDataset(src_tokens, self.args.max_positions)
dataset = {
'id': IdDataset(),
'net_input': {
'src_tokens':
RightPadDataset(
src_tokens,
pad_idx=self.source_dictionary.pad(),
),
'src_lengths':
NumelDataset(src_tokens, reduce=False),
},
'nsentences': NumSamplesDataset(),
'ntokens': NumelDataset(src_tokens, reduce=True),
}
if not self.args.regression_target:
label_dataset = make_dataset('label', self.target_dictionary)
if label_dataset is not None:
dataset.update(target=OffsetTokensDataset(
StripTokenDataset(
label_dataset,
id_to_strip=self.target_dictionary.eos(),
),
offset=-self.target_dictionary.nspecial,
))
else:
label_path = f"{get_path('label', split)}.label"
if os.path.exists(label_path):
dataset.update(target=RawLabelDataset(
[float(x.strip()) for x in open(label_path).readlines()]))
nested_dataset = NestedDictionaryDataset(
dataset,
sizes=[src_tokens.sizes],
)
if self.args.no_shuffle:
dataset = nested_dataset
else:
dataset = SortDataset(
nested_dataset,
# shuffle
sort_order=[shuffle],
)
print(f"| Loaded {split} with #samples: {len(dataset)}")
self.datasets[split] = dataset
return self.datasets[split]
def load_dataset(self, split, combine=False, **kwargs):
"""Load a given dataset split (e.g., train, valid, test)."""
def get_path(type, split):
return os.path.join(self.args.data, type, split)
def make_dataset(type, dictionary):
split_path = get_path(type, split)
dataset = data_utils.load_indexed_dataset(
split_path,
dictionary,
self.args.dataset_impl,
combine=combine,
)
return dataset
src_tokens = make_dataset('data', self.source_dictionary)
if self.args.init_token is not None:
src_tokens = PrependTokenDataset(src_tokens, self.args.init_token)
with data_utils.numpy_seed(self.args.seed):
shuffle = np.random.permutation(len(src_tokens))
if self.args.truncate_sequence:
src_tokens = TruncateDataset(src_tokens, self.args.max_positions)
dataset = {
'id': IdDataset(),
'net_input': {
'src_tokens':
RightPadDataset(
src_tokens,
pad_idx=self.source_dictionary.pad(),
),
'src_lengths':
NumelDataset(src_tokens, reduce=False),
},
'nsentences': NumSamplesDataset(),
'ntokens': NumelDataset(src_tokens, reduce=True),
}
label_dataset = make_dataset('label', self.label_dictionary)
if label_dataset is not None:
dataset.update(target=OffsetTokensDataset(
StripTokenDataset(
label_dataset,
id_to_strip=self.label_dictionary.eos(),
),
offset=-self.label_dictionary.nspecial,
))
nested_dataset = NestedDictionaryDataset(
dataset,
sizes=[src_tokens.sizes],
)
if self.args.no_shuffle:
dataset = nested_dataset
else:
dataset = SortDataset(
nested_dataset,
# shuffle
sort_order=[shuffle],
)
logger.info("Loaded {0} with #samples: {1}".format(
split, len(dataset)))
self.datasets[split] = dataset
return self.datasets[split]
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.cfg.data)
assert len(paths) > 0
if split != getattr(self.cfg, "train_subset", None):
# if not training data set, use the first shard for valid and test
paths = paths[:1]
data_path = paths[(epoch - 1) % len(paths)]
# infer langcode
src, tgt = self.cfg.source_lang, self.cfg.target_lang
prefix = os.path.join(data_path, '{}.{}-{}.'.format(split, src, tgt))
src_dataset = data_utils.load_indexed_dataset(prefix + src,
self.src_dict,
self.cfg.dataset_impl)
tag_dataset = data_utils.load_indexed_dataset(prefix + tgt,
self.tag_dict,
self.cfg.dataset_impl)
src_dataset = StripTokenDataset(
src_dataset, id_to_strip=self.source_dictionary.eos())
tag_dataset = StripTokenDataset(tag_dataset,
id_to_strip=self.tag_dictionary.eos())
tag_pad = self.source_dictionary.pad()
tag_offset = tag_pad + 1
dataset = {
'id':
IdDataset(),
'net_input': {
'src_tokens':
RightPadDataset(src_dataset,
pad_idx=self.source_dictionary.pad()),
'src_lengths':
NumelDataset(src_dataset, reduce=False),
},
'nsentences':
NumSamplesDataset(),
'ntokens':
NumelDataset(src_dataset, reduce=True),
'target':
RightPadDataset(
OffsetTokensDataset(tag_dataset,
offset=-self.tag_dictionary.nspecial +
tag_offset),
pad_idx=tag_pad,
),
}
dataset = NestedDictionaryDataset(
dataset,
sizes=[src_dataset.sizes],
)
logger.info(
str([self.src_dict[k]
for k in dataset[0]['net_input.src_tokens']]))
logger.info(
str([
self.tag_dict[k + self.tag_dictionary.nspecial - tag_offset]
for k in dataset[0]['target']
]))
self.datasets[split] = dataset
def load_ape_dataset(
data_path,
split,
src_dict,
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,
num_buckets=0,
input_type='src_only',
src_type="src",
):
"""
ignoring src and tgt name. Assume $split.src, $split.mt, and $split.pe exist
"""
src = src_type
mt = "mt"
tgt = "pe"
term = "term"
src_factor = src_type + "_embed"
mt_factor = "mt_embed"
def split_exists(split, lang, data_path):
filename = os.path.join(data_path, '{}.{}'.format(split, lang))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
def load_dataset(lang,
lang_dict,
prefix,
dataset_length,
sample_ratios=None):
"""
Function to load additional dataset and deal with all parameters.
Easier than copying redudant code for each dataset.
Requires src_dataset to provide the length and sample_ratios.
"""
lang_datasets = []
lang_dataset = data_utils.load_indexed_dataset(prefix + lang,
lang_dict, dataset_impl)
if lang_dataset is not None:
lang_datasets.append(lang_dataset)
assert dataset_length == len(lang_datasets) or len(lang_datasets) == 0
if dataset_length == 1:
lang_dataset = lang_datasets[0] if len(lang_datasets) > 0 else None
else:
assert sample_ratios is not None
if len(lang_datasets) > 0:
lang_dataset = ConcatDataset(lang_datasets, sample_ratios)
else:
lang_dataset = None
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(
lang_dict, "bos_index")
if lang_dataset is not None:
lang_dataset = PrependTokenDataset(lang_dataset,
lang_dict.bos())
eos = None
if append_source_id:
if lang_dataset is not None:
lang_dataset = AppendTokenDataset(
lang_dataset, lang_dict.index('[{}]'.format(lang)))
lang_dataset_sizes = lang_dataset.sizes if lang_dataset is not None else None
return lang_dataset, lang_dataset_sizes
src_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else '')
# infer langcode
if split_exists(split_k, src, data_path):
prefix = os.path.join(data_path, '{}.'.format(split_k))
elif split_exists(split_k, mt, data_path):
prefix = os.path.join(data_path, '{}.'.format(split_k))
elif split_exists(split_k, tgt, data_path):
prefix = os.path.join(data_path, '{}.'.format(split_k))
elif split_exists(split_k, term, data_path):
prefix = os.path.join(data_path, '{}.'.format(split_k))
elif split_exists(split_k, src_factor, data_path):
prefix = os.path.join(data_path, '{}.'.format(split_k))
else:
if k > 0:
break
else:
raise FileNotFoundError('Dataset not found: {} ({})'.format(
split, data_path))
src_dataset = data_utils.load_indexed_dataset(prefix + src, src_dict,
dataset_impl)
if truncate_source:
src_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
)
src_datasets.append(src_dataset)
if not combine:
break
dataset_length = len(src_datasets)
sample_ratios = None
if len(src_datasets) == 1:
src_dataset = src_datasets[0]
else:
sample_ratios = [1] * len(src_datasets)
sample_ratios[0] = upsample_primary
src_dataset = ConcatDataset(src_datasets, sample_ratios)
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(
tgt_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
eos = None
if append_source_id:
src_dataset = AppendTokenDataset(src_dataset,
src_dict.index('[{}]'.format(src)))
eos = tgt_dict.index('[{}]'.format(tgt))
align_dataset = None
mt_dataset, mt_dataset_sizes = load_dataset(mt,
tgt_dict,
prefix,
dataset_length,
sample_ratios=sample_ratios)
tgt_dataset, tgt_dataset_sizes = load_dataset(tgt,
tgt_dict,
prefix,
dataset_length,
sample_ratios=sample_ratios)
term_dataset, term_dataset_sizes = load_dataset(
term, tgt_dict, prefix, dataset_length, sample_ratios=sample_ratios)
src_factor_dataset, src_factor_dataset_sizes = load_dataset(
src_factor,
tgt_dict,
prefix,
dataset_length,
sample_ratios=sample_ratios)
mt_factor_dataset, mt_factor_dataset_sizes = load_dataset(
mt_factor,
tgt_dict,
prefix,
dataset_length,
sample_ratios=sample_ratios)
logger.info('{} {} {} examples'.format(data_path, split_k,
len(src_datasets[-1])))
return APEDataset(src_dataset,
src_dataset.sizes,
src_dict,
tgt_dataset,
tgt_dataset_sizes,
tgt_dict,
mt_dataset,
mt_dataset_sizes,
term_dataset,
term_dataset_sizes,
src_factor_dataset,
src_factor_dataset_sizes,
mt_factor_dataset,
mt_factor_dataset_sizes,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
align_dataset=align_dataset,
eos=eos,
num_buckets=num_buckets,
input_type=input_type)
def load_langpair_dataset(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,
explicit_str_att=False):
def split_exists(split, src, tgt, lang, data_path):
filename = os.path.join(data_path,
'{}.{}-{}.{}'.format(split, src, tgt, lang))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
src_datasets = []
tgt_datasets = []
sent_id_datasets = []
chains_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else '')
# infer langcode
if split_exists(split_k, src, tgt, src, data_path):
prefix = os.path.join(data_path,
'{}.{}-{}.'.format(split_k, src, tgt))
elif split_exists(split_k, tgt, src, src, data_path):
prefix = os.path.join(data_path,
'{}.{}-{}.'.format(split_k, tgt, src))
else:
if k > 0:
break
else:
raise FileNotFoundError('Dataset not found: {} ({})'.format(
split, data_path))
pre_src_dataset = data_utils.load_indexed_dataset(
prefix + src, src_dict, dataset_impl)
if truncate_source:
src_dataset = AppendTokenDataset(TruncateDataset(
StripTokenDataset(pre_src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
split=split)
src_datasets.append(src_dataset)
else:
src_datasets.append(pre_src_dataset)
if split_exists(split_k, src, tgt, src, data_path):
prefix = os.path.join(data_path,
'{}.{}-{}.'.format(split_k, src, tgt))
elif split_exists(split_k, tgt, src, src, data_path):
prefix = os.path.join(data_path,
'{}.{}-{}.'.format(split_k, tgt, src))
else:
if k > 0:
break
else:
raise FileNotFoundError('Dataset not found: {} ({})'.format(
split, data_path))
sent_id_dataset = SentIdsRawDataset(prefix + 'source.sentids')
if truncate_source:
sent_id_dataset = AppendLastTokenDataset(TruncateNDimDataset(
StripTokenFromMaskDataset(sent_id_dataset, pre_src_dataset,
src_dict.eos()),
max_source_positions - 1,
dim=1),
split=split)
sent_id_datasets.append(sent_id_dataset)
if explicit_str_att:
chains_dataset = ChainsDataset(prefix + 'source.chains')
chains_datasets.append(chains_dataset)
tgt_dataset = data_utils.load_indexed_dataset(prefix + tgt, tgt_dict,
dataset_impl)
if tgt_dataset is not None:
tgt_datasets.append(tgt_dataset)
logger.info('{} {} {}-{} {} examples'.format(data_path, split_k,
src, tgt,
len(src_datasets[-1])))
if not combine:
break
assert (len(src_datasets) == len(tgt_datasets) and len(src_datasets)
== len(sent_id_datasets)) or len(tgt_datasets) == 0
if len(src_datasets) == 1:
src_dataset = src_datasets[0]
sent_id_dataset = sent_id_datasets[0]
chains_dataset = chains_datasets
if explicit_str_att:
chains_dataset = chains_datasets[0]
tgt_dataset = tgt_datasets[0] if len(tgt_datasets) > 0 else None
else:
sample_ratios = [1] * len(src_datasets)
sample_ratios[0] = upsample_primary
# src_dataset = ConcatDataset(src_datasets, sample_ratios)
# if len(tgt_datasets) > 0:
# tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
# else:
# tgt_dataset = None
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(
tgt_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
sent_id_dataset = PrependFirstTokenDataset(sent_id_dataset)
if tgt_dataset is not None:
tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
eos = None
if append_source_id:
src_dataset = AppendTokenDataset(src_dataset,
src_dict.index('[{}]'.format(src)),
split=split)
sent_id_dataset = AppendLastTokenDataset(sent_id_dataset, split=split)
if tgt_dataset is not None:
tgt_dataset = AppendTokenDataset(
tgt_dataset, tgt_dict.index('[{}]'.format(tgt)))
eos = tgt_dict.index('[{}]'.format(tgt))
align_dataset = None
# if load_alignments:
# align_path = os.path.join(data_path, '{}.align.{}-{}'.format(split, src, tgt))
# chains = torch.load(load_alignments)
# if indexed_dataset.dataset_exists(align_path, impl=dataset_impl):
# align_dataset = data_utils.load_indexed_dataset(align_path, None, dataset_impl)
tgt_dataset_sizes = tgt_dataset.sizes if tgt_dataset is not None else None
return StructSumDataset(src_dataset,
src_dataset.sizes,
src_dict,
tgt_dataset,
tgt_dataset_sizes,
tgt_dict,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
max_source_positions=max_source_positions,
max_target_positions=max_target_positions,
align_dataset=align_dataset,
eos=eos,
src_sent_ids=sent_id_dataset,
split=split,
chains_dataset=chains_dataset,
explicit_str_att=explicit_str_att)
def load_dataset(self, split, combine=False, **kwargs):
"""Load a given dataset split (e.g., train, valid, test)."""
def get_path(key, split):
return os.path.join(self.cfg.data, key, split)
def make_dataset(key, dictionary):
split_path = get_path(key, split)
try:
dataset = data_utils.load_indexed_dataset(
split_path,
dictionary,
combine=combine,
)
except Exception as e:
if "StorageException: [404] Path not found" in str(e):
logger.warning(f"dataset {e} not found")
dataset = None
else:
raise e
return dataset
input0 = make_dataset("input0", self.source_dictionary)
assert input0 is not None, "could not find dataset: {}".format(
get_path("input0", split))
input1 = make_dataset("input1", self.source_dictionary)
if self.cfg.init_token is not None:
input0 = PrependTokenDataset(input0, self.cfg.init_token)
if input1 is None:
src_tokens = input0
else:
if self.cfg.separator_token is not None:
input1 = PrependTokenDataset(input1, self.cfg.separator_token)
src_tokens = ConcatSentencesDataset(input0, input1)
with data_utils.numpy_seed(self.cfg.seed):
shuffle = np.random.permutation(len(src_tokens))
src_tokens = maybe_shorten_dataset(
src_tokens,
split,
self.cfg.shorten_data_split_list,
self.cfg.shorten_method,
self.max_positions(),
self.cfg.seed,
)
dataset = {
"id": IdDataset(),
"net_input": {
"src_tokens":
RightPadDataset(
src_tokens,
pad_idx=self.source_dictionary.pad(),
),
"src_lengths":
NumelDataset(src_tokens, reduce=False),
},
"nsentences": NumSamplesDataset(),
"ntokens": NumelDataset(src_tokens, reduce=True),
}
if self.cfg.add_prev_output_tokens:
prev_tokens_dataset = RightPadDataset(
RollDataset(src_tokens, 1),
pad_idx=self.dictionary.pad(),
)
dataset["net_input"].update(
prev_output_tokens=prev_tokens_dataset, )
if not self.cfg.regression_target:
label_dataset = make_dataset("label", self.label_dictionary)
if label_dataset is not None:
dataset.update(target=OffsetTokensDataset(
StripTokenDataset(
label_dataset,
id_to_strip=self.label_dictionary.eos(),
),
offset=-self.label_dictionary.nspecial,
))
else:
label_path = "{0}.label".format(get_path("label", split))
if os.path.exists(label_path):
def parse_regression_target(i, line):
values = line.split()
assert (
len(values) == self.cfg.num_classes
), f'expected num_classes={self.cfg.num_classes} regression target values on line {i}, found: "{line}"'
return [float(x) for x in values]
with open(label_path) as h:
dataset.update(target=RawLabelDataset([
parse_regression_target(i, line.strip())
for i, line in enumerate(h.readlines())
]))
nested_dataset = NestedDictionaryDataset(
dataset,
sizes=[src_tokens.sizes],
)
if self.cfg.no_shuffle:
dataset = nested_dataset
else:
dataset = SortDataset(
nested_dataset,
# shuffle
sort_order=[shuffle],
)
logger.info("Loaded {0} with #samples: {1}".format(
split, len(dataset)))
self.datasets[split] = dataset
return self.datasets[split]
def load_pos_langpair_dataset(
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
):
# Check the existence of the file
def split_exists(split, src, tgt, lang, data_path):
filename = os.path.join(
data_path, '{}.{}-{}.{}'.format(split, src, tgt, lang))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
src_datasets = []
tgt_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else '')
# infer langcode (from a->b or from b->a)
if split_exists(split_k, src, tgt, src, data_path):
prefix = os.path.join(
data_path, '{}.{}-{}.'.format(split_k, src, tgt))
elif split_exists(split_k, tgt, src, src, data_path):
prefix = os.path.join(
data_path, '{}.{}-{}.'.format(split_k, tgt, src))
else:
if k > 0:
break
else:
raise FileNotFoundError(
'Dataset not found: {} ({})'.format(split, data_path))
src_dataset = data_utils.load_indexed_dataset(
prefix + src, src_dict, dataset_impl)
if truncate_source:
src_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
)
src_datasets.append(src_dataset)
tgt_dataset = data_utils.load_indexed_dataset(
prefix + tgt, tgt_dict, dataset_impl)
if tgt_dataset is not None:
tgt_datasets.append(tgt_dataset)
logger.info('{} {} {}-{} {} examples'.format(
data_path, split_k, src, tgt, len(src_datasets[-1])
))
if not combine:
break
assert len(src_datasets) == len(tgt_datasets) or len(tgt_datasets) == 0
if len(src_datasets) == 1:
src_dataset = src_datasets[0]
tgt_dataset = tgt_datasets[0] if len(tgt_datasets) > 0 else None
else:
sample_ratios = [1] * len(src_datasets)
sample_ratios[0] = upsample_primary
src_dataset = ConcatDataset(src_datasets, sample_ratios)
if len(tgt_datasets) > 0:
tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
else:
tgt_dataset = None
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(
tgt_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
if tgt_dataset is not None:
tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
eos = None
if append_source_id:
src_dataset = AppendTokenDataset(
src_dataset, src_dict.index('[{}]'.format(src)))
if tgt_dataset is not None:
tgt_dataset = AppendTokenDataset(
tgt_dataset, tgt_dict.index('[{}]'.format(tgt)))
eos = tgt_dict.index('[{}]'.format(tgt))
align_dataset = None
if load_alignments:
align_path = os.path.join(
data_path, '{}.align.{}-{}'.format(split, src, tgt))
if indexed_dataset.dataset_exists(align_path, impl=dataset_impl):
align_dataset = data_utils.load_indexed_dataset(
align_path, None, dataset_impl)
tgt_dataset_sizes = tgt_dataset.sizes if tgt_dataset is not None else None
# Load POS Graph
def graph_exist(data_path, split, src, tgt, lang):
existence = True
row_path = os.path.join(data_path, '{}.{}-{}.{}'.format(split, src, tgt, src)) + '.row'
col_path = os.path.join(data_path, '{}.{}-{}.{}'.format(split, src, tgt, src)) + '.col'
anchor_path = os.path.join(data_path, '{}.{}-{}.{}'.format(split, src, tgt, src)) + '.anchor'
if(not os.path.exists(row_path)):
existence = False
elif(not os.path.exists(col_path)):
existence = False
elif(not os.path.exists(anchor_path)):
existence = False
return existence
pos_graphs_l = []
pos_anchors_l = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else '')
existence = graph_exist(data_path, split_k, src, tgt, src)
if(not existence):
if(k == 0):
raise FileNotFoundError('POS Graph Dataset not found')
if(k > 0):
break
pos_rows = codecs.open(os.path.join(
data_path, '{}.{}-{}.{}'.format(split_k, src, tgt, src)) + '.row', 'r', 'utf-8').readlines()
pos_cols = codecs.open(os.path.join(
data_path, '{}.{}-{}.{}'.format(split_k, src, tgt, src)) + '.col', 'r', 'utf-8').readlines()
pos_graphs = []
print('Loading graphs' + '.' * 50)
assert len(pos_cols) == len(pos_rows)
pbar = tqdm(total=len(pos_cols))
for n, (row, col) in enumerate(zip(pos_rows, pos_cols)):
pos_row = [eval(i) for i in row.strip().split()]
pos_col = [eval(i) for i in col.strip().split()]
pos_graphs.append((pos_row, pos_col))
pbar.update()
pbar.close()
pos_anchors = codecs.open(os.path.join(
data_path, '{}.{}-{}.{}'.format(split_k, src, tgt, src)) + '.anchor', 'r', 'utf-8').readlines()
anchors = []
for line in pos_anchors:
anchors.append([eval(i) for i in line.strip().split()])
pos_graphs_l.extend(pos_graphs)
pos_anchors_l.extend(anchors)
assert (len(pos_anchors_l) == len(pos_graphs_l)) and (len(src_dataset.sizes) == len(pos_anchors_l))
return POSGraphLanguagePairDataset(
src_dataset, src_dataset.sizes, src_dict, pos_anchors_l, pos_graphs_l,
tgt_dataset, tgt_dataset_sizes, tgt_dict,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
max_source_positions=max_source_positions,
max_target_positions=max_target_positions,
align_dataset=align_dataset, eos=eos
)
def load_generation_pair_dataset(
data_path, split,
tgt,
src_dict,
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, common_eos=None
):
def split_exists(split, src, tgt, lang, data_path):
filename = os.path.join(data_path, '{}.{}-{}.{}'.format(split, src, tgt, lang))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
src_datasets = []
tgt_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else '')
# infer langcode
if split_exists(split_k, "src", "tgt", "src", data_path):
prefix = os.path.join(data_path, '{}.{}-{}.'.format(split_k, "src", "tgt"))
elif split_exists(split_k, "tgt", "src", "src", data_path):
prefix = os.path.join(data_path, '{}.{}-{}.'.format(split_k, "tgt", "src"))
else:
if k > 0:
break
else:
raise FileNotFoundError('Dataset not found: {} ({})'.format(split, data_path))
src_dataset = data_utils.load_indexed_dataset(prefix + "src", src_dict, dataset_impl)
if truncate_source:
src_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
)
src_datasets.append(src_dataset)
tgt_dataset = data_utils.load_indexed_dataset(prefix + "tgt", tgt_dict, dataset_impl)
if tgt_dataset is not None:
tgt_datasets.append(tgt_dataset)
logger.info('{} {} {}-{} {} examples'.format(
data_path, split_k, "src", "tgt", len(src_datasets[-1])
))
if not combine:
break
assert len(src_datasets) == len(tgt_datasets) or len(tgt_datasets) == 0
if len(src_datasets) == 1:
src_dataset = src_datasets[0]
tgt_dataset = tgt_datasets[0] if len(tgt_datasets) > 0 else None
else:
sample_ratios = [1] * len(src_datasets)
sample_ratios[0] = upsample_primary
src_dataset = ConcatDataset(src_datasets, sample_ratios)
if len(tgt_datasets) > 0:
tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
else:
tgt_dataset = None
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(tgt_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
if tgt_dataset is not None:
tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
eos = None
if append_source_id:
if common_eos is not None:
src_dataset = AppendTokenDataset(src_dataset, src_dict.index('[{}]'.format(common_eos)))
if tgt_dataset is not None:
tgt_dataset = AppendTokenDataset(tgt_dataset, tgt_dict.index('[{}]'.format(common_eos)))
eos = tgt_dict.index('[{}]'.format(common_eos))
bos = tgt_dict.index('[{}]'.format(tgt))
align_dataset = None
if load_alignments:
align_path = os.path.join(data_path, '{}.align.{}-{}'.format(split, "src", "tgt"))
if indexed_dataset.dataset_exists(align_path, impl=dataset_impl):
align_dataset = data_utils.load_indexed_dataset(align_path, None, dataset_impl)
tgt_dataset_sizes = tgt_dataset.sizes if tgt_dataset is not None else None
return GenerationPairDataset(
src_dataset, src_dataset.sizes, src_dict,
tgt_dataset, tgt_dataset_sizes, tgt_dict,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
max_source_positions=max_source_positions,
max_target_positions=max_target_positions,
align_dataset=align_dataset, eos=eos, bos=bos
)
def load_glue_data(
task_path,
mydict,
mode='train'
): #一个大列表，每个item是一个文档矩阵，矩阵里面每个item是一个node的数值 ，for token_id 和
# dataset = data_utils.load_indexed_dataset(path,mydict,'mmap',combine=False,)
# dataset = TokenBlockDataset(dataset,dataset.sizes,512 - 1,pad=mydict.pad(),eos=mydict.eos(), break_mode='complete',)
# dataset = PrependTokenDataset(dataset, mydict.bos())
#dataset=[]
#input1=open(input_path1,'r').readlines()#[:10000]
#label=open(label_path,'r').readlines()
input0 = data_utils.load_indexed_dataset(
os.path.join(task_path, 'input0', mode),
mydict,
'mmap',
combine=False,
)
assert input0 is not None, 'could not find dataset: {}'.format(
get_path(type, split))
input1 = data_utils.load_indexed_dataset(
os.path.join(task_path, 'input1', mode),
mydict,
'mmap',
combine=False,
)
input0 = PrependTokenDataset(input0, mydict.bos())
if input1 is None:
src_tokens = input0
else:
input1 = PrependTokenDataset(input1, mydict.eos())
src_tokens = ConcatSentencesDataset(input0, input1)
if not 'STS-B' in task_path:
label_dictionary = Dictionary.load(
os.path.join(task_path, 'label', 'dict.txt'))
label_dictionary.add_symbol('<mask>')
#label_dataset = make_dataset('label', label_dictionary)
label_dataset = data_utils.load_indexed_dataset(
os.path.join(task_path, 'label', mode),
label_dictionary,
'mmap',
combine=False,
)
if label_dataset is not None:
label = OffsetTokensDataset(
StripTokenDataset(
label_dataset,
id_to_strip=label_dictionary.eos(),
),
offset=-label_dictionary.nspecial,
)
else:
label_path = "{0}.label".format(os.path.join(task_path, 'label', mode))
if os.path.exists(label_path):
def parse_regression_target(i, line):
values = line.split()
assert len(values) == 1, \
f'expected num_classes={self.args.num_classes} regression target values on line {i}, found: "{line}"'
return [float(x) for x in values]
with open(label_path) as h:
label = RawLabelDataset([
parse_regression_target(i, line.strip())
for i, line in enumerate(h.readlines())
])
print('data size: ', len(src_tokens), len(label))
assert len(src_tokens) == len(label)
# with data_utils.numpy_seed(self.args.seed):
# shuffle = np.random.permutation(len(src_tokens))
# src_tokens = maybe_shorten_dataset(
# src_tokens,
# split,
# self.args.shorten_data_split_list,
# self.args.shorten_method,
# self.args.max_positions,
# self.args.seed,
# )
# input_data1=[]
# input_data2=[]
#label_list=[]
# for line in input1:
# if len(line.strip())==0:
# input_data1.append([])
# else:
# line = line.strip().split(' ')
# input_data1.append([int(x) for x in line])
# if input_path2:
# input2=open(input_path2,'r').readlines()
# for line in input2:
# if len(line.strip())==0:
# input_data2.append([])
# else:
# line = line.strip().split(' ')
# input_data2.append([int(x) for x in line])
# if task=='QNLI':
# for line in label:
# line = line.strip()
# if line=='entailment':
# label_list.append(int(1))
# else:
# assert line=='not_entailment'
# label_list.append(int(0))
# else:
# for line in label:
# line = line.strip()
# label_list.append(int(line))
# print('data length: ',len(input_data1),len(input_data2))
# assert len(input_data1)==len(label_list)
# if len(input_data2)!=0:
# assert len(input_data1)==len(input_data2)
return src_tokens, label
def load_langpair_with_additional_data_dataset(
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,
add_dir=None,
add_lang=None,
add_dict=None,
userdirname=None,
):
def split_exists(split, src, tgt, lang, data_path):
filename = os.path.join(data_path,
'{}.{}-{}.{}'.format(split, src, tgt, lang))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
src_datasets = []
tgt_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else '')
# infer langcode
if split_exists(split_k, src, tgt, src, data_path):
prefix = os.path.join(data_path,
'{}.{}-{}.'.format(split_k, src, tgt))
elif split_exists(split_k, tgt, src, src, data_path):
prefix = os.path.join(data_path,
'{}.{}-{}.'.format(split_k, tgt, src))
else:
if k > 0:
break
else:
raise FileNotFoundError('Dataset not found: {} ({})'.format(
split, data_path))
src_dataset = data_utils.load_indexed_dataset(prefix + src, src_dict,
dataset_impl)
if truncate_source:
src_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
)
src_datasets.append(src_dataset)
tgt_datasets.append(
data_utils.load_indexed_dataset(prefix + tgt, tgt_dict,
dataset_impl))
print('| {} {} {}-{} {} examples'.format(data_path, split_k, src, tgt,
len(src_datasets[-1])))
if not combine:
break
assert len(src_datasets) == len(tgt_datasets)
if len(src_datasets) == 1:
src_dataset, tgt_dataset = src_datasets[0], tgt_datasets[0]
else:
sample_ratios = [1] * len(src_datasets)
sample_ratios[0] = upsample_primary
src_dataset = ConcatDataset(src_datasets, sample_ratios)
tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(
tgt_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
align_dataset = None
if load_alignments:
align_path = os.path.join(data_path,
'{}.align.{}-{}'.format(split, src, tgt))
if indexed_dataset.dataset_exists(align_path, impl=dataset_impl):
align_dataset = data_utils.load_indexed_dataset(
align_path, None, dataset_impl)
add_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else '')
additional_data_path = f'{data_path}/{add_dir}'
# infer langcode
if split_exists(split_k, add_lang, 'None', add_lang,
additional_data_path):
prefix = os.path.join(
additional_data_path,
'{}.{}-{}.'.format(split_k, add_lang, 'None'))
else:
if k > 0:
break
else:
raise FileNotFoundError('Dataset not found: {} ({})'.format(
split, additional_data_path))
add_dataset = data_utils.load_indexed_dataset(prefix + add_lang,
add_dict, dataset_impl)
if truncate_source:
add_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(add_dataset, add_dict.eos()),
max_source_positions - 1,
),
add_dict.eos(),
)
add_datasets.append(add_dataset)
print('| {} {} {}-{} {} examples'.format(data_path, split_k,
add_lang, 'None',
len(add_datasets[-1])))
if not combine:
break
if len(add_datasets) == 1:
add_dataset = add_datasets[0]
else:
raise Exception
# sample_ratios = [1] * len(src_datasets)
# sample_ratios[0] = upsample_primary
# src_dataset = ConcatDataset(src_datasets, sample_ratios)
# tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
if add_dataset:
import sys
module_parent, module_name = os.path.split(
os.path.abspath(userdirname))
add_user_module(userdirname)
return sys.modules[
module_name].data.LanguagePairWithAdditionalDataDataset(
src_dataset,
src_dataset.sizes,
src_dict,
tgt_dataset,
tgt_dataset.sizes,
tgt_dict,
add_dataset,
add_dataset.sizes,
add_dict,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
max_source_positions=max_source_positions,
max_target_positions=max_target_positions,
align_dataset=align_dataset,
)
else:
return LanguagePairDataset(
src_dataset,
src_dataset.sizes,
src_dict,
tgt_dataset,
tgt_dataset.sizes,
tgt_dict,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
max_source_positions=max_source_positions,
max_target_positions=max_target_positions,
align_dataset=align_dataset,
)
def load_lang_dataset(
self,
data_path,
split,
src,
src_dict,
tgt,
tgt_dict,
combine,
dataset_impl,
upsample_primary,
max_source_positions,
prepend_bos=False,
load_alignments=False,
truncate_source=False,
):
src_datasets = []
tgt_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else "")
# infer langcode
if self.split_exists(split_k, src, tgt, src, data_path, dataset_impl):
prefix = os.path.join(data_path, "{}.{}-{}.".format(split_k, src, tgt))
elif self.split_exists(split_k, tgt, src, src, data_path, dataset_impl):
prefix = os.path.join(data_path, "{}.{}-{}.".format(split_k, tgt, src))
else:
if k > 0:
break
else:
logger.error(
f"Dataset not found: {data_path}, {split_k}, {src}, {tgt}"
)
raise FileNotFoundError(
"Dataset not found: {} ({})".format(split, data_path)
)
src_dataset = self.load_data(prefix + src, src_dict, dataset_impl)
if truncate_source:
src_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
)
src_datasets.append(src_dataset)
tgt_datasets.append(self.load_data(prefix + tgt, tgt_dict, dataset_impl))
logger.info(
"{} {} {}-{} {} examples".format(
data_path, split_k, src, tgt, len(src_datasets[-1])
)
)
if not combine:
break
assert len(src_datasets) == len(tgt_datasets)
if len(src_datasets) == 1:
src_dataset, tgt_dataset = src_datasets[0], tgt_datasets[0]
else:
sample_ratios = [1] * len(src_datasets)
sample_ratios[0] = upsample_primary
src_dataset = ConcatDataset(src_datasets, sample_ratios)
tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(tgt_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
align_dataset = None
if load_alignments:
align_path = os.path.join(
data_path, "{}.align.{}-{}".format(split, src, tgt)
)
if indexed_dataset.dataset_exists(align_path, impl=dataset_impl):
align_dataset = data_utils.load_indexed_dataset(
align_path, None, dataset_impl
)
return src_dataset, tgt_dataset, align_dataset
def load_seq_sql_dataset(data_path, split, src, src_dict, prev_src_dict, sql,
sql_dict, prev_sql_dict, encoder_embed_path,
encoder_embed_dim, decoder_embed_path,
decoder_embed_dim, encoder_random_embedding_path,
decoder_random_embedding_path, dataset_impl,
upsample_primary, left_pad_source, left_pad_target,
max_source_positions, max_target_positions,
truncate_source, prepend_bos):
src_datasets = []
sql_datasets = []
prefix = os.path.join(data_path, split)
src_dataset = data_utils.load_indexed_dataset(prefix + '.' + src, src_dict,
dataset_impl)
#col_sizes = get_col_sizes(prefix + '.col')
if truncate_source:
src_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
)
src_datasets.append(src_dataset)
sql_datasets.append(
data_utils.load_indexed_dataset(prefix + '.' + sql, sql_dict,
dataset_impl))
assert len(src_datasets) == len(sql_datasets)
if len(src_datasets) == 1:
src_dataset, sql_dataset = src_datasets[0], sql_datasets[0]
else: #not implemented
sample_ratios = [1] * len(src_datasets)
sample_ratios[0] = upsample_primary
src_dataset = ConcatDataset(src_datasets, sample_ratios)
sql_dataset = ConcatDataset(sql_datasets, sample_ratios)
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(
sql_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
sql_dataset = PrependTokenDataset(sql_dataset, sql_dict.bos())
return Seq2SqlPairDataSet(
src_dataset,
src_dataset.sizes,
src_dict,
prev_src_dict,
sql_dataset,
sql_dataset.sizes,
sql_dict,
prev_sql_dict,
encoder_embed_path,
encoder_embed_dim,
decoder_embed_path,
decoder_embed_dim,
encoder_random_embedding_path,
decoder_random_embedding_path,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
max_source_positions=max_source_positions,
max_target_positions=max_target_positions,
)
def load_langpair_dataset(
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, srcda=False, srcda_choice='uniform',
tgtda=False, tgtda_choice='uniform'
):
def split_exists(split, src, tgt, lang, data_path):
filename = os.path.join(data_path, '{}.{}-{}.{}'.format(split, src, tgt, lang))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
src_datasets = []
tgt_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else '')
# infer langcode
if split_exists(split_k, src, tgt, src, data_path):
prefix = os.path.join(data_path, '{}.{}-{}.'.format(split_k, src, tgt))
elif split_exists(split_k, tgt, src, src, data_path):
prefix = os.path.join(data_path, '{}.{}-{}.'.format(split_k, tgt, src))
else:
if k > 0:
break
else:
raise FileNotFoundError('Dataset not found: {} ({})'.format(split, data_path))
src_dataset = data_utils.load_indexed_dataset(prefix + src, src_dict, dataset_impl)
if truncate_source:
src_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
)
src_datasets.append(src_dataset)
tgt_datasets.append(
data_utils.load_indexed_dataset(prefix + tgt, tgt_dict, dataset_impl)
)
print('| {} {} {}-{} {} examples'.format(data_path, split_k, src, tgt, len(src_datasets[-1])))
if not combine:
break
assert len(src_datasets) == len(tgt_datasets)
if len(src_datasets) == 1:
src_dataset, tgt_dataset = src_datasets[0], tgt_datasets[0]
else:
sample_ratios = [1] * len(src_datasets)
sample_ratios[0] = upsample_primary
src_dataset = ConcatDataset(src_datasets, sample_ratios)
tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(tgt_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
align_dataset = None
if load_alignments:
align_path = os.path.join(data_path, '{}.align.{}-{}'.format(split, src, tgt))
if indexed_dataset.dataset_exists(align_path, impl=dataset_impl):
align_dataset = data_utils.load_indexed_dataset(align_path, None, dataset_impl)
return LanguagePairDatasetDA(
src_dataset, src_dataset.sizes, src_dict,
tgt_dataset, tgt_dataset.sizes, tgt_dict,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
max_source_positions=max_source_positions,
max_target_positions=max_target_positions,
align_dataset=align_dataset,
srcda=srcda, srcda_choice=srcda_choice,
tgtda=tgtda, tgtda_choice=tgtda_choice
)
def lang_dataset(lang):
input0 = make_dataset('input0', lang, self.source_dictionary)
assert input0 is not None, 'could not find dataset: {}'.format(
get_path('input0', lang, split))
input1 = make_dataset('input1', lang, self.source_dictionary)
if self.args.init_token is not None:
input0 = PrependTokenDataset(input0, self.args.init_token)
if input1 is None:
src_tokens = input0
else:
if self.args.separator_token is not None:
input1 = PrependTokenDataset(input1,
self.args.separator_token)
src_tokens = ConcatSentencesDataset(input0, input1)
with data_utils.numpy_seed(self.args.seed):
shuffle = np.random.permutation(len(src_tokens))
if self.args.truncate_sequence:
src_tokens = TruncateDataset(src_tokens,
self.args.max_positions)
dataset = {
'id': IdDataset(),
'net_input': {
'src_tokens':
RightPadDataset(
src_tokens,
pad_idx=self.source_dictionary.pad(),
),
'src_lengths':
NumelDataset(src_tokens, reduce=False),
},
'nsentences': NumSamplesDataset(),
'ntokens': NumelDataset(src_tokens, reduce=True),
}
if not self.args.regression_target:
label_dataset = make_dataset('label', lang,
self.target_dictionary)
if label_dataset is not None:
dataset.update(target=OffsetTokensDataset(
StripTokenDataset(
label_dataset,
id_to_strip=self.target_dictionary.eos(),
),
offset=-self.target_dictionary.nspecial,
))
else:
label_path = "{0}.label".format(get_path('label', lang, split))
if os.path.exists(label_path):
dataset.update(target=RawLabelDataset([
float(x.strip()) for x in open(label_path).readlines()
]))
nested_dataset = NestedDictionaryDataset(
dataset,
sizes=[src_tokens.sizes],
)
if self.args.no_shuffle:
dataset = nested_dataset
else:
dataset = SortDataset(
nested_dataset,
# shuffle
sort_order=[shuffle],
)
print("| Loaded {0} with #samples: {1}".format(
split, len(dataset)))
return dataset
def load_dataset(self, split, combine=False, **kwargs):
"""Load a given dataset split (e.g., train, valid, test)."""
def get_path(type, split):
return os.path.join(self.args.data, type, split)
def make_dataset(type, dictionary):
split_path = get_path(type, split)
dataset = data_utils.load_indexed_dataset(
split_path,
dictionary,
self.args.dataset_impl,
combine=combine,
)
return dataset
input0 = make_dataset('input0', self.source_dictionary)
assert input0 is not None, 'could not find dataset: {}'.format(get_path(type, split))
input1 = make_dataset('input1', self.source_dictionary)
if self.args.init_token is not None:
input0 = PrependTokenDataset(input0, self.args.init_token)
if input1 is None:
src_tokens = input0
else:
if self.args.separator_token is not None:
input1 = PrependTokenDataset(input1, self.args.separator_token)
src_tokens = ConcatSentencesDataset(input0, input1)
with data_utils.numpy_seed(self.args.seed):
shuffle = np.random.permutation(len(src_tokens))
src_tokens = maybe_shorten_dataset(
src_tokens,
split,
self.args.shorten_data_split_whitelist,
self.args.shorten_method,
self.args.max_positions,
self.args.seed,
)
dataset = {
'id': IdDataset(),
'net_input': {
'src_tokens': RightPadDataset(
src_tokens,
pad_idx=self.source_dictionary.pad(),
),
'src_lengths': NumelDataset(src_tokens, reduce=False),
},
'nsentences': NumSamplesDataset(),
'ntokens': NumelDataset(src_tokens, reduce=True),
}
if self.args.add_prev_output_tokens:
prev_tokens_dataset = RightPadDataset(
RollDataset(src_tokens, 1),
pad_idx=self.dictionary.pad(),
)
dataset['net_input'].update(
prev_output_tokens=prev_tokens_dataset,
)
if not self.args.regression_target:
label_dataset = make_dataset('label', self.label_dictionary)
if label_dataset is not None:
dataset.update(
target=OffsetTokensDataset(
StripTokenDataset(
label_dataset,
id_to_strip=self.label_dictionary.eos(),
),
offset=-self.label_dictionary.nspecial,
)
)
else:
label_path = "{0}.label".format(get_path('label', split))
if os.path.exists(label_path):
def parse_regression_target(i, line):
values = line.split()
assert len(values) == self.args.num_classes, \
f'expected num_classes={self.args.num_classes} regression target values on line {i}, found: "{line}"'
return [float(x) for x in values]
dataset.update(
target=RawLabelDataset([
parse_regression_target(i, line.strip()) for i, line in enumerate(open(label_path).readlines())
])
)
nested_dataset = NestedDictionaryDataset(
dataset,
sizes=[src_tokens.sizes],
)
if self.args.no_shuffle:
dataset = nested_dataset
else:
dataset = SortDataset(
nested_dataset,
# shuffle
sort_order=[shuffle],
)
logger.info("Loaded {0} with #samples: {1}".format(split, len(dataset)))
self.datasets[split] = dataset
return self.datasets[split]
def load_langpair_dataset(
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,
num_buckets=0,
shuffle=True,
pad_to_multiple=1,
# Masked LM parameters.
mask_idx: int = 0,
seed: int = 1,
mask_prob: float = 0.01,
leave_unmasked_prob: float = 0.0,
random_token_prob: float = 0.0,
freq_weighted_replacement: bool = False,
mask_whole_words: torch.Tensor = None,
mask_multiple_length: int = 1,
mask_stdev: float = 0.0,
):
def split_exists(split, src, tgt, lang, data_path):
filename = os.path.join(
data_path, "{}.{}-{}.{}".format(split, src, tgt, lang))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
src_datasets = []
tgt_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else "")
# infer langcode
if split_exists(split_k, src, tgt, src, data_path):
prefix = os.path.join(
data_path, "{}.{}-{}.".format(split_k, src, tgt))
elif split_exists(split_k, tgt, src, src, data_path):
prefix = os.path.join(
data_path, "{}.{}-{}.".format(split_k, tgt, src))
else:
if k > 0:
break
else:
raise FileNotFoundError(
"Dataset not found: {} ({})".format(split, data_path)
)
src_dataset = data_utils.load_indexed_dataset(
prefix + src, src_dict, dataset_impl
)
if truncate_source:
src_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
)
src_datasets.append(src_dataset)
tgt_dataset = data_utils.load_indexed_dataset(
prefix + tgt, tgt_dict, dataset_impl
)
if tgt_dataset is not None:
tgt_datasets.append(tgt_dataset)
logger.info(
"{} {} {}-{} {} examples".format(
data_path, split_k, src, tgt, len(src_datasets[-1])
)
)
if not combine:
break
# logger.info('Length of Source DataSets: {}'.format(len(src_datasets)))
assert len(src_datasets) == len(tgt_datasets) or len(tgt_datasets) == 0
if len(src_datasets) == 1:
src_dataset = src_datasets[0]
tgt_dataset = tgt_datasets[0] if len(tgt_datasets) > 0 else None
else:
sample_ratios = [1] * len(src_datasets)
sample_ratios[0] = upsample_primary
src_dataset = ConcatDataset(src_datasets, sample_ratios)
if len(tgt_datasets) > 0:
tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
else:
tgt_dataset = None
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(
tgt_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
if tgt_dataset is not None:
tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
eos = None
if append_source_id:
src_dataset = AppendTokenDataset(
src_dataset, src_dict.index("[{}]".format(src))
)
if tgt_dataset is not None:
tgt_dataset = AppendTokenDataset(
tgt_dataset, tgt_dict.index("[{}]".format(tgt))
)
eos = tgt_dict.index("[{}]".format(tgt))
align_dataset = None
if load_alignments:
align_path = os.path.join(
data_path, "{}.align.{}-{}".format(split, src, tgt))
if indexed_dataset.dataset_exists(align_path, impl=dataset_impl):
align_dataset = data_utils.load_indexed_dataset(
align_path, None, dataset_impl
)
tgt_dataset_sizes = tgt_dataset.sizes if tgt_dataset is not None else None
# mask source dataset.
src_dataset, masked_src_dataset = MaskTokensDataset.apply_mask(
src_dataset,
src_dict,
pad_idx=src_dict.pad(),
mask_idx=mask_idx,
seed=seed,
mask_prob=mask_prob,
leave_unmasked_prob=leave_unmasked_prob,
random_token_prob=random_token_prob,
freq_weighted_replacement=freq_weighted_replacement,
mask_whole_words=mask_whole_words,
mask_multiple_length=mask_multiple_length,
mask_stdev=mask_stdev,
)
# Print samples.
# if split == 'valid':
# print(src_dataset[1])
# print(masked_src_dataset[1])
return LanguagePairDataset(
src_dataset,
src_dataset.sizes,
src_dict,
tgt_dataset,
tgt_dataset_sizes,
tgt_dict,
# for Mask LM loss calculation.
masked_src_dataset,
masked_src_dataset.sizes,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
align_dataset=align_dataset,
eos=eos,
num_buckets=num_buckets,
shuffle=shuffle,
pad_to_multiple=pad_to_multiple,
)
def pos_loader(data_path,
split,
src,
src_dict,
tgt,
tgt_dict,
anchor,
anchor_dict,
combine,
dataset_impl,
upsample_primary,
left_pad_source,
left_pad_target,
max_source_positions,
max_target_positions,
prepend_bos=False,
truncate_source=False,
append_source_id=False):
# Check the existence of the file
def split_exists(split, src, tgt, lang, data_path):
filename = os.path.join(data_path,
'{}.{}-{}.{}'.format(split, src, tgt, lang))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
src_datasets = []
tgt_datasets = []
anchor_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else '')
# infer langcode (from a->b or from b->a)
if split_exists(split_k, src, tgt, src, data_path):
prefix = os.path.join(data_path,
'{}.{}-{}.'.format(split_k, src, tgt))
elif split_exists(split_k, tgt, src, src, data_path):
prefix = os.path.join(data_path,
'{}.{}-{}.'.format(split_k, tgt, src))
else:
if k > 0:
break
else:
raise FileNotFoundError('Dataset not found: {} ({})'.format(
split, data_path))
src_dataset = data_utils.load_indexed_dataset(prefix + src, src_dict,
dataset_impl)
if truncate_source:
src_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
)
src_datasets.append(src_dataset)
tgt_dataset = data_utils.load_indexed_dataset(prefix + tgt, tgt_dict,
dataset_impl)
if tgt_dataset is not None:
tgt_datasets.append(tgt_dataset)
anchor_prefix = os.path.join(data_path, anchor,
'{}.{}-{}.'.format(split_k, anchor, tgt))
anchor_dataset = data_utils.load_indexed_dataset(
anchor_prefix + anchor, anchor_dict, dataset_impl)
if anchor_dataset is not None:
anchor_datasets.append(anchor_dataset)
logger.info('{} {} {}-{} {} examples'.format(data_path, split_k,
src, tgt,
len(src_datasets[-1])))
if not combine:
break
assert len(src_datasets) == len(tgt_datasets) or len(tgt_datasets) == 0
# None is not avaliable for anchors
assert len(src_datasets) == len(anchor_datasets)
if len(src_datasets) == 1:
src_dataset = src_datasets[0]
tgt_dataset = tgt_datasets[0] if len(tgt_datasets) > 0 else None
anchor_dataset = anchor_datasets[0]
else:
sample_ratios = [1] * len(src_datasets)
sample_ratios[0] = upsample_primary
src_dataset = ConcatDataset(src_datasets, sample_ratios)
if len(tgt_datasets) > 0:
tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
else:
tgt_dataset = None
anchor_dataset = ConcatDataset(anchor_datasets, sample_ratios)
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(
tgt_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
if tgt_dataset is not None:
tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
eos = None
if append_source_id:
src_dataset = AppendTokenDataset(src_dataset,
src_dict.index('[{}]'.format(src)))
if tgt_dataset is not None:
tgt_dataset = AppendTokenDataset(
tgt_dataset, tgt_dict.index('[{}]'.format(tgt)))
eos = tgt_dict.index('[{}]'.format(tgt))
tgt_dataset_sizes = tgt_dataset.sizes if tgt_dataset is not None else None
return POSGraphLanguagePairDatasetb(
src_dataset,
src_dataset.sizes,
src_dict,
anchor_dataset,
anchor_dataset.sizes,
anchor_dict,
tgt_dataset,
tgt_dataset_sizes,
tgt_dict,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
max_source_positions=max_source_positions,
max_target_positions=max_target_positions,
eos=eos)
def load_langpair_dataset(
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,
num_buckets=0,
shuffle=True,
pad_to_multiple=1,
prepend_bos_src=None,
bert_model_name=None,
bart_model_name=None,
electra_model_name=None,
electra_pretrain=False,
denoising=False,
masking=False,
extra_data=False,
input_mapping=False,
mask_ratio=None,
random_ratio=None,
insert_ratio=None,
rotate_ratio=None,
permute_sentence_ratio=None,
):
def split_exists(split, src, tgt, lang, data_path):
filename = os.path.join(data_path,
"{}.{}-{}.{}".format(split, src, tgt, lang))
return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
src_datasets = []
tgt_datasets = []
bert_tokenizer = BertTokenizer.from_pretrained(bert_model_name,
do_lower_case=False)
if denoising:
bart_tokenizer = AutoTokenizer.from_pretrained(bart_model_name,
do_lower_case=False)
#bart_tokenizer = BartTokenizer.from_pretrained(bart_model_name, do_lower_case=False)
if electra_pretrain:
electra_tokenizer = ElectraTokenizer.from_pretrained(
electra_model_name)
srcbert_datasets = []
extra_datasets = []
extra_bert_datasets = []
extra_bert_mapping_datasets = []
extra_bart_datasets = []
extra_bart_mapping_datasets = []
if denoising:
srcbart_datasets = []
if electra_pretrain:
srcelectra_datasets = []
for k in itertools.count():
split_k = split + (str(k) if k > 0 else "")
# infer langcode
if split_exists(split_k, src, tgt, src, data_path):
prefix = os.path.join(data_path,
"{}.{}-{}.".format(split_k, src, tgt))
bertprefix = os.path.join(
data_path, '{}.bert.{}-{}.'.format(split_k, src, tgt))
bert_mapping_prefix = os.path.join(
data_path, '{}.bert.map.{}-{}.'.format(split_k, src, tgt))
if denoising:
bartprefix = os.path.join(
data_path, '{}.bart.{}-{}.'.format(split_k, src, tgt))
bart_mapping_prefix = os.path.join(
data_path, '{}.bart.map.{}-{}.'.format(split_k, src, tgt))
if electra_pretrain:
electraprefix = os.path.join(
data_path, '{}.electra.{}-{}.'.format(split_k, src, tgt))
electra_mapping_prefix = os.path.join(
data_path,
'{}.electra.map.{}-{}.'.format(split_k, src, tgt))
if extra_data:
extraprefix = os.path.join(
data_path, '{}.extra.{}-{}.'.format(split_k, src, tgt))
extra_bert_prefix = os.path.join(
data_path,
'{}.extra.bert.{}-{}.'.format(split_k, src, tgt))
extra_bert_mapping_prefix = os.path.join(
data_path,
'{}.extra.bert.map.{}-{}.'.format(split_k, src, tgt))
extra_bart_prefix = os.path.join(
data_path,
'{}.extra.bart.{}-{}.'.format(split_k, src, tgt))
extra_bart_mapping_prefix = os.path.join(
data_path,
'{}.extra.bart.map.{}-{}.'.format(split_k, src, tgt))
elif split_exists(split_k, tgt, src, src, data_path):
prefix = os.path.join(data_path,
"{}.{}-{}.".format(split_k, tgt, src))
bertprefix = os.path.join(
data_path, '{}.bert.{}-{}.'.format(split_k, tgt, src))
bert_mapping_prefix = os.path.join(
data_path, '{}.bert.map.{}-{}.'.format(split_k, src, tgt))
if denoising:
bartprefix = os.path.join(
data_path, '{}.bart.{}-{}.'.format(split_k, tgt, src))
bart_mapping_prefix = os.path.join(
data_path, '{}.bart.map.{}-{}.'.format(split_k, src, tgt))
if electra_pretrain:
electraprefix = os.path.join(
data_path, '{}.electra.{}-{}.'.format(split_k, src, tgt))
electra_mapping_prefix = os.path.join(
data_path,
'{}.electra.map.{}-{}.'.format(split_k, src, tgt))
if extra_data:
extraprefix = os.path.join(
data_path, '{}.extra.{}-{}.'.format(split_k, src, tgt))
extra_bert_prefix = os.path.join(
data_path,
'{}.extra.bert.{}-{}.'.format(split_k, src, tgt))
extra_bert_mapping_prefix = os.path.join(
data_path,
'{}.extra.bert.map.{}-{}.'.format(split_k, src, tgt))
extra_bart_prefix = os.path.join(
data_path,
'{}.extra.bart.{}-{}.'.format(split_k, src, tgt))
extra_bart_mapping_prefix = os.path.join(
data_path,
'{}.extra.bart.map.{}-{}.'.format(split_k, src, tgt))
else:
if k > 0:
break
else:
raise FileNotFoundError("Dataset not found: {} ({})".format(
split, data_path))
src_dataset = data_utils.load_indexed_dataset(prefix + src, src_dict,
dataset_impl)
if truncate_source:
src_dataset = AppendTokenDataset(
TruncateDataset(
StripTokenDataset(src_dataset, src_dict.eos()),
max_source_positions - 1,
),
src_dict.eos(),
)
src_datasets.append(src_dataset)
tgt_dataset = data_utils.load_indexed_dataset(prefix + tgt, tgt_dict,
dataset_impl)
if tgt_dataset is not None:
tgt_datasets.append(tgt_dataset)
# srcbert_datasets.append(indexed_dataset.make_dataset(bertprefix + src, impl=dataset_impl,
# fix_lua_indexing=True, ))
# if denoising:
# srcbart_datasets.append(indexed_dataset.make_dataset(bartprefix + src, impl=dataset_impl,
# fix_lua_indexing=True, ))
# if extra_data:
# extra_datasets.append(indexed_dataset.make_dataset(extraprefix + src, impl=dataset_impl,
# fix_lua_indexing=True, ))
srcbert_datasets.append(
data_utils.load_indexed_dataset(
bertprefix + src,
dataset_impl=dataset_impl,
))
if denoising:
srcbart_datasets.append(
data_utils.load_indexed_dataset(
bartprefix + src,
dataset_impl=dataset_impl,
))
if electra_pretrain:
srcelectra_datasets.append(
data_utils.load_indexed_dataset(
electraprefix + src,
dataset_impl=dataset_impl,
))
if extra_data and split == 'train':
extra_datasets.append(
data_utils.load_indexed_dataset(
extraprefix + src,
dataset_impl=dataset_impl,
))
extra_bert_datasets.append(
data_utils.load_indexed_dataset(
extra_bert_prefix + src,
dataset_impl=dataset_impl,
))
extra_bert_mapping_datasets.append(
data_utils.load_indexed_dataset(
extra_bert_mapping_prefix + src,
dataset_impl=dataset_impl,
))
extra_bart_datasets.append(
data_utils.load_indexed_dataset(
extra_bart_prefix + src,
dataset_impl=dataset_impl,
))
extra_bart_mapping_datasets.append(
data_utils.load_indexed_dataset(
extra_bart_mapping_prefix + src,
dataset_impl=dataset_impl,
))
#import pdb; pdb.set_trace()
assert extra_datasets != [] or extra_bert_datasets != [] or extra_bert_mapping_datasets != [] or extra_bart_datasets != [] or extra_bart_mapping_datasets != []
#extra_datasets = extra_datasets[0]
#import pdb; pdb.set_trace()
src_datasets[-1] = PrependTokenDataset(src_datasets[-1],
token=src_dict.bos_index)
if extra_data and split == 'train':
extra_datasets[-1] = PrependTokenDataset(extra_datasets[-1],
token=src_dict.bos_index)
if denoising is True:
if input_mapping is True and split == 'train':
bart_mapping_dataset = data_utils.load_indexed_dataset(
bart_mapping_prefix + src, dataset_impl=dataset_impl)
else:
bart_mapping_dataset = None
src_datasets[-1] = DenoisingBartDataset(
src_datasets[-1],
src_datasets[-1].sizes,
src_dict,
srcbart_datasets[-1],
srcbart_datasets[-1].sizes,
bart_tokenizer,
map_dataset=bart_mapping_dataset,
mask_ratio=mask_ratio,
random_ratio=random_ratio,
insert_ratio=insert_ratio,
rotate_ratio=rotate_ratio,
permute_sentence_ratio=permute_sentence_ratio,
)
if electra_pretrain is True:
if input_mapping is True and split == 'train':
electra_mapping_dataset = data_utils.load_indexed_dataset(
electra_mapping_prefix + src, dataset_impl=dataset_impl)
else:
electra_mapping_dataset = None
src_datasets[-1] = ElectrapretrainDataset(
src_datasets[-1],
src_datasets[-1].sizes,
src_dict,
srcelectra_datasets[-1],
srcelectra_datasets[-1].sizes,
electra_tokenizer,
map_dataset=electra_mapping_dataset,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
max_source_positions=max_source_positions,
max_target_positions=max_target_positions,
)
if masking is True:
if input_mapping is True and split == 'train':
#bert_mapping_dataset = indexed_dataset.make_dataset(bert_mapping_prefix + src, impl=dataset_impl, fix_lua_indexing=True)
bert_mapping_dataset = data_utils.load_indexed_dataset(
bert_mapping_prefix + src, dataset_impl=dataset_impl)
else:
bert_mapping_dataset = None
src_datasets[-1] = MaskingDataset(
src_datasets[-1],
src_datasets[-1].sizes,
src_dict,
srcbert_datasets[-1],
srcbert_datasets[-1].sizes,
bert_tokenizer,
map_dataset=bert_mapping_dataset,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
max_source_positions=max_source_positions,
max_target_positions=max_target_positions,
)
if extra_data is True and split == 'train':
assert input_mapping is True
src_datasets[-1] = MaskingExtraDataset(
src_datasets[-1],
src_datasets[-1].sizes,
src_dict,
extra_datasets[-1],
extra_datasets[-1].sizes,
extra_bert_datasets[-1],
extra_bert_datasets[-1].sizes,
bert_tokenizer,
map_dataset=extra_bert_mapping_datasets[-1],
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
max_source_positions=max_source_positions,
max_target_positions=max_target_positions,
)
src_datasets[-1] = DenoisingBartExtraDataset(
src_datasets[-1],
src_datasets[-1].sizes,
src_dict,
extra_datasets[-1],
extra_datasets[-1].sizes,
extra_bart_datasets[-1],
extra_bart_datasets[-1].sizes,
bart_tokenizer,
map_dataset=extra_bart_mapping_datasets[-1],
)
logger.info("{} {} {}-{} {} examples".format(data_path, split_k,
src, tgt,
len(src_datasets[-1])))
if not combine:
break
assert len(src_datasets) == len(tgt_datasets) or len(tgt_datasets) == 0
if len(src_datasets) == 1:
src_dataset = src_datasets[0]
tgt_dataset = tgt_datasets[0] if len(tgt_datasets) > 0 else None
# srcbert_datasets = srcbert_datasets[0]
# if denoising:
# srcbart_datasets = srcbart_datasets[0]
else:
sample_ratios = [1] * len(src_datasets)
sample_ratios[0] = upsample_primary
src_dataset = ConcatDataset(src_datasets, sample_ratios)
if len(tgt_datasets) > 0:
tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
else:
tgt_dataset = None
if prepend_bos:
assert hasattr(src_dict, "bos_index") and hasattr(
tgt_dict, "bos_index")
src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
if tgt_dataset is not None:
tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
elif prepend_bos_src is not None:
logger.info(f"prepending src bos: {prepend_bos_src}")
src_dataset = PrependTokenDataset(src_dataset, prepend_bos_src)
eos = None
if append_source_id:
src_dataset = AppendTokenDataset(src_dataset,
src_dict.index("[{}]".format(src)))
if tgt_dataset is not None:
tgt_dataset = AppendTokenDataset(
tgt_dataset, tgt_dict.index("[{}]".format(tgt)))
eos = tgt_dict.index("[{}]".format(tgt))
align_dataset = None
if load_alignments:
align_path = os.path.join(data_path,
"{}.align.{}-{}".format(split, src, tgt))
if indexed_dataset.dataset_exists(align_path, impl=dataset_impl):
align_dataset = data_utils.load_indexed_dataset(
align_path, None, dataset_impl)
tgt_dataset_sizes = tgt_dataset.sizes if tgt_dataset is not None else None
src_bart_dataset = None
src_bert_dataset = None
src_electra_dataset = None
return LanguagePairDataset(
src_dataset,
src_dataset.sizes,
src_dict,
tgt_dataset,
tgt_dataset_sizes,
tgt_dict,
masking,
src_bert_dataset,
denoising,
src_bart_dataset,
src_electra_dataset,
#extra_datasets,
left_pad_source=left_pad_source,
left_pad_target=left_pad_target,
align_dataset=align_dataset,
eos=eos,
num_buckets=num_buckets,
shuffle=shuffle,
pad_to_multiple=pad_to_multiple,
)