Fix Ruff rule B008

tensorflow · Oct 27, 2024 · f011508 · f011508
1 parent f5224b2
commit f011508
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 4 deletions.
diff --git a/tfx/dsl/component/experimental/decorators_test.py b/tfx/dsl/component/experimental/decorators_test.py
@@ -14,6 +14,7 @@
 """Tests for tfx.dsl.components.base.decorators."""
 
 
+from __future__ import annotations
 import pytest
 import os
 from typing import Any, Dict, List, Optional
@@ -141,8 +142,10 @@ def verify_beam_pipeline_arg(a: int) -> OutputDict(b=float):  # pytype: disable=
 
 def verify_beam_pipeline_arg_non_none_default_value(
     a: int,
-    beam_pipeline: BeamComponentParameter[beam.Pipeline] = beam.Pipeline(),
+    beam_pipeline: BeamComponentParameter[beam.Pipeline] = None,
 ) -> OutputDict(b=float):  # pytype: disable=invalid-annotation,wrong-arg-types
+  if beam_pipeline is None:
+    beam_pipeline = beam.Pipeline()
   del beam_pipeline
   return {'b': float(a)}
 

diff --git a/tfx/dsl/component/experimental/decorators_typeddict_test.py b/tfx/dsl/component/experimental/decorators_typeddict_test.py
@@ -14,6 +14,7 @@
 """Tests for tfx.dsl.components.base.decorators."""
 
 
+from __future__ import annotations
 import pytest
 import os
 from typing import Any, Dict, List, Optional, TypedDict
@@ -141,8 +142,10 @@ def verify_beam_pipeline_arg(a: int) -> TypedDict('Output6', dict(b=float)):  #
 
 def verify_beam_pipeline_arg_non_none_default_value(
     a: int,
-    beam_pipeline: BeamComponentParameter[beam.Pipeline] = beam.Pipeline(),
+    beam_pipeline: BeamComponentParameter[beam.Pipeline]  | None = None,
 ) -> TypedDict('Output7', dict(b=float)):  # pytype: disable=wrong-arg-types
+  if beam_pipeline is None:
+      beam_pipeline = beam.Pipeline()
   del beam_pipeline
   return {'b': float(a)}
 

diff --git a/tfx/examples/bert/utils/bert_models.py b/tfx/examples/bert/utils/bert_models.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Configurable fine-tuning BERT models for various tasks."""
 
+from __future__ import annotations
 from typing import Optional, List, Union
 
 import tensorflow as tf
@@ -59,8 +60,7 @@ def build_bert_classifier(bert_layer: tf.keras.layers.Layer,
 
 def compile_bert_classifier(
     model: tf.keras.Model,
-    loss: tf.keras.losses.Loss = tf.keras.losses.SparseCategoricalCrossentropy(
-        from_logits=True),
+    loss: tf.keras.losses.Loss | None = None,
     learning_rate: float = 2e-5,
     metrics: Optional[List[Union[str, tf.keras.metrics.Metric]]] = None):
   """Compile the BERT classifier using suggested parameters.
@@ -79,6 +79,9 @@ def compile_bert_classifier(
   Returns:
     None.
   """
+  if loss is None:
+    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+
   if metrics is None:
     metrics = ["sparse_categorical_accuracy"]