Move tf2_encoder_checkpoint_converter to public.

PiperOrigin-RevId: 283374562

official/nlp/bert/tf1_to_keras_checkpoint_converter.py

-# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-r"""Convert checkpoints created by Estimator (tf1) to be Keras compatible.
-
-Keras manages variable names internally, which results in subtly different names
-for variables between the Estimator and Keras version.
-The script should be used with TF 1.x.
-
-Usage:
-
-  python checkpoint_convert.py \
-      --checkpoint_from_path="/path/to/checkpoint" \
-      --checkpoint_to_path="/path/to/new_checkpoint"
-"""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl import app
-
-import tensorflow as tf  # TF 1.x
-from official.nlp.bert import tf1_checkpoint_converter_lib
-
-
-flags = tf.flags
-
-FLAGS = flags.FLAGS
-
-## Required parameters
-flags.DEFINE_string("checkpoint_from_path", None,
-                    "Source BERT checkpoint path.")
-flags.DEFINE_string("checkpoint_to_path", None,
-                    "Destination BERT checkpoint path.")
-flags.DEFINE_string(
-    "exclude_patterns", None,
-    "Comma-delimited string of a list of patterns to exclude"
-    " variables from source checkpoint.")
-flags.DEFINE_integer(
-    "num_heads", -1,
-    "The number of attention heads, used to reshape variables. If it is -1, "
-    "we do not reshape variables."
-)
-flags.DEFINE_boolean(
-    "create_v2_checkpoint", False,
-    "Whether to create a checkpoint compatible with KerasBERT V2 modeling code."
-)
-
-
-def main(_):
-  exclude_patterns = None
-  if FLAGS.exclude_patterns:
-    exclude_patterns = FLAGS.exclude_patterns.split(",")
-
-  if FLAGS.create_v2_checkpoint:
-    name_replacements = tf1_checkpoint_converter_lib.BERT_V2_NAME_REPLACEMENTS
-    permutations = tf1_checkpoint_converter_lib.BERT_V2_PERMUTATIONS
-  else:
-    name_replacements = tf1_checkpoint_converter_lib.BERT_NAME_REPLACEMENTS
-    permutations = tf1_checkpoint_converter_lib.BERT_PERMUTATIONS
-
-  tf1_checkpoint_converter_lib.convert(FLAGS.checkpoint_from_path,
-                                       FLAGS.checkpoint_to_path,
-                                       FLAGS.num_heads, name_replacements,
-                                       permutations, exclude_patterns)
-
-
-if __name__ == "__main__":
-  flags.mark_flag_as_required("checkpoint_from_path")
-  flags.mark_flag_as_required("checkpoint_to_path")
-  app.run(main)

official/nlp/bert/tf2_checkpoint_converter.py → official/nlp/bert/tf2_encoder_checkpoint_converter.py

@@ -12,83 +12,97 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""A converter for BERT name-based checkpoint to object-based checkpoint.
+"""A converter from a V1 BERT encoder checkpoint to a V2 encoder checkpoint.
 
-The conversion will yield objected-oriented checkpoint for TF2 Bert models,
-when BergConfig.backward_compatible is true.
-The variable/tensor shapes matches TF1 BERT model, but backward compatiblity
-introduces unnecessary reshape compuation.
+The conversion will yield an object-oriented checkpoint that can be used
+to restore a TransformerEncoder object.
 """
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import os
+
 from absl import app
 from absl import flags
 
-import tensorflow as tf  # TF 1.x
+import tensorflow as tf
+from official.modeling import activations
 from official.nlp import bert_modeling as modeling
+from official.nlp.bert import tf1_checkpoint_converter_lib
+from official.nlp.modeling import networks
 
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("bert_config_file", None,
                     "Bert configuration file to define core bert layers.")
 flags.DEFINE_string(
-    "init_checkpoint", None,
-    "Initial checkpoint (usually from a pre-trained BERT model).")
-flags.DEFINE_string("converted_checkpoint", None,
-                    "Path to objected-based V2 checkpoint.")
-flags.DEFINE_bool(
-    "export_bert_as_layer", False,
-    "Whether to use a layer rather than a model inside the checkpoint.")
+    "checkpoint_to_convert", None,
+    "Initial checkpoint from a pretrained BERT model core (that is, only the "
+    "BertModel, with no task heads.)")
+flags.DEFINE_string("converted_checkpoint_path", None,
+                    "Name for the created object-based V2 checkpoint.")
 
 
-def create_bert_model(bert_config):
+def _create_bert_model(cfg):
   """Creates a BERT keras core model from BERT configuration.
 
   Args:
