def test_load(self):
tfrecord_file = self._get_tfrecord_file()
meta_data_file = self._get_meta_data_file()
dataset, meta_data = dataloader.load(tfrecord_file, meta_data_file,
self.model_spec)
for i, (input_ids, label_ids) in enumerate(dataset):
self.assertEqual(i, 0)
self.assertTrue((input_ids.numpy() == [0, 1, 2, 3]).all())
self.assertTrue((label_ids.numpy() == [0]).all())
self.assertEqual(meta_data['size'], 1)
self.assertEqual(meta_data['num_classes'], 1)
self.assertEqual(meta_data['index_to_label'], ['0'])
def load(tfrecord_file, meta_data_file, model_spec):
"""Gets `TextClassifierDataLoader` object from tfrecord file and metadata file."""
dataset, meta_data = dataloader.load(tfrecord_file, meta_data_file,
model_spec)
tf.compat.v1.logging.info(
'Load preprocessed data and metadata from %s and %s '
'with size: %d, num_classes: %d', tfrecord_file, meta_data_file,
meta_data['size'], meta_data['num_classes'])
return TextClassifierDataLoader(dataset, meta_data['size'],
meta_data['num_classes'],
meta_data['index_to_label'])