Internal change

PiperOrigin-RevId: 300477605

official/nlp/albert/export_albert_tfhub.py

@@ -78,7 +78,6 @@ def export_albert_tfhub(albert_config: configs.AlbertConfig,
 
 
 def main(_):
-  assert tf.version.VERSION.startswith('2.')
   albert_config = configs.AlbertConfig.from_json_file(
       FLAGS.albert_config_file)
   export_albert_tfhub(albert_config, FLAGS.model_checkpoint_path,

official/nlp/albert/export_albert_tfhub_test.py

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

official/nlp/albert/run_classifier.py

@@ -33,7 +33,6 @@ FLAGS = flags.FLAGS
 
 def main(_):
   # Users should always run this script under TF 2.x
-  assert tf.version.VERSION.startswith('2.')
 
   with tf.io.gfile.GFile(FLAGS.input_meta_data_path, 'rb') as reader:
     input_meta_data = json.loads(reader.read().decode('utf-8'))

official/nlp/albert/run_squad.py

@@ -80,7 +80,6 @@ def export_squad(model_export_path, input_meta_data):
 
 def main(_):
   # Users should always run this script under TF 2.x
-  assert tf.version.VERSION.startswith('2.')
 
   with tf.io.gfile.GFile(FLAGS.input_meta_data_path, 'rb') as reader:
     input_meta_data = json.loads(reader.read().decode('utf-8'))

official/nlp/albert/tf2_albert_encoder_checkpoint_converter.py

@@ -122,7 +122,6 @@ def convert_checkpoint(bert_config, output_path, v1_checkpoint):
 
 
 def main(_):
-  assert tf.version.VERSION.startswith('2.')
   output_path = FLAGS.converted_checkpoint_path
   v1_checkpoint = FLAGS.checkpoint_to_convert
   albert_config = configs.AlbertConfig.from_json_file(FLAGS.albert_config_file)

official/nlp/bert/export_tfhub.py

@@ -77,7 +77,6 @@ def export_bert_tfhub(bert_config: configs.BertConfig,
 
 
 def main(_):
-  assert tf.version.VERSION.startswith('2.')
   bert_config = configs.BertConfig.from_json_file(FLAGS.bert_config_file)
   export_bert_tfhub(bert_config, FLAGS.model_checkpoint_path, FLAGS.export_path,
                     FLAGS.vocab_file)

official/nlp/bert/export_tfhub_test.py

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

official/nlp/bert/run_classifier.py

@@ -403,7 +403,6 @@ def run_bert(strategy,
 
 def main(_):
   # Users should always run this script under TF 2.x
-  assert tf.version.VERSION.startswith('2.')
 
   with tf.io.gfile.GFile(FLAGS.input_meta_data_path, 'rb') as reader:
     input_meta_data = json.loads(reader.read().decode('utf-8'))

official/nlp/bert/run_pretraining.py

@@ -159,7 +159,6 @@ def run_bert_pretrain(strategy):
 
 def main(_):
   # Users should always run this script under TF 2.x
-  assert tf.version.VERSION.startswith('2.')
   gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_param)
   if not FLAGS.model_dir:
     FLAGS.model_dir = '/tmp/bert20/'

official/nlp/bert/run_squad.py

@@ -76,7 +76,6 @@ def export_squad(model_export_path, input_meta_data):
 
 def main(_):
   # Users should always run this script under TF 2.x
-  assert tf.version.VERSION.startswith('2.')
 
   with tf.io.gfile.GFile(FLAGS.input_meta_data_path, 'rb') as reader:
     input_meta_data = json.loads(reader.read().decode('utf-8'))

official/nlp/bert/tf2_encoder_checkpoint_converter.py

@@ -98,7 +98,7 @@ def convert_checkpoint(bert_config, output_path, v1_checkpoint):
 
 
 def main(_):
-  assert tf.version.VERSION.startswith('2.')
+  tf.enable_v2_behavior()
   output_path = FLAGS.converted_checkpoint_path
   v1_checkpoint = FLAGS.checkpoint_to_convert
   bert_config = configs.BertConfig.from_json_file(FLAGS.bert_config_file)

official/nlp/modeling/layers/transformer_test.py

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

official/nlp/modeling/models/bert_classifier.py

@@ -19,7 +19,7 @@ from __future__ import division
 # from __future__ import google_type_annotations
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v2 as tf
 
 from official.nlp.modeling import networks
 

official/nlp/modeling/models/bert_classifier_test.py

@@ -18,7 +18,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v2 as tf
 
 from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling import networks

official/nlp/modeling/models/bert_pretrainer.py

@@ -20,7 +20,7 @@ from __future__ import division
 from __future__ import print_function
 
 import copy
-import tensorflow as tf
+import tensorflow.compat.v2 as tf
 
 from official.nlp.modeling import networks
 

official/nlp/modeling/models/bert_pretrainer_test.py

@@ -18,7 +18,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v2 as tf
 
 from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling import networks

official/nlp/modeling/models/bert_span_labeler.py

@@ -19,7 +19,7 @@ from __future__ import division
 # from __future__ import google_type_annotations
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v2 as tf
 
 from official.nlp.modeling import networks
 

official/nlp/modeling/models/bert_span_labeler_test.py

@@ -18,7 +18,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v2 as tf
 
 from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling import networks

official/nlp/modeling/networks/albert_transformer_encoder_test.py

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

official/nlp/modeling/networks/encoder_scaffold_test.py

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