Don't require transformed_metadata to exist when using TFTransformOut…

…put.vocabulary_size_by_name PiperOrigin-RevId: 421615884
tensorflow · Jan 13, 2022 · 10ad0a4 · 10ad0a4
1 parent 963f2c4
commit 10ad0a4
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/tensorflow_transform/output_wrapper.py b/tensorflow_transform/output_wrapper.py
@@ -88,8 +88,7 @@ def transformed_metadata(self) -> dataset_metadata.DatasetMetadata:
     """A DatasetMetadata."""
     if self._transformed_metadata is None:
       self._transformed_metadata = metadata_io.read_metadata(
-          os.path.join(self._transform_output_dir,
-                       self.TRANSFORMED_METADATA_DIR))
+          self._transformed_metadata_dir)
     return self._transformed_metadata
 
   @property
@@ -108,6 +107,11 @@ def _exported_as_v1(self) -> bool:
           self.transform_savedmodel_dir)
     return self._exported_as_v1_value
 
+  @property
+  def _transformed_metadata_dir(self) -> str:
+    return os.path.join(self._transform_output_dir,
+                        self.TRANSFORMED_METADATA_DIR)
+
   def transformed_feature_spec(self) -> Dict[str, common_types.FeatureSpecType]:
     """Returns a feature_spec for the transformed features.
 
@@ -150,8 +154,7 @@ def vocabulary_file_by_name(self, vocab_filename: str) -> Optional[str]:
     Args:
       vocab_filename: The vocabulary name to lookup.
     """
-    mapping_path = os.path.join(self._transform_output_dir,
-                                self.TRANSFORMED_METADATA_DIR, self.ASSET_MAP)
+    mapping_path = os.path.join(self._transformed_metadata_dir, self.ASSET_MAP)
 
     mapping = {}
     if tf.io.gfile.exists(mapping_path):
@@ -181,8 +184,12 @@ def _vocabulary_size_from_annotations(self,
     if not common.IS_ANNOTATIONS_PB_AVAILABLE:
       return None
 
+    try:
+      schema = self.transformed_metadata.schema
+    except IOError:
+      return None
+
     from tensorflow_transform import annotations_pb2  # pylint: disable=g-import-not-at-top
-    schema = self.transformed_metadata.schema
     for annotation in schema.annotation.extra_metadata:
       message = annotations_pb2.VocabularyMetadata()
       annotation.Unpack(message)