def _write_dtd_to_gcs(storage_client_fn, outdir, mae_tag_categories,
task_name):
"""Write the DTD config file."""
storage_client = storage_client_fn()
dtd_contents = mae.generate_dtd(mae_tag_categories, task_name)
bucket_name, blob_dir = split_gcs_name(outdir)
bucket = storage_client.get_bucket(bucket_name)
blob = bucket.blob(posixpath.join(blob_dir, 'classification.dtd'))
blob.upload_from_string(dtd_contents)
def testGenerateDtd(self):
mae_tag_categories = [{'name': 'type1'}, {'name': 'type2'}]
dtd = mae.generate_dtd(mae_tag_categories, 'task_name')
expected_dtd = """<!ENTITY name "task_name">
<!ELEMENT type1 ( #PCDATA ) >
<!ATTLIST type1 id ID prefix="type1" #REQUIRED >
<!ELEMENT type2 ( #PCDATA ) >
<!ATTLIST type2 id ID prefix="type2" #REQUIRED >
"""
self.assertEqual(dtd, expected_dtd)
def _write_dtd(storage_client_fn, outdir, mae_tag_categories, task_name):
if outdir.startswith('gs://'):
return _write_dtd_to_gcs(storage_client_fn, outdir, mae_tag_categories,
task_name)
with open(os.path.join(outdir, 'classification.dtd'), 'w') as f:
f.write(mae.generate_dtd(mae_tag_categories, task_name))