Internal change

PiperOrigin-RevId: 326286926

official/nlp/albert/run_squad.py

@@ -22,6 +22,7 @@ import json
 import os
 import time
 
+# Import libraries
 from absl import app
 from absl import flags
 from absl import logging

official/nlp/albert/tf2_albert_encoder_checkpoint_converter.py

@@ -23,6 +23,7 @@ from __future__ import print_function
 
 import os
 
+# Import libraries
 from absl import app
 from absl import flags
 

official/nlp/bert/bert_models_test.py

@@ -51,7 +51,9 @@ class BertModelsTest(tf.test.TestCase):
     self.assertIsInstance(encoder, networks.TransformerEncoder)
 
     # model has one scalar output: loss value.
-    self.assertEqual(model.output.shape.as_list(), [None,])
+    self.assertEqual(model.output.shape.as_list(), [
+        None,
+    ])
 
     # Expect two output from encoder: sequence and classification output.
     self.assertIsInstance(encoder.output, list)

official/nlp/bert/common_flags.py

@@ -73,9 +73,10 @@ def define_common_bert_flags():
       'If specified, init_checkpoint flag should not be used.')
   flags.DEFINE_bool('hub_module_trainable', True,
                     'True to make keras layers in the hub module trainable.')
-  flags.DEFINE_string('sub_model_export_name', None,
-                      'If set, `sub_model` checkpoints are exported into '
-                      'FLAGS.model_dir/FLAGS.sub_model_export_name.')
+  flags.DEFINE_string(
+      'sub_model_export_name', None,
+      'If set, `sub_model` checkpoints are exported into '
+      'FLAGS.model_dir/FLAGS.sub_model_export_name.')
 
   flags_core.define_log_steps()
 

official/nlp/bert/configs.py

@@ -20,6 +20,7 @@ from __future__ import print_function
 
 import copy
 import json
+
 import six
 import tensorflow as tf
 
@@ -105,4 +106,3 @@ class BertConfig(object):
   def to_json_string(self):
     """Serializes this instance to a JSON string."""
     return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
-

official/nlp/bert/export_tfhub.py

@@ -18,6 +18,7 @@ from __future__ import division
 # from __future__ import google_type_annotations
 from __future__ import print_function
 
+# Import libraries
 from absl import app
 from absl import flags
 from absl import logging

official/nlp/bert/export_tfhub_test.py

@@ -91,6 +91,7 @@ class ExportTfhubTest(tf.test.TestCase):
       outputs = np.concatenate(
           [hub_layer(inputs, training=training)[0] for _ in range(num_runs)])
       return np.mean(np.std(outputs, axis=0))
+
     self.assertLess(_dropout_mean_stddev(training=False), 1e-6)
     self.assertGreater(_dropout_mean_stddev(training=True), 1e-3)
 

official/nlp/bert/model_saving_utils.py

