Fix Ruff B008 errors

janasangeetha · janasangeetha · commit 1ccbc833e85e · 2024-11-06T14:33:21.000+05:30
diff --git a/tfx/dsl/component/experimental/decorators_test.py b/tfx/dsl/component/experimental/decorators_test.py
@@ -140,8 +140,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] = 0,
 ) -> OutputDict(b=float):  # pytype: disable=invalid-annotation,wrong-arg-types
+  if beam_pipeline == 0:
+    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
@@ -140,8 +140,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] = 0,
 ) -> TypedDict('Output7', dict(b=float)):  # pytype: disable=wrong-arg-types
+  if beam_pipeline == 0:
+    beam_pipeline = beam.Pipeline()
   del beam_pipeline
   return {'b': float(a)}
 
@@ -807,4 +809,4 @@ def testListOfArtifacts(self):
         ],
     )
 
-    beam_dag_runner.BeamDagRunner().run(test_pipeline)
+    beam_dag_runner.BeamDagRunner().run(test_pipeline)
diff --git a/tfx/examples/bert/utils/bert_models.py b/tfx/examples/bert/utils/bert_models.py
@@ -59,16 +59,15 @@ 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.
 
   Args:
     model: A keras model. Most likely the output of build_bert_classifier.
-    loss: tf.keras.losses. The suggested loss function expects integer labels
-      (e.g. 0, 1, 2). If the labels are one-hot encoded, consider using
+    loss: Default None will use tf.keras.losses. The suggested loss function expects
+      integer labels (e.g. 0, 1, 2). If the labels are one-hot encoded, consider using
       tf.keras.lossesCategoricalCrossEntropy with from_logits set to true.
     learning_rate: Suggested learning rate to be used in
       tf.keras.optimizer.Adam. The three suggested learning_rates for
@@ -79,6 +78,8 @@ 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"]
 
