def test_iterator_single_dataset(self):
r"""Tests iterating over a single dataset.
"""
data = MonoTextData(self._test_hparams)
data_iterator = DataIterator(data)
data_iterator.switch_to_dataset(dataset_name="data")
iterator = data_iterator.get_iterator()
i = 1001
for idx, batch in enumerate(iterator):
self.assertEqual(batch.batch_size, self._test_hparams['batch_size'])
np.testing.assert_array_equal(batch['length'], [1, 1])
for example in batch['text']:
self.assertEqual(example[0], str(i))
i += 1
self.assertEqual(i, 2001)
def test_iterator_multi_datasets(self):
r"""Tests iterating over multiple datasets.
"""
train = MonoTextData(self._train_hparams)
test = MonoTextData(self._test_hparams)
train_batch_size = self._train_hparams["batch_size"]
test_batch_size = self._test_hparams["batch_size"]
data_iterator = DataIterator({"train": train, "test": test})
data_iterator.switch_to_dataset(dataset_name="train")
iterator = data_iterator.get_iterator()
for idx, val in enumerate(iterator):
self.assertEqual(len(val), train_batch_size)
number = idx * train_batch_size + 1
self.assertEqual(val.text[0], [str(number)])
# numbers: 1 - 2000, first 4 vocab entries are special tokens
self.assertEqual(val.text_ids[0], torch.tensor(number + 3))
data_iterator.switch_to_dataset(dataset_name="test")
iterator = data_iterator.get_iterator()
for idx, val in enumerate(iterator):
self.assertEqual(len(val), test_batch_size)
number = idx * test_batch_size + 1001
self.assertEqual(val.text[0], [str(number)])
self.assertEqual(val.text_ids[0], torch.tensor(number + 3))
# test `get_iterator` interface
for idx, val in enumerate(data_iterator.get_iterator('train')):
self.assertEqual(len(val), train_batch_size)
number = idx * train_batch_size + 1
self.assertEqual(val.text[0], [str(number)])
self.assertEqual(val.text_ids[0], torch.tensor(number + 3))
# test exception for invalid dataset name
with self.assertRaises(ValueError) as context:
data_iterator.switch_to_dataset('val')
self.assertTrue('not found' in str(context.exception))