Fix xlnet imports

PiperOrigin-RevId: 268968289

official/nlp/xlnet/preprocess_pretrain_data.py

@@ -23,6 +23,7 @@ import json
 import os
 import random
 
+from absl import app
 from absl import flags
 import absl.logging as _logging  # pylint: disable=unused-import
 
@@ -994,4 +995,4 @@ if __name__ == "__main__":
                        "using multiple workers to identify each worker.")
 
   tf.logging.set_verbosity(tf.logging.INFO)
-  absl_app.run(create_data)
+  app.run(create_data)

official/nlp/xlnet/run_classifier.py

@@ -27,13 +27,13 @@ from absl import logging
 import numpy as np
 import tensorflow as tf
 # pylint: disable=unused-import
-# Initialize TPU System.
 from official.nlp import xlnet_config
 from official.nlp import xlnet_modeling as modeling
 from official.nlp.xlnet import common_flags
 from official.nlp.xlnet import data_utils
 from official.nlp.xlnet import optimization
 from official.nlp.xlnet import training_utils
+from official.utils.misc import tpu_lib
 
 flags.DEFINE_integer("n_class", default=2, help="Number of classes.")
 
@@ -135,7 +135,7 @@ def main(unused_argv):
   if FLAGS.strategy_type == "mirror":
     strategy = tf.distribute.MirroredStrategy()
   elif FLAGS.strategy_type == "tpu":
-
+    # Initialize TPU System.
     cluster_resolver = tpu_lib.tpu_initialize(FLAGS.tpu)
     strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)
     use_remote_tpu = True

official/nlp/xlnet/run_pretrain.py

@@ -26,13 +26,13 @@ from absl import flags
 from absl import logging
 import tensorflow as tf
 # pylint: disable=unused-import
-# Initialize TPU System.
 from official.nlp import xlnet_config
 from official.nlp import xlnet_modeling as modeling
 from official.nlp.xlnet import common_flags
 from official.nlp.xlnet import data_utils
 from official.nlp.xlnet import optimization
 from official.nlp.xlnet import training_utils
+from official.utils.misc import tpu_lib
 
 flags.DEFINE_integer(
     "mask_alpha", default=6, help="How many tokens to form a group.")
@@ -66,7 +66,7 @@ def main(unused_argv):
   if FLAGS.strategy_type == "mirror":
     strategy = tf.distribute.MirroredStrategy()
   elif FLAGS.strategy_type == "tpu":
-
+    # Initialize TPU System.
     cluster_resolver = tpu_lib.tpu_initialize(FLAGS.tpu)
     strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)
     use_remote_tpu = True

official/nlp/xlnet/run_squad.py

@@ -30,7 +30,6 @@ from absl import logging
 
 import tensorflow as tf
 # pylint: disable=unused-import
-# Initialize TPU System.
 from official.nlp import xlnet_config
 from official.nlp import xlnet_modeling as modeling
 from official.nlp.xlnet import common_flags
@@ -38,6 +37,7 @@ from official.nlp.xlnet import data_utils
 from official.nlp.xlnet import optimization
 from official.nlp.xlnet import squad_utils
 from official.nlp.xlnet import training_utils
+from official.utils.misc import tpu_lib
 
 flags.DEFINE_string(
     "test_feature_path", default=None, help="Path to feature of test set.")
@@ -227,7 +227,7 @@ def main(unused_argv):
   if FLAGS.strategy_type == "mirror":
     strategy = tf.distribute.MirroredStrategy()
   elif FLAGS.strategy_type == "tpu":
-
+    # Initialize TPU System.
     cluster_resolver = tpu_lib.tpu_initialize(FLAGS.tpu)
     strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)
     use_remote_tpu = True