def test_load_dataset(self):
batch_tokens_size = 100
train_data_path = os.path.join(self.get_temp_dir(), 'train.tf_record')
_create_fake_dataset(train_data_path)
data_config = wmt_dataloader.WMTDataConfig(
input_path=train_data_path,
max_seq_length=35,
global_batch_size=batch_tokens_size,
is_training=True,
static_batch=False)
dataset = wmt_dataloader.WMTDataLoader(data_config).load()
examples = next(iter(dataset))
inputs, targets = examples['inputs'], examples['targets']
logging.info('dynamic inputs=%s targets=%s', inputs, targets)
data_config = wmt_dataloader.WMTDataConfig(
input_path=train_data_path,
max_seq_length=35,
global_batch_size=batch_tokens_size,
is_training=True,
static_batch=True)
dataset = wmt_dataloader.WMTDataLoader(data_config).load()
examples = next(iter(dataset))
inputs, targets = examples['inputs'], examples['targets']
logging.info('static inputs=%s targets=%s', inputs, targets)
self.assertEqual(inputs.shape, (2, 35))
self.assertEqual(targets.shape, (2, 35))
def test_load_dataset_raise_invalid_window(self):
batch_tokens_size = 10 # this is too small to form buckets.
train_data_path = os.path.join(self.get_temp_dir(), 'train.tf_record')
_create_fake_dataset(train_data_path)
data_config = wmt_dataloader.WMTDataConfig(
input_path=train_data_path,
max_seq_length=100,
global_batch_size=batch_tokens_size)
with self.assertRaisesRegex(
ValueError,
'The token budget, global batch size, is too small.*'):
_ = wmt_dataloader.WMTDataLoader(data_config).load()
def test_load_dataset_raise_invalid_window(self):
batch_tokens_size = 10 # this is too small to form buckets.
data_config = wmt_dataloader.WMTDataConfig(
input_path=self._record_train_input_path,
max_seq_length=100,
global_batch_size=batch_tokens_size,
is_training=True,
static_batch=False,
src_lang='en',
tgt_lang='reverse_en',
sentencepiece_model_path=self._sentencepeice_model_path)
with self.assertRaisesRegex(
ValueError, 'The token budget, global batch size, is too small.*'):
_ = wmt_dataloader.WMTDataLoader(data_config).load()
def test_load_dataset(
self, is_training, static_batch, batch_size, expected_shape):
data_config = wmt_dataloader.WMTDataConfig(
input_path=self._record_train_input_path
if is_training else self._record_test_input_path,
max_seq_length=35,
global_batch_size=batch_size,
is_training=is_training,
static_batch=static_batch,
src_lang='en',
tgt_lang='reverse_en',
sentencepiece_model_path=self._sentencepeice_model_path)
dataset = wmt_dataloader.WMTDataLoader(data_config).load()
examples = next(iter(dataset))
inputs, targets = examples['inputs'], examples['targets']
self.assertEqual(inputs.shape, expected_shape)
self.assertEqual(targets.shape, expected_shape)