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))