@@ -38,13 +38,13 @@ def export_bert_model(model_export_path: typing.Text,
       checkpoint_dir: Path from which model weights will be loaded, if
         specified.
       restore_model_using_load_weights: Whether to use checkpoint.restore() API
-        for custom checkpoint or to use model.load_weights() API.
-        There are 2 different ways to save checkpoints. One is using
-        tf.train.Checkpoint and another is using Keras model.save_weights().
-        Custom training loop implementation uses tf.train.Checkpoint API
-        and Keras ModelCheckpoint callback internally uses model.save_weights()
-        API. Since these two API's cannot be used toghether, model loading logic
-        must be take into account how model checkpoint was saved.
+        for custom checkpoint or to use model.load_weights() API. There are 2
+        different ways to save checkpoints. One is using tf.train.Checkpoint and
+        another is using Keras model.save_weights(). Custom training loop
+        implementation uses tf.train.Checkpoint API and Keras ModelCheckpoint
+        callback internally uses model.save_weights() API. Since these two API's
+        cannot be used toghether, model loading logic must be take into account
+        how model checkpoint was saved.
 
   Raises:
     ValueError when either model_export_path or model is not specified.

official/nlp/bert/model_training_utils.py

@@ -164,8 +164,8 @@ def run_customized_training_loop(
       custom_callbacks: A list of Keras Callbacks objects to run during
         training. More specifically, `on_train_begin(), on_train_end(),
         on_batch_begin()`, `on_batch_end()`, `on_epoch_begin()`,
-        `on_epoch_end()` methods are invoked during training.
-        Note that some metrics may be missing from `logs`.
+        `on_epoch_end()` methods are invoked during training. Note that some
+        metrics may be missing from `logs`.
       run_eagerly: Whether to run model training in pure eager execution. This
         should be disable for TPUStrategy.
       sub_model_export_name: If not None, will export `sub_model` returned by
@@ -458,8 +458,7 @@ def run_customized_training_loop(
     callback_list.on_train_begin()
     while current_step < total_training_steps and not model.stop_training:
       if current_step % steps_per_epoch == 0:
-        callback_list.on_epoch_begin(
-            int(current_step / steps_per_epoch) + 1)
+        callback_list.on_epoch_begin(int(current_step / steps_per_epoch) + 1)
 
       # Training loss/metric are taking average over steps inside micro
       # training loop. We reset the their values before each round.

official/nlp/bert/model_training_utils_test.py

@@ -139,9 +139,9 @@ class RecordingCallback(tf.keras.callbacks.Callback):
 
   def __init__(self):
     self.batch_begin = []  # (batch, logs)
-    self.batch_end = []    # (batch, logs)
+    self.batch_end = []  # (batch, logs)
     self.epoch_begin = []  # (epoch, logs)
-    self.epoch_end = []    # (epoch, logs)
+    self.epoch_end = []  # (epoch, logs)
 
   def on_batch_begin(self, batch, logs=None):
     self.batch_begin.append((batch, logs))
@@ -212,17 +212,19 @@ class ModelTrainingUtilsTest(tf.test.TestCase, parameterized.TestCase):
     # Two checkpoints should be saved after two epochs.
     files = map(os.path.basename,
                 tf.io.gfile.glob(os.path.join(model_dir, 'ctl_step_*index')))
-    self.assertCountEqual(['ctl_step_20.ckpt-1.index',
-                           'ctl_step_40.ckpt-2.index'], files)
+    self.assertCountEqual(
+        ['ctl_step_20.ckpt-1.index', 'ctl_step_40.ckpt-2.index'], files)
 
     # Three submodel checkpoints should be saved after two epochs (one after
     # each epoch plus one final).
-    files = map(os.path.basename,
-                tf.io.gfile.glob(os.path.join(model_dir,
-                                              'my_submodel_name*index')))
-    self.assertCountEqual(['my_submodel_name.ckpt-3.index',
-                           'my_submodel_name_step_20.ckpt-1.index',
-                           'my_submodel_name_step_40.ckpt-2.index'], files)
+    files = map(
+        os.path.basename,
+        tf.io.gfile.glob(os.path.join(model_dir, 'my_submodel_name*index')))
+    self.assertCountEqual([
+        'my_submodel_name.ckpt-3.index',
+        'my_submodel_name_step_20.ckpt-1.index',
+        'my_submodel_name_step_40.ckpt-2.index'
+    ], files)
 
     self.assertNotEmpty(
         tf.io.gfile.glob(

official/nlp/bert/run_classifier.py

@@ -22,6 +22,7 @@ import json
 import math
 import os
 
+# Import libraries
 from absl import app
 from absl import flags
 from absl import logging

official/nlp/bert/run_pretraining.py

@@ -17,6 +17,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+# Import libraries
 from absl import app
 from absl import flags
 from absl import logging

official/nlp/bert/run_squad.py

@@ -22,6 +22,7 @@ import json
 import os
 import time
 
+# Import libraries
 from absl import app
 from absl import flags
 from absl import logging

official/nlp/bert/run_squad_helper.py

@@ -20,6 +20,7 @@ from __future__ import print_function
 import collections
 import json
 import os
+
 from absl import flags
 from absl import logging
 import tensorflow as tf
@@ -39,10 +40,10 @@ from official.utils.misc import keras_utils
 def define_common_squad_flags():
   """Defines common flags used by SQuAD tasks."""
   flags.DEFINE_enum(
-      'mode', 'train_and_eval',
-      ['train_and_eval', 'train_and_predict',
-       'train', 'eval', 'predict', 'export_only'],
-      'One of {"train_and_eval", "train_and_predict", '
+      'mode', 'train_and_eval', [
+          'train_and_eval', 'train_and_predict', 'train', 'eval', 'predict',
+          'export_only'
+      ], 'One of {"train_and_eval", "train_and_predict", '
       '"train", "eval", "predict", "export_only"}. '
       '`train_and_eval`: train & predict to json files & compute eval metrics. '
       '`train_and_predict`: train & predict to json files. '
@@ -60,12 +61,12 @@ def define_common_squad_flags():
   # Model training specific flags.
   flags.DEFINE_integer('train_batch_size', 32, 'Total batch size for training.')
   # Predict processing related.
-  flags.DEFINE_string('predict_file', None,
-                      'SQuAD prediction json file path. '
-                      '`predict` mode supports multiple files: one can use '
-                      'wildcard to specify multiple files and it can also be '
-                      'multiple file patterns separated by comma. Note that '
-                      '`eval` mode only supports a single predict file.')
+  flags.DEFINE_string(
+      'predict_file', None, 'SQuAD prediction json file path. '
+      '`predict` mode supports multiple files: one can use '
+      'wildcard to specify multiple files and it can also be '
+      'multiple file patterns separated by comma. Note that '
+      '`eval` mode only supports a single predict file.')
   flags.DEFINE_bool(
       'do_lower_case', True,
       'Whether to lower case the input text. Should be True for uncased '
@@ -97,10 +98,7 @@ def define_common_squad_flags():
 FLAGS = flags.FLAGS
 
 
-def squad_loss_fn(start_positions,
-                  end_positions,
-                  start_logits,
-                  end_logits):
+def squad_loss_fn(start_positions, end_positions, start_logits, end_logits):
   """Returns sparse categorical crossentropy for start/end logits."""
   start_loss = tf.keras.losses.sparse_categorical_crossentropy(
       start_positions, start_logits, from_logits=True)
@@ -118,11 +116,8 @@ def get_loss_fn():
     start_positions = labels['start_positions']
     end_positions = labels['end_positions']
     start_logits, end_logits = model_outputs
-    return squad_loss_fn(
-        start_positions,
-        end_positions,
-        start_logits,
-        end_logits)
+    return squad_loss_fn(start_positions, end_positions, start_logits,
+                         end_logits)
 
   return _loss_fn
 
@@ -182,11 +177,8 @@ def get_squad_model_to_predict(strategy, bert_config, checkpoint_path,
   return squad_model
 
 
-def predict_squad_customized(strategy,
-                             input_meta_data,
-                             predict_tfrecord_path,
-                             num_steps,
-                             squad_model):
+def predict_squad_customized(strategy, input_meta_data, predict_tfrecord_path,
+                             num_steps, squad_model):
   """Make predictions using a Bert-based squad model."""
   predict_dataset_fn = get_dataset_fn(
       predict_tfrecord_path,
@@ -259,8 +251,7 @@ def train_squad(strategy,
         hub_module_trainable=FLAGS.hub_module_trainable)
     optimizer = optimization.create_optimizer(FLAGS.learning_rate,
                                               steps_per_epoch * epochs,
-                                              warmup_steps,
-                                              FLAGS.end_lr,
+                                              warmup_steps, FLAGS.end_lr,
                                               FLAGS.optimizer_type)
 
     squad_model.optimizer = performance.configure_optimizer(
@@ -344,8 +335,9 @@ def prediction_output_squad(strategy, input_meta_data, tokenizer, squad_lib,
   logging.info('  Batch size = %d', FLAGS.predict_batch_size)
 
   num_steps = int(dataset_size / FLAGS.predict_batch_size)
-  all_results = predict_squad_customized(
-      strategy, input_meta_data, eval_writer.filename, num_steps, squad_model)
+  all_results = predict_squad_customized(strategy, input_meta_data,
+                                         eval_writer.filename, num_steps,
+                                         squad_model)
 
   all_predictions, all_nbest_json, scores_diff_json = (
       squad_lib.postprocess_output(
@@ -362,8 +354,12 @@ def prediction_output_squad(strategy, input_meta_data, tokenizer, squad_lib,
   return all_predictions, all_nbest_json, scores_diff_json
 
 
-def dump_to_files(all_predictions, all_nbest_json, scores_diff_json,
-                  squad_lib, version_2_with_negative, file_prefix=''):
+def dump_to_files(all_predictions,
+                  all_nbest_json,
+                  scores_diff_json,
+                  squad_lib,
+                  version_2_with_negative,
+                  file_prefix=''):
   """Save output to json files."""
   output_prediction_file = os.path.join(FLAGS.model_dir,
                                         '%spredictions.json' % file_prefix)
@@ -452,8 +448,7 @@ def eval_squad(strategy,
     dataset_json = json.load(reader)
     pred_dataset = dataset_json['data']
   if input_meta_data.get('version_2_with_negative', False):
-    eval_metrics = squad_evaluate_v2_0.evaluate(pred_dataset,
-                                                all_predictions,
+    eval_metrics = squad_evaluate_v2_0.evaluate(pred_dataset, all_predictions,
                                                 scores_diff_json)
   else:
     eval_metrics = squad_evaluate_v1_1.evaluate(pred_dataset, all_predictions)

official/nlp/bert/serving.py

@@ -22,11 +22,11 @@ import tensorflow as tf
 from official.nlp.bert import bert_models
 from official.nlp.bert import configs
 
-flags.DEFINE_integer("sequence_length", None,
-                     "Sequence length to parse the tf.Example. If "
-                     "sequence_length > 0, add a signature for serialized "
-                     "tf.Example and define the parsing specification by the "
-                     "sequence_length.")
+flags.DEFINE_integer(
+    "sequence_length", None, "Sequence length to parse the tf.Example. If "
+    "sequence_length > 0, add a signature for serialized "
+    "tf.Example and define the parsing specification by the "
+    "sequence_length.")
 flags.DEFINE_string("bert_config_file", None,
                     "Bert configuration file to define core bert layers.")
 flags.DEFINE_string("model_checkpoint_path", None,

official/nlp/bert/squad_evaluate_v1_1.py

@@ -31,6 +31,7 @@ import re
 import string
 
 # pylint: disable=g-bad-import-order
+
 from absl import logging
 # pylint: enable=g-bad-import-order
 

official/nlp/bert/tf1_checkpoint_converter_lib.py

@@ -164,7 +164,6 @@ def convert(checkpoint_from_path,
         new_shape = _get_new_shape(new_var_name, tensor.shape, num_heads)
       if new_shape:
         tf.logging.info("Veriable %s has a shape change from %s to %s",
-
                         var_name, tensor.shape, new_shape)
         tensor = np.reshape(tensor, new_shape)
 

official/nlp/bert/tf2_encoder_checkpoint_converter.py

@@ -49,6 +49,7 @@ def _create_bert_model(cfg):
 
   Args:
     cfg: A `BertConfig` to create the core model.
+
   Returns:
     A TransformerEncoder netowork.
   """

official/nlp/data/create_finetuning_data.py

@@ -22,6 +22,7 @@ import functools
 import json
 import os
 
+# Import libraries
 from absl import app
 from absl import flags
 import tensorflow as tf

official/nlp/data/create_pretraining_data.py

@@ -21,6 +21,7 @@ import collections
 import itertools
 import random
 
+# Import libraries
 from absl import app
 from absl import flags
 from absl import logging