def _infer_metadata_from_saved_model(
    saved_model_dir: str,
    use_tf_compat_v1: bool) -> dataset_metadata.DatasetMetadata:
  """Infers a DatasetMetadata for outputs of a SavedModel."""
  if use_tf_compat_v1:
    return _infer_metadata_from_saved_model_v1(saved_model_dir)
  else:
    return _infer_metadata_from_saved_model_v2(saved_model_dir)
Infers a DatasetMetadata for outputs of a SavedModel.
tensorflow_transform/beam/impl.py