Address comments

googleapis · Jun 30, 2021 · f9f4ad4 · f9f4ad4
1 parent 279492f
commit f9f4ad4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 16 deletions.
diff --git a/google/cloud/aiplatform/explain/metadata/tf/v2/saved_model_metadata_builder.py b/google/cloud/aiplatform/explain/metadata/tf/v2/saved_model_metadata_builder.py
@@ -30,7 +30,7 @@ def __init__(
         self,
         model_path: str,
         signature_name: Optional[str] = None,
-        outputs_to_explain: Optional[List[str]] = (),
+        outputs_to_explain: Optional[List[str]] = None,
         **kwargs
     ) -> None:
         """Initializes a SavedModelMetadataBuilder object.
@@ -93,26 +93,18 @@ def _infer_metadata_entries_from_model(
               Inferred input metadata and output metadata from the model.
 
         Raises:
-            ValueError if specified name is not found in signature outputs.
+              ValueError if specified name is not found in signature outputs.
         """
 
         loaded_sig = self._loaded_model.signatures[signature_name]
         _, input_sig = loaded_sig.structured_input_signature
         output_sig = loaded_sig.structured_outputs
         input_mds = {}
         for name, tensor_spec in input_sig.items():
-            if tensor_spec.dtype.is_floating:
-                input_mds[
-                    name
-                ] = explanation_metadata.ExplanationMetadata.InputMetadata(
-                    input_tensor_name=name
-                )
-            else:
-                input_mds[
-                    name
-                ] = explanation_metadata.ExplanationMetadata.InputMetadata(
-                    input_tensor_name=name, modality="categorical",
-                )
+            input_mds[name] = explanation_metadata.ExplanationMetadata.InputMetadata(
+                input_tensor_name=name,
+                modality=None if tensor_spec.dtype.is_floating else "categorical",
+            )
 
         output_mds = {}
         for name in output_sig:

diff --git a/setup.py b/setup.py
@@ -36,8 +36,8 @@
 ]
 metadata_extra_require = ["pandas >= 1.0.0"]
 xai_extra_require = ["tensorflow-cpu>=2.3.0, <=2.5.0"]
-full_extra_require = (
-    tensorboard_extra_require + metadata_extra_require + xai_extra_require
+full_extra_require = list(
+    set(tensorboard_extra_require + metadata_extra_require + xai_extra_require)
 )
 testing_extra_require = full_extra_require + ["grpcio-testing ~= 1.34.0"]