Internal change

PiperOrigin-RevId: 300477605

official/nlp/modeling/networks/transformer_encoder_test.py

@@ -200,5 +200,4 @@ class TransformerEncoderTest(keras_parameterized.TestCase):
 
 
 if __name__ == "__main__":
-  assert tf.version.VERSION.startswith('2.')
   tf.test.main()

official/nlp/optimization.py

@@ -20,7 +20,7 @@ from __future__ import print_function
 
 import re
 
-import tensorflow as tf
+import tensorflow.compat.v2 as tf
 
 
 class WarmUp(tf.keras.optimizers.schedules.LearningRateSchedule):

official/nlp/transformer/model_utils_test.py

@@ -59,5 +59,4 @@ class ModelUtilsTest(tf.test.TestCase):
 
 
 if __name__ == "__main__":
-  assert tf.version.VERSION.startswith('2.')
   tf.test.main()

official/nlp/transformer/transformer_layers_test.py

@@ -94,5 +94,4 @@ class TransformerLayersTest(tf.test.TestCase):
 
 
 if __name__ == "__main__":
-  tf.compat.v1.enable_v2_behavior()
   tf.test.main()

official/nlp/transformer/transformer_main.py

@@ -474,7 +474,6 @@ def main(_):
 
 
 if __name__ == "__main__":
-  tf.compat.v1.enable_v2_behavior()
   logging.set_verbosity(logging.INFO)
   misc.define_transformer_flags()
   app.run(main)

official/nlp/transformer/transformer_main_test.py

@@ -187,5 +187,4 @@ class TransformerTaskTest(tf.test.TestCase):
 
 
 if __name__ == '__main__':
-  tf.compat.v1.enable_v2_behavior()
   tf.test.main()

official/nlp/transformer/transformer_test.py

@@ -65,5 +65,4 @@ class TransformerV2Test(tf.test.TestCase):
 
 
 if __name__ == "__main__":
-  tf.compat.v1.enable_v2_behavior()
   tf.test.main()

official/nlp/transformer/utils/tokenizer.py

@@ -26,7 +26,7 @@ import unicodedata
 import numpy as np
 import six
 from six.moves import xrange  # pylint: disable=redefined-builtin
-import tensorflow as tf
+import tensorflow.compat.v2 as tf
 
 PAD = "<pad>"
 PAD_ID = 0

official/nlp/transformer/utils/tokenizer_test.py

@@ -17,7 +17,7 @@
 import collections
 import tempfile
 
-import tensorflow as tf  # pylint: disable=g-bad-import-order
+import tensorflow.compat.v2 as tf  # pylint: disable=g-bad-import-order
 
 from official.nlp.transformer.utils import tokenizer
 

official/nlp/xlnet/preprocess_classification_data.py

@@ -455,5 +455,5 @@ def main(_):
 
 if __name__ == "__main__":
 
-  assert tf.version.VERSION.startswith('2.')
+  tf.enable_v2_behavior()
   app.run(main)

official/nlp/xlnet/run_classifier.py

@@ -193,5 +193,4 @@ def main(unused_argv):
 
 
 if __name__ == "__main__":
-  assert tf.version.VERSION.startswith('2.')
   app.run(main)

official/nlp/xlnet/run_pretrain.py

@@ -153,5 +153,4 @@ def main(unused_argv):
 
 
 if __name__ == "__main__":
-  assert tf.version.VERSION.startswith('2.')
   app.run(main)

official/nlp/xlnet/run_squad.py

@@ -301,5 +301,4 @@ def main(unused_argv):
 
 
 if __name__ == "__main__":
-  assert tf.version.VERSION.startswith('2.')
   app.run(main)

official/nlp/xlnet/xlnet_modeling_test.py

@@ -49,5 +49,4 @@ class PositionalEmbeddingLayerTest(tf.test.TestCase):
     self.assertAllClose(pos_emb, target)
 
 if __name__ == "__main__":
-  assert tf.version.VERSION.startswith('2.')
   tf.test.main()