-    bert_config: A BertConfig` to create the core model.
+    cfg: A `BertConfig` to create the core model.
   Returns:
     A keras model.
   """
-  max_seq_length = bert_config.max_position_embeddings
-
-  # Adds input layers just as placeholders.
-  input_word_ids = tf.keras.layers.Input(
-      shape=(max_seq_length,), dtype=tf.int32, name="input_word_ids")
-  input_mask = tf.keras.layers.Input(
-      shape=(max_seq_length,), dtype=tf.int32, name="input_mask")
-  input_type_ids = tf.keras.layers.Input(
-      shape=(max_seq_length,), dtype=tf.int32, name="input_type_ids")
-  core_model = modeling.get_bert_model(
-      input_word_ids,
-      input_mask,
-      input_type_ids,
-      config=bert_config,
-      name="bert_model",
-      float_type=tf.float32)
-  return core_model
-
-
-def convert_checkpoint():
-  """Converts a name-based matched TF V1 checkpoint to TF V2 checkpoint."""
-  bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
-  core_model = create_bert_model(bert_config)
-
-  # Uses streaming-restore in eager model to read V1 name-based checkpoints.
-  core_model.load_weights(FLAGS.init_checkpoint)
-  if FLAGS.export_bert_as_layer:
-    bert_layer = core_model.get_layer("bert_model")
-    checkpoint = tf.train.Checkpoint(bert_layer=bert_layer)
-  else:
-    checkpoint = tf.train.Checkpoint(model=core_model)
-
-  checkpoint.save(FLAGS.converted_checkpoint)
+  bert_encoder = networks.TransformerEncoder(
+      vocab_size=cfg.vocab_size,
+      hidden_size=cfg.hidden_size,
+      num_layers=cfg.num_hidden_layers,
+      num_attention_heads=cfg.num_attention_heads,
+      intermediate_size=cfg.intermediate_size,
+      activation=activations.gelu,
+      dropout_rate=cfg.hidden_dropout_prob,
+      attention_dropout_rate=cfg.attention_probs_dropout_prob,
+      sequence_length=cfg.max_position_embeddings,
+      type_vocab_size=cfg.type_vocab_size,
+      initializer=tf.keras.initializers.TruncatedNormal(
+          stddev=cfg.initializer_range))
+
+  return bert_encoder
+
+
+def convert_checkpoint(bert_config, output_path, v1_checkpoint):
+  """Converts a V1 checkpoint into an OO V2 checkpoint."""
+  output_dir, _ = os.path.split(output_path)
+
+  # Create a temporary V1 name-converted checkpoint in the output directory.
+  temporary_checkpoint_dir = os.path.join(output_dir, "temp_v1")
+  temporary_checkpoint = os.path.join(temporary_checkpoint_dir, "ckpt")
+  tf1_checkpoint_converter_lib.convert(
+      checkpoint_from_path=v1_checkpoint,
+      checkpoint_to_path=temporary_checkpoint,
+      num_heads=bert_config.num_attention_heads,
+      name_replacements=tf1_checkpoint_converter_lib.BERT_V2_NAME_REPLACEMENTS,
+      permutations=tf1_checkpoint_converter_lib.BERT_V2_PERMUTATIONS,
+      exclude_patterns=["adam", "Adam"])
+
+  # Create a V2 checkpoint from the temporary checkpoint.
+  model = _create_bert_model(bert_config)
+  tf1_checkpoint_converter_lib.create_v2_checkpoint(model, temporary_checkpoint,
+                                                    output_path)
+
+  # Clean up the temporary checkpoint, if it exists.
+  try:
+    tf.io.gfile.rmtree(temporary_checkpoint_dir)
+  except tf.errors.OpError:
+    # If it doesn't exist, we don't need to clean it up; continue.
+    pass
 
 
 def main(_):
-  tf.enable_eager_execution()
-  convert_checkpoint()
+  assert tf.version.VERSION.startswith('2.')
+  output_path = FLAGS.converted_checkpoint_path
+  v1_checkpoint = FLAGS.checkpoint_to_convert
+  bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
+  convert_checkpoint(bert_config, output_path, v1_checkpoint)
 
 
 if __name__ == "__main__":