def get_payload_format(examples: types.Artifact) -> int:
"""Returns the payload format of Examples artifact.
If Examples artifact does not contain the "payload_format" custom property,
it is made before tfx supports multiple payload format, and can regard as
tf.Example format.
Args:
examples: A standard_artifacts.Examples artifact.
Returns:
payload_format: One of the enums in example_gen_pb2.PayloadFormat.
"""
assert examples.type is standard_artifacts.Examples, (
'examples must be of type standard_artifacts.Examples')
if examples.has_custom_property(
example_gen_utils.PAYLOAD_FORMAT_PROPERTY_NAME):
return example_gen_pb2.PayloadFormat.Value(
examples.get_string_custom_property(
example_gen_utils.PAYLOAD_FORMAT_PROPERTY_NAME))
else:
logging.warning(
'Examples artifact does not have %s custom property. '
'Falling back to %s',
example_gen_utils.PAYLOAD_FORMAT_PROPERTY_NAME,
example_gen_pb2.PayloadFormat.Name(_DEFAULT_PAYLOAD_FORMAT))
return _DEFAULT_PAYLOAD_FORMAT
def get_file_format(examples: types.Artifact) -> str:
"""Returns the file format of Examples artifact.
If Examples artifact does not contain the "file_format" custom property,
it is made by OSS ExampleGen and can be treated as 'tfrecords_gzip' format.
Args:
examples: A standard_artifacts.Examples artifact.
Returns:
One of the file format that tfx_bsl understands.
"""
assert examples.type_name == standard_artifacts.Examples.TYPE_NAME, (
'examples must be of type standard_artifacts.Examples')
if examples.has_custom_property(
example_gen_utils.FILE_FORMAT_PROPERTY_NAME):
return examples.get_string_custom_property(
example_gen_utils.FILE_FORMAT_PROPERTY_NAME)
else:
return _DEFAULT_FILE_FORMAT
