def test_metadata_correctness(self):
# Get metadata.
metadata = model_conversion_beam_utils.get_pipeline_metadata(
self.base_experiment_dir, self.xids, self.output_dir,
self.conversion_types, 'suffix')
# Check correctness.
self.assertLen(
metadata,
len(self.xids) * len(self.folders) * len(self.conversion_types))
xid_counts = {k: 0 for k in self.xids}
model_num_counts = {k: 0 for k in range(len(self.folders))}
conversion_type_counts = {k: 0 for k in self.conversion_types}
for m in metadata:
xid_counts[m.xid] += 1
model_num_counts[m.model_num] += 1
conversion_type_counts[m.conversion_type] += 1
for xid in self.xids:
self.assertEqual(xid_counts[xid], 6)
for k in range(3):
self.assertEqual(model_num_counts[k], 6)
for ctype in conversion_type_counts:
self.assertEqual(conversion_type_counts[ctype], 9)
# Check output uniqueness.
all_output_filenames = [m.output_filename for m in metadata]
self.assertEqual(len(set(all_output_filenames)),
len(all_output_filenames))
def main(unused_argv):
beam_options = None
# Get metadata for conversion.
metadata = utils.get_pipeline_metadata(FLAGS.base_experiment_dir,
FLAGS.xids, FLAGS.output_dir,
FLAGS.conversion_types,
FLAGS.output_suffix)
if not metadata:
raise ValueError(
f'No data found: {FLAGS.base_experiment_dir}, {FLAGS.xids}')
logging.info('%i models in %i xids.', len(metadata), len(FLAGS.xids))
# Check that models don't already exist, and create directories if necessary.
for m in metadata:
utils.sanity_check_output_filename(m.output_filename)
logging.info('Starting to create flume pipeline...')
def _convert_and_write_model(m):
utils.convert_and_write_model(m,
include_frontend=FLAGS.include_frontend,
sanity_check=FLAGS.sanity_check)
return m
# Make and run beam pipeline.
with beam.Pipeline(beam_options) as root:
_ = (root
| 'MakeMetadataCollection' >> beam.Create(metadata)
| 'ConvertAndWriteModelsToDisk' >>
beam.Map(_convert_and_write_model))
def test_metadata_sanity(self):
# Get metadata.
metadata = model_conversion_beam_utils.get_pipeline_metadata(
self.base_experiment_dir, self.xids, self.output_dir,
self.conversion_types)
for m in metadata:
model_conversion_beam_utils.sanity_check_output_filename(
m.output_filename)