Delete export_tfhub_lib_test.py

export_tfhub_lib_test.py

@@ -1,1080 +0,0 @@
-# Copyright 2024 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.
-
-"""Tests export_tfhub_lib."""
-
-import os
-import tempfile
-
-from absl.testing import parameterized
-import numpy as np
-import tensorflow as tf, tf_keras
-from tensorflow import estimator as tf_estimator
-import tensorflow_hub as hub
-import tensorflow_text as text
-
-from sentencepiece import SentencePieceTrainer
-from official.legacy.bert import configs
-from official.modeling import tf_utils
-from official.nlp.configs import encoders
-from official.nlp.modeling import layers
-from official.nlp.modeling import models
-from official.nlp.tools import export_tfhub_lib
-
-
-def _get_bert_config_or_encoder_config(use_bert_config,
-                                       hidden_size,
-                                       num_hidden_layers,
-                                       encoder_type="albert",
-                                       vocab_size=100):
-  """Generates config args for export_tfhub_lib._create_model().
-
-  Args:
-    use_bert_config: bool. If True, returns legacy BertConfig.
-    hidden_size: int.
-    num_hidden_layers: int.
-    encoder_type: str. Can be ['albert', 'bert', 'bert_v2']. If use_bert_config
-      == True, then model_type is not used.
-    vocab_size: int.
-
-  Returns:
-    bert_config, encoder_config. Only one is not None. If
-      `use_bert_config` == True, the first config is valid. Otherwise
-      `bert_config` == None.
-  """
-  if use_bert_config:
-    bert_config = configs.BertConfig(
-        vocab_size=vocab_size,
-        hidden_size=hidden_size,
-        intermediate_size=32,
-        max_position_embeddings=128,
-        num_attention_heads=2,
-        num_hidden_layers=num_hidden_layers)
-    encoder_config = None
-  else:
-    bert_config = None
-    if encoder_type == "albert":
-      encoder_config = encoders.EncoderConfig(
-          type="albert",
-          albert=encoders.AlbertEncoderConfig(
-              vocab_size=vocab_size,
-              embedding_width=16,
-              hidden_size=hidden_size,
-              intermediate_size=32,
-              max_position_embeddings=128,
-              num_attention_heads=2,
-              num_layers=num_hidden_layers,
-              dropout_rate=0.1))
-    else:
-      # encoder_type can be 'bert' or 'bert_v2'.
-      model_config = encoders.BertEncoderConfig(
-          vocab_size=vocab_size,
-          embedding_size=16,
-          hidden_size=hidden_size,
-          intermediate_size=32,
-          max_position_embeddings=128,
-          num_attention_heads=2,
-          num_layers=num_hidden_layers,
-          dropout_rate=0.1)
-      kwargs = {"type": encoder_type, encoder_type: model_config}
-      encoder_config = encoders.EncoderConfig(**kwargs)
-
-  return bert_config, encoder_config
-
-
-def _get_vocab_or_sp_model_dummy(temp_dir, use_sp_model):
-  """Returns tokenizer asset args for export_tfhub_lib.export_model()."""
-  dummy_file = os.path.join(temp_dir, "dummy_file.txt")
-  with tf.io.gfile.GFile(dummy_file, "w") as f:
-    f.write("dummy content")
-  if use_sp_model:
-    vocab_file, sp_model_file = None, dummy_file
-  else:
-    vocab_file, sp_model_file = dummy_file, None
-  return vocab_file, sp_model_file
-
-
-def _read_asset(asset: tf.saved_model.Asset):
-  return tf.io.gfile.GFile(asset.asset_path.numpy()).read()
-
-
-def _find_lambda_layers(layer):
-  """Returns list of all Lambda layers in a Keras model."""
-  if isinstance(layer, tf_keras.layers.Lambda):
-    return [layer]
-  elif hasattr(layer, "layers"):  # It's nested, like a Model.
-    result = []
-    for l in layer.layers:
-      result += _find_lambda_layers(l)
-    return result
-  else:
-    return []
-
-
-class ExportModelTest(tf.test.TestCase, parameterized.TestCase):
-  """Tests exporting a Transformer Encoder model as a SavedModel.
-
-  This covers export from an Encoder checkpoint to a SavedModel without
-  the .mlm subobject. This is no longer preferred, but still useful
-    for models like Electra that are trained without the MLM task.
-
-  The export code is generic. This test focuses on two main cases
-  (the most important ones in practice when this was written in 2020):
-    - BERT built from a legacy BertConfig, for use with BertTokenizer.
-    - ALBERT built from an EncoderConfig (as a representative of all other
-      choices beyond BERT, for use with SentencepieceTokenizer (the one
-      alternative to BertTokenizer).
-  """
-
-  @parameterized.named_parameters(
-      ("Bert_Legacy", True, None), ("Albert", False, "albert"),
-      ("BertEncoder", False, "bert"), ("BertEncoderV2", False, "bert_v2"))
-  def test_export_model(self, use_bert, encoder_type):
-    # Create the encoder and export it.
-    hidden_size = 16
-    num_hidden_layers = 1
-    bert_config, encoder_config = _get_bert_config_or_encoder_config(
-        use_bert,
-        hidden_size=hidden_size,
-        num_hidden_layers=num_hidden_layers,
-        encoder_type=encoder_type)
-    bert_model, encoder = export_tfhub_lib._create_model(
-        bert_config=bert_config, encoder_config=encoder_config, with_mlm=False)
-    self.assertEmpty(
-        _find_lambda_layers(bert_model),
-        "Lambda layers are non-portable since they serialize Python bytecode.")
-    model_checkpoint_dir = os.path.join(self.get_temp_dir(), "checkpoint")
-    checkpoint = tf.train.Checkpoint(encoder=encoder)
-    checkpoint.save(os.path.join(model_checkpoint_dir, "test"))
-    model_checkpoint_path = tf.train.latest_checkpoint(model_checkpoint_dir)
-
-    vocab_file, sp_model_file = _get_vocab_or_sp_model_dummy(
-        self.get_temp_dir(), use_sp_model=not use_bert)
-    export_path = os.path.join(self.get_temp_dir(), "hub")
-    export_tfhub_lib.export_model(
-        export_path=export_path,
-        bert_config=bert_config,
-        encoder_config=encoder_config,
-        model_checkpoint_path=model_checkpoint_path,
-        with_mlm=False,
-        vocab_file=vocab_file,
-        sp_model_file=sp_model_file,
-        do_lower_case=True)
-
-    # Restore the exported model.
-    hub_layer = hub.KerasLayer(export_path, trainable=True)
-
-    # Check legacy tokenization data.
-    if use_bert:
-      self.assertTrue(hub_layer.resolved_object.do_lower_case.numpy())
-      self.assertEqual("dummy content",
-                       _read_asset(hub_layer.resolved_object.vocab_file))
-      self.assertFalse(hasattr(hub_layer.resolved_object, "sp_model_file"))
-    else:
-      self.assertFalse(hasattr(hub_layer.resolved_object, "do_lower_case"))
-      self.assertFalse(hasattr(hub_layer.resolved_object, "vocab_file"))
-      self.assertEqual("dummy content",
-                       _read_asset(hub_layer.resolved_object.sp_model_file))
-
-    # Check restored weights.
-    self.assertEqual(
-        len(bert_model.trainable_weights), len(hub_layer.trainable_weights))
-    for source_weight, hub_weight in zip(bert_model.trainable_weights,
-                                         hub_layer.trainable_weights):
-      self.assertAllClose(source_weight.numpy(), hub_weight.numpy())
-
-    # Check computation.
-    seq_length = 10
-    dummy_ids = np.zeros((2, seq_length), dtype=np.int32)
-    input_dict = dict(
-        input_word_ids=dummy_ids,
-        input_mask=dummy_ids,
-        input_type_ids=dummy_ids)
-    hub_output = hub_layer(input_dict)
-    source_output = bert_model(input_dict)
-    encoder_output = encoder(input_dict)
-    self.assertEqual(hub_output["pooled_output"].shape, (2, hidden_size))
-    self.assertEqual(hub_output["sequence_output"].shape,
-                     (2, seq_length, hidden_size))
-    self.assertLen(hub_output["encoder_outputs"], num_hidden_layers)
-
-    for key in ("pooled_output", "sequence_output", "encoder_outputs"):
-      self.assertAllClose(source_output[key], hub_output[key])
-      self.assertAllClose(source_output[key], encoder_output[key])
-
-    # The "default" output of BERT as a text representation is pooled_output.
-    self.assertAllClose(hub_output["pooled_output"], hub_output["default"])
-
-    # Test that training=True makes a difference (activates dropout).
-    def _dropout_mean_stddev(training, num_runs=20):
-      input_ids = np.array([[14, 12, 42, 95, 99]], np.int32)
-      input_dict = dict(
-          input_word_ids=input_ids,
-          input_mask=np.ones_like(input_ids),
-          input_type_ids=np.zeros_like(input_ids))
-      outputs = np.concatenate([
-          hub_layer(input_dict, training=training)["pooled_output"]
-          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)
-
-    # Test propagation of seq_length in shape inference.
-    input_word_ids = tf_keras.layers.Input(shape=(seq_length,), dtype=tf.int32)
-    input_mask = tf_keras.layers.Input(shape=(seq_length,), dtype=tf.int32)
-    input_type_ids = tf_keras.layers.Input(shape=(seq_length,), dtype=tf.int32)
-    input_dict = dict(
-        input_word_ids=input_word_ids,
-        input_mask=input_mask,
-        input_type_ids=input_type_ids)
-    output_dict = hub_layer(input_dict)
-    pooled_output = output_dict["pooled_output"]
-    sequence_output = output_dict["sequence_output"]
-    encoder_outputs = output_dict["encoder_outputs"]
-
-    self.assertEqual(pooled_output.shape.as_list(), [None, hidden_size])
-    self.assertEqual(sequence_output.shape.as_list(),
-                     [None, seq_length, hidden_size])
-    self.assertLen(encoder_outputs, num_hidden_layers)
-
-
-class ExportModelWithMLMTest(tf.test.TestCase, parameterized.TestCase):
-  """Tests exporting a Transformer Encoder model as a SavedModel.
-
-  This covers export from a Pretrainer checkpoint to a SavedModel including
-  the .mlm subobject, which is the preferred way since 2020.
-
-  The export code is generic. This test focuses on two main cases
-  (the most important ones in practice when this was written in 2020):
-    - BERT built from a legacy BertConfig, for use with BertTokenizer.
-    - ALBERT built from an EncoderConfig (as a representative of all other
-      choices beyond BERT, for use with SentencepieceTokenizer (the one
-      alternative to BertTokenizer).
-  """
-
-  def test_copy_pooler_dense_to_encoder(self):
-    encoder_config = encoders.EncoderConfig(
-        type="bert",
-        bert=encoders.BertEncoderConfig(
-            hidden_size=24, intermediate_size=48, num_layers=2))
-    cls_heads = [
-        layers.ClassificationHead(
-            inner_dim=24, num_classes=2, name="next_sentence")
-    ]
-    encoder = encoders.build_encoder(encoder_config)
-    pretrainer = models.BertPretrainerV2(
-        encoder_network=encoder,
-        classification_heads=cls_heads,
-        mlm_activation=tf_utils.get_activation(
-            encoder_config.get().hidden_activation))
-    # Makes sure the pretrainer variables are created.
-    _ = pretrainer(pretrainer.inputs)
-    checkpoint = tf.train.Checkpoint(**pretrainer.checkpoint_items)
-    model_checkpoint_dir = os.path.join(self.get_temp_dir(), "checkpoint")
-    checkpoint.save(os.path.join(model_checkpoint_dir, "test"))
-
-    vocab_file, sp_model_file = _get_vocab_or_sp_model_dummy(
-        self.get_temp_dir(), use_sp_model=True)
-    export_path = os.path.join(self.get_temp_dir(), "hub")
-    export_tfhub_lib.export_model(
-        export_path=export_path,
-        encoder_config=encoder_config,
-        model_checkpoint_path=tf.train.latest_checkpoint(model_checkpoint_dir),
-        with_mlm=True,
-        copy_pooler_dense_to_encoder=True,
-        vocab_file=vocab_file,
-        sp_model_file=sp_model_file,
-        do_lower_case=True)
-    # Restores a hub KerasLayer.
-    hub_layer = hub.KerasLayer(export_path, trainable=True)
-    dummy_ids = np.zeros((2, 10), dtype=np.int32)
-    input_dict = dict(
-        input_word_ids=dummy_ids,
-        input_mask=dummy_ids,
-        input_type_ids=dummy_ids)
-    hub_pooled_output = hub_layer(input_dict)["pooled_output"]
-    encoder_outputs = encoder(input_dict)
-    # Verify that hub_layer's pooled_output is the same as the output of next
-    # sentence prediction's dense layer.
-    pretrained_pooled_output = cls_heads[0].dense(
-        (encoder_outputs["sequence_output"][:, 0, :]))
-    self.assertAllClose(hub_pooled_output, pretrained_pooled_output)
-    # But the pooled_output between encoder and hub_layer are not the same.
-    encoder_pooled_output = encoder_outputs["pooled_output"]
-    self.assertNotAllClose(hub_pooled_output, encoder_pooled_output)
-
-  @parameterized.named_parameters(
-      ("Bert", True),
-      ("Albert", False),
-  )
-  def test_export_model_with_mlm(self, use_bert):
-    # Create the encoder and export it.
-    hidden_size = 16
-    num_hidden_layers = 2
-    bert_config, encoder_config = _get_bert_config_or_encoder_config(
-        use_bert, hidden_size, num_hidden_layers)
-    bert_model, pretrainer = export_tfhub_lib._create_model(
-        bert_config=bert_config, encoder_config=encoder_config, with_mlm=True)
-    self.assertEmpty(
-        _find_lambda_layers(bert_model),
-        "Lambda layers are non-portable since they serialize Python bytecode.")
-    bert_model_with_mlm = bert_model.mlm
-    model_checkpoint_dir = os.path.join(self.get_temp_dir(), "checkpoint")
-
-    checkpoint = tf.train.Checkpoint(**pretrainer.checkpoint_items)
-
-    checkpoint.save(os.path.join(model_checkpoint_dir, "test"))
-    model_checkpoint_path = tf.train.latest_checkpoint(model_checkpoint_dir)
-
-    vocab_file, sp_model_file = _get_vocab_or_sp_model_dummy(
-        self.get_temp_dir(), use_sp_model=not use_bert)
-    export_path = os.path.join(self.get_temp_dir(), "hub")
-    export_tfhub_lib.export_model(
-        export_path=export_path,
-        bert_config=bert_config,
-        encoder_config=encoder_config,
-        model_checkpoint_path=model_checkpoint_path,
-        with_mlm=True,
-        vocab_file=vocab_file,
-        sp_model_file=sp_model_file,
-        do_lower_case=True)
-
-    # Restore the exported model.
-    hub_layer = hub.KerasLayer(export_path, trainable=True)
-
-    # Check legacy tokenization data.
-    if use_bert:
-      self.assertTrue(hub_layer.resolved_object.do_lower_case.numpy())
-      self.assertEqual("dummy content",
-                       _read_asset(hub_layer.resolved_object.vocab_file))
-      self.assertFalse(hasattr(hub_layer.resolved_object, "sp_model_file"))
-    else:
-      self.assertFalse(hasattr(hub_layer.resolved_object, "do_lower_case"))
-      self.assertFalse(hasattr(hub_layer.resolved_object, "vocab_file"))
-      self.assertEqual("dummy content",
-                       _read_asset(hub_layer.resolved_object.sp_model_file))
-
-    # Check restored weights.
-    # Note that we set `_auto_track_sub_layers` to False when exporting the
-    # SavedModel, so hub_layer has the same number of weights as bert_model;
-    # otherwise, hub_layer will have extra weights from its `mlm` subobject.
-    self.assertEqual(
-        len(bert_model.trainable_weights), len(hub_layer.trainable_weights))
-    for source_weight, hub_weight in zip(bert_model.trainable_weights,
-                                         hub_layer.trainable_weights):
-      self.assertAllClose(source_weight, hub_weight)
-
-    # Check computation.
-    seq_length = 10
-    dummy_ids = np.zeros((2, seq_length), dtype=np.int32)
-    input_dict = dict(
-        input_word_ids=dummy_ids,
-        input_mask=dummy_ids,
-        input_type_ids=dummy_ids)
-    hub_outputs_dict = hub_layer(input_dict)
-    source_outputs_dict = bert_model(input_dict)
-    encoder_outputs_dict = pretrainer.encoder_network(
-        [dummy_ids, dummy_ids, dummy_ids])
-    self.assertEqual(hub_outputs_dict["pooled_output"].shape, (2, hidden_size))
-    self.assertEqual(hub_outputs_dict["sequence_output"].shape,
-                     (2, seq_length, hidden_size))
-    for output_key in ("pooled_output", "sequence_output", "encoder_outputs"):
-      self.assertAllClose(source_outputs_dict[output_key],
-                          hub_outputs_dict[output_key])
-      self.assertAllClose(source_outputs_dict[output_key],
-                          encoder_outputs_dict[output_key])
-
-    # The "default" output of BERT as a text representation is pooled_output.
-    self.assertAllClose(hub_outputs_dict["pooled_output"],
-                        hub_outputs_dict["default"])
-
-    # Test that training=True makes a difference (activates dropout).
-    def _dropout_mean_stddev(training, num_runs=20):
-      input_ids = np.array([[14, 12, 42, 95, 99]], np.int32)
-      input_dict = dict(
-          input_word_ids=input_ids,
-          input_mask=np.ones_like(input_ids),
-          input_type_ids=np.zeros_like(input_ids))
-      outputs = np.concatenate([
-          hub_layer(input_dict, training=training)["pooled_output"]
-          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)
-
-    # Checks sub-object `mlm`.
-    self.assertTrue(hasattr(hub_layer.resolved_object, "mlm"))
-
-    self.assertLen(hub_layer.resolved_object.mlm.trainable_variables,
-                   len(bert_model_with_mlm.trainable_weights))
-    self.assertLen(hub_layer.resolved_object.mlm.trainable_variables,
-                   len(pretrainer.trainable_weights))
-    for source_weight, hub_weight, pretrainer_weight in zip(
-        bert_model_with_mlm.trainable_weights,
-        hub_layer.resolved_object.mlm.trainable_variables,
-        pretrainer.trainable_weights):
-      self.assertAllClose(source_weight, hub_weight)
-      self.assertAllClose(source_weight, pretrainer_weight)
-
-    max_predictions_per_seq = 4
-    mlm_positions = np.zeros((2, max_predictions_per_seq), dtype=np.int32)
-    input_dict = dict(
-        input_word_ids=dummy_ids,
-        input_mask=dummy_ids,
-        input_type_ids=dummy_ids,
-        masked_lm_positions=mlm_positions)
-    hub_mlm_outputs_dict = hub_layer.resolved_object.mlm(input_dict)
-    source_mlm_outputs_dict = bert_model_with_mlm(input_dict)
-    for output_key in ("pooled_output", "sequence_output", "mlm_logits",
-                       "encoder_outputs"):
-      self.assertAllClose(hub_mlm_outputs_dict[output_key],
-                          source_mlm_outputs_dict[output_key])
-
-    pretrainer_mlm_logits_output = pretrainer(input_dict)["mlm_logits"]
-    self.assertAllClose(hub_mlm_outputs_dict["mlm_logits"],
-                        pretrainer_mlm_logits_output)
-
-    # Test that training=True makes a difference (activates dropout).
-    def _dropout_mean_stddev_mlm(training, num_runs=20):
-      input_ids = np.array([[14, 12, 42, 95, 99]], np.int32)
-      mlm_position_ids = np.array([[1, 2, 3, 4]], np.int32)
-      input_dict = dict(
-          input_word_ids=input_ids,
-          input_mask=np.ones_like(input_ids),
-          input_type_ids=np.zeros_like(input_ids),
-          masked_lm_positions=mlm_position_ids)
-      outputs = np.concatenate([
-          hub_layer.resolved_object.mlm(input_dict,
-                                        training=training)["pooled_output"]
-          for _ in range(num_runs)
-      ])
-      return np.mean(np.std(outputs, axis=0))
-
-    self.assertLess(_dropout_mean_stddev_mlm(training=False), 1e-6)
-    self.assertGreater(_dropout_mean_stddev_mlm(training=True), 1e-3)
-
-    # Test propagation of seq_length in shape inference.
-    input_word_ids = tf_keras.layers.Input(shape=(seq_length,), dtype=tf.int32)
-    input_mask = tf_keras.layers.Input(shape=(seq_length,), dtype=tf.int32)
-    input_type_ids = tf_keras.layers.Input(shape=(seq_length,), dtype=tf.int32)
-    input_dict = dict(
-        input_word_ids=input_word_ids,
-        input_mask=input_mask,
-        input_type_ids=input_type_ids)
-    hub_outputs_dict = hub_layer(input_dict)
-    self.assertEqual(hub_outputs_dict["pooled_output"].shape.as_list(),
-                     [None, hidden_size])
-    self.assertEqual(hub_outputs_dict["sequence_output"].shape.as_list(),
-                     [None, seq_length, hidden_size])
-
-
-_STRING_NOT_TO_LEAK = "private_path_component_"
-
-
-class ExportPreprocessingTest(tf.test.TestCase, parameterized.TestCase):
-
-  def _make_vocab_file(self, vocab, filename="vocab.txt", add_mask_token=False):
-    """Creates wordpiece vocab file with given words plus special tokens.
-
-    The tokens of the resulting model are, in this order:
-        [PAD], [UNK], [CLS], [SEP], [MASK]*, ...vocab...
-    *=if requested by args.
-
-    This function also accepts wordpieces that start with the ## continuation
-    marker, but avoiding those makes this function interchangeable with
-    _make_sp_model_file(), up to the extra dimension returned by BertTokenizer.
-
-    Args:
-      vocab: a list of strings with the words or wordpieces to put into the
-        model's vocabulary. Do not include special tokens here.
-      filename: Optionally, a filename (relative to the temporary directory
-        created by this function).
-      add_mask_token: an optional bool, whether to include a [MASK] token.
-
-    Returns:
-      The absolute filename of the created vocab file.
-    """
-    full_vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]"
-                 ] + ["[MASK]"] * add_mask_token + vocab
-    path = os.path.join(
-        tempfile.mkdtemp(
-            dir=self.get_temp_dir(),  # New subdir each time.
-            prefix=_STRING_NOT_TO_LEAK),
-        filename)
-    with tf.io.gfile.GFile(path, "w") as f:
-      f.write("\n".join(full_vocab + [""]))
-    return path
-
-  def _make_sp_model_file(self, vocab, prefix="spm", add_mask_token=False):
-    """Creates Sentencepiece word model with given words plus special tokens.
-
-    The tokens of the resulting model are, in this order:
-        <pad>, <unk>, [CLS], [SEP], [MASK]*, ...vocab..., <s>, </s>
-    *=if requested by args.
-
-    The words in the input vocab are plain text, without the whitespace marker.
-    That makes this function interchangeable with _make_vocab_file().
-
-    Args:
-      vocab: a list of strings with the words to put into the model's
-        vocabulary. Do not include special tokens here.
-      prefix: an optional string, to change the filename prefix for the model
-        (relative to the temporary directory created by this function).
-      add_mask_token: an optional bool, whether to include a [MASK] token.
-
-    Returns:
-      The absolute filename of the created Sentencepiece model file.
-    """
-    model_prefix = os.path.join(
-        tempfile.mkdtemp(dir=self.get_temp_dir()),  # New subdir each time.
-        prefix)
-    input_file = model_prefix + "_train_input.txt"
-    # Create input text for training the sp model from the tokens provided.
-    # Repeat tokens, the earlier the more, because they are sorted by frequency.
-    input_text = []
-    for i, token in enumerate(vocab):
-      input_text.append(" ".join([token] * (len(vocab) - i)))
-    with tf.io.gfile.GFile(input_file, "w") as f:
-      f.write("\n".join(input_text + [""]))
-    control_symbols = "[CLS],[SEP]"
-    full_vocab_size = len(vocab) + 6  # <pad>, <unk>, [CLS], [SEP], <s>, </s>.
-    if add_mask_token:
-      control_symbols += ",[MASK]"
-      full_vocab_size += 1
-    flags = dict(
-        model_prefix=model_prefix,
-        model_type="word",
-        input=input_file,
-        pad_id=0,
-        unk_id=1,
-        control_symbols=control_symbols,
-        vocab_size=full_vocab_size,
-        bos_id=full_vocab_size - 2,
-        eos_id=full_vocab_size - 1)
-    SentencePieceTrainer.Train(" ".join(
-        ["--{}={}".format(k, v) for k, v in flags.items()]))
-    return model_prefix + ".model"
-
-  def _do_export(self,
-                 vocab,
-                 do_lower_case,
-                 default_seq_length=128,
-                 tokenize_with_offsets=True,
-                 use_sp_model=False,
-                 experimental_disable_assert=False,
-                 add_mask_token=False):
-    """Runs SavedModel export and returns the export_path."""
-    export_path = tempfile.mkdtemp(dir=self.get_temp_dir())
-    vocab_file = sp_model_file = None
-    if use_sp_model:
-      sp_model_file = self._make_sp_model_file(
-          vocab, add_mask_token=add_mask_token)
-    else:
-      vocab_file = self._make_vocab_file(vocab, add_mask_token=add_mask_token)
-    export_tfhub_lib.export_preprocessing(
-        export_path,
-        vocab_file=vocab_file,
-        sp_model_file=sp_model_file,
-        do_lower_case=do_lower_case,
-        tokenize_with_offsets=tokenize_with_offsets,
-        default_seq_length=default_seq_length,
-        experimental_disable_assert=experimental_disable_assert)
-    # Invalidate the original filename to verify loading from the SavedModel.
-    tf.io.gfile.remove(sp_model_file or vocab_file)
-    return export_path
-
-  def test_no_leaks(self):
-    """Tests not leaking the path to the original vocab file."""
-    path = self._do_export(["d", "ef", "abc", "xy"],
-                           do_lower_case=True,
-                           use_sp_model=False)
-    with tf.io.gfile.GFile(os.path.join(path, "saved_model.pb"), "rb") as f:
-      self.assertFalse(  # pylint: disable=g-generic-assert
-          _STRING_NOT_TO_LEAK.encode("ascii") in f.read())
-
-  @parameterized.named_parameters(("Bert", False), ("Sentencepiece", True))
-  def test_exported_callables(self, use_sp_model):
-    preprocess = tf.saved_model.load(
-        self._do_export(
-            ["d", "ef", "abc", "xy"],
-            do_lower_case=True,
-            # TODO(b/181866850): drop this.
-            tokenize_with_offsets=not use_sp_model,
-            # TODO(b/175369555): drop this.
-            experimental_disable_assert=True,
-            use_sp_model=use_sp_model))
-
-    def fold_dim(rt):
-      """Removes the word/subword distinction of BertTokenizer."""
-      return rt if use_sp_model else rt.merge_dims(1, 2)
-
-    # .tokenize()
-    inputs = tf.constant(["abc d ef", "ABC D EF d"])
-    token_ids = preprocess.tokenize(inputs)
-    self.assertAllEqual(
-        fold_dim(token_ids), tf.ragged.constant([[6, 4, 5], [6, 4, 5, 4]]))
-
-    special_tokens_dict = {
-        k: v.numpy().item()  # Expecting eager Tensor, converting to Python.
-        for k, v in preprocess.tokenize.get_special_tokens_dict().items()
-    }
-    self.assertDictEqual(
-        special_tokens_dict,
-        dict(
-            padding_id=0,
-            start_of_sequence_id=2,
-            end_of_segment_id=3,
-            vocab_size=4 + 6 if use_sp_model else 4 + 4))
-
-    # .tokenize_with_offsets()
-    if use_sp_model:
-      # TODO(b/181866850): Enable tokenize_with_offsets when it works and test.
-      self.assertFalse(hasattr(preprocess, "tokenize_with_offsets"))
-    else:
-      token_ids, start_offsets, limit_offsets = (
-          preprocess.tokenize_with_offsets(inputs))
-      self.assertAllEqual(
-          fold_dim(token_ids), tf.ragged.constant([[6, 4, 5], [6, 4, 5, 4]]))
-      self.assertAllEqual(
-          fold_dim(start_offsets), tf.ragged.constant([[0, 4, 6], [0, 4, 6,
-                                                                   9]]))
-      self.assertAllEqual(
-          fold_dim(limit_offsets), tf.ragged.constant([[3, 5, 8], [3, 5, 8,
-                                                                   10]]))
-      self.assertIs(preprocess.tokenize.get_special_tokens_dict,
-                    preprocess.tokenize_with_offsets.get_special_tokens_dict)
-
-    # Root callable.
-    bert_inputs = preprocess(inputs)
-    self.assertAllEqual(bert_inputs["input_word_ids"].shape.as_list(), [2, 128])
-    self.assertAllEqual(
-        bert_inputs["input_word_ids"][:, :10],
-        tf.constant([[2, 6, 4, 5, 3, 0, 0, 0, 0, 0],
-                     [2, 6, 4, 5, 4, 3, 0, 0, 0, 0]]))
-    self.assertAllEqual(bert_inputs["input_mask"].shape.as_list(), [2, 128])
-    self.assertAllEqual(
-        bert_inputs["input_mask"][:, :10],
-        tf.constant([[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
-                     [1, 1, 1, 1, 1, 1, 0, 0, 0, 0]]))
-    self.assertAllEqual(bert_inputs["input_type_ids"].shape.as_list(), [2, 128])
-    self.assertAllEqual(
-        bert_inputs["input_type_ids"][:, :10],
-        tf.constant([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-                     [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]))
-
-    # .bert_pack_inputs()
-    inputs_2 = tf.constant(["d xy", "xy abc"])
-    token_ids_2 = preprocess.tokenize(inputs_2)
-    bert_inputs = preprocess.bert_pack_inputs([token_ids, token_ids_2],
-                                              seq_length=256)
-    self.assertAllEqual(bert_inputs["input_word_ids"].shape.as_list(), [2, 256])
-    self.assertAllEqual(
-        bert_inputs["input_word_ids"][:, :10],
-        tf.constant([[2, 6, 4, 5, 3, 4, 7, 3, 0, 0],
-                     [2, 6, 4, 5, 4, 3, 7, 6, 3, 0]]))
-    self.assertAllEqual(bert_inputs["input_mask"].shape.as_list(), [2, 256])
-    self.assertAllEqual(
-        bert_inputs["input_mask"][:, :10],
-        tf.constant([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
-                     [1, 1, 1, 1, 1, 1, 1, 1, 1, 0]]))
-    self.assertAllEqual(bert_inputs["input_type_ids"].shape.as_list(), [2, 256])
-    self.assertAllEqual(
-        bert_inputs["input_type_ids"][:, :10],
-        tf.constant([[0, 0, 0, 0, 0, 1, 1, 1, 0, 0],
-                     [0, 0, 0, 0, 0, 0, 1, 1, 1, 0]]))
-
-  # For BertTokenizer only: repeat relevant parts for do_lower_case=False,
-  # default_seq_length=10, experimental_disable_assert=False,
-  # tokenize_with_offsets=False, and without folding the word/subword dimension.
-  def test_cased_length10(self):
-    preprocess = tf.saved_model.load(
-        self._do_export(["d", "##ef", "abc", "ABC"],
-                        do_lower_case=False,
-                        default_seq_length=10,
-                        tokenize_with_offsets=False,
-                        use_sp_model=False,
-                        experimental_disable_assert=False))
-    inputs = tf.constant(["abc def", "ABC DEF"])
-    token_ids = preprocess.tokenize(inputs)
-    self.assertAllEqual(token_ids,
-                        tf.ragged.constant([[[6], [4, 5]], [[7], [1]]]))
-
-    self.assertFalse(hasattr(preprocess, "tokenize_with_offsets"))
-
-    bert_inputs = preprocess(inputs)
-    self.assertAllEqual(
-        bert_inputs["input_word_ids"],
-        tf.constant([[2, 6, 4, 5, 3, 0, 0, 0, 0, 0],
-                     [2, 7, 1, 3, 0, 0, 0, 0, 0, 0]]))
-    self.assertAllEqual(
-        bert_inputs["input_mask"],
-        tf.constant([[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
-                     [1, 1, 1, 1, 0, 0, 0, 0, 0, 0]]))
-    self.assertAllEqual(
-        bert_inputs["input_type_ids"],
-        tf.constant([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-                     [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]))
-
-    inputs_2 = tf.constant(["d ABC", "ABC abc"])
-    token_ids_2 = preprocess.tokenize(inputs_2)
-    bert_inputs = preprocess.bert_pack_inputs([token_ids, token_ids_2])
-    # Test default seq_length=10.
-    self.assertAllEqual(
-        bert_inputs["input_word_ids"],
-        tf.constant([[2, 6, 4, 5, 3, 4, 7, 3, 0, 0],
-                     [2, 7, 1, 3, 7, 6, 3, 0, 0, 0]]))
-    self.assertAllEqual(
-        bert_inputs["input_mask"],
-        tf.constant([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
-                     [1, 1, 1, 1, 1, 1, 1, 0, 0, 0]]))
-    self.assertAllEqual(
-        bert_inputs["input_type_ids"],
-        tf.constant([[0, 0, 0, 0, 0, 1, 1, 1, 0, 0],
-                     [0, 0, 0, 0, 1, 1, 1, 0, 0, 0]]))
-
-  # XLA requires fixed shapes for tensors found in graph mode.
-  # Statically known shapes in Python are a particularly firm way to
-  # guarantee that, and they are generally more convenient to work with.
-  # We test that the exported SavedModel plays well with TF's shape
-  # inference when applied to fully or partially known input shapes.
-  @parameterized.named_parameters(("Bert", False), ("Sentencepiece", True))
-  def test_shapes(self, use_sp_model):
-    preprocess = tf.saved_model.load(
-        self._do_export(
-            ["abc", "def"],
-            do_lower_case=True,
-            # TODO(b/181866850): drop this.
-            tokenize_with_offsets=not use_sp_model,
-            # TODO(b/175369555): drop this.
-            experimental_disable_assert=True,
-            use_sp_model=use_sp_model))
-
-    def expected_bert_input_shapes(batch_size, seq_length):
-      return dict(
-          input_word_ids=[batch_size, seq_length],
-          input_mask=[batch_size, seq_length],
-          input_type_ids=[batch_size, seq_length])
-
-    for batch_size in [7, None]:
-      if use_sp_model:
-        token_out_shape = [batch_size, None]  # No word/subword distinction.
-      else:
-        token_out_shape = [batch_size, None, None]
-      self.assertEqual(
-          _result_shapes_in_tf_function(preprocess.tokenize,
-                                        tf.TensorSpec([batch_size], tf.string)),
-          token_out_shape, "with batch_size=%s" % batch_size)
-      # TODO(b/181866850): Enable tokenize_with_offsets when it works and test.
-      if use_sp_model:
-        self.assertFalse(hasattr(preprocess, "tokenize_with_offsets"))
-      else:
-        self.assertEqual(
-            _result_shapes_in_tf_function(
-                preprocess.tokenize_with_offsets,
-                tf.TensorSpec([batch_size], tf.string)), [token_out_shape] * 3,
-            "with batch_size=%s" % batch_size)
-      self.assertEqual(
-          _result_shapes_in_tf_function(
-              preprocess.bert_pack_inputs,
-              [tf.RaggedTensorSpec([batch_size, None, None], tf.int32)] * 2,
-              seq_length=256), expected_bert_input_shapes(batch_size, 256),
-          "with batch_size=%s" % batch_size)
-      self.assertEqual(
-          _result_shapes_in_tf_function(preprocess,
-                                        tf.TensorSpec([batch_size], tf.string)),
-          expected_bert_input_shapes(batch_size, 128),
-          "with batch_size=%s" % batch_size)
-
-  @parameterized.named_parameters(("Bert", False), ("Sentencepiece", True))
-  def test_reexport(self, use_sp_model):
-    """Test that preprocess keeps working after another save/load cycle."""
-    path1 = self._do_export(
-        ["d", "ef", "abc", "xy"],
-        do_lower_case=True,
-        default_seq_length=10,
-        tokenize_with_offsets=False,
-        experimental_disable_assert=True,  # TODO(b/175369555): drop this.
-        use_sp_model=use_sp_model)
-    path2 = path1.rstrip("/") + ".2"
-    model1 = tf.saved_model.load(path1)
-    tf.saved_model.save(model1, path2)
-    # Delete the first SavedModel to test that the sceond one loads by itself.
-    # https://github.com/tensorflow/tensorflow/issues/46456 reports such a
-    # failure case for BertTokenizer.
-    tf.io.gfile.rmtree(path1)
-    model2 = tf.saved_model.load(path2)
-
-    inputs = tf.constant(["abc d ef", "ABC D EF d"])
-    bert_inputs = model2(inputs)
-    self.assertAllEqual(
-        bert_inputs["input_word_ids"],
-        tf.constant([[2, 6, 4, 5, 3, 0, 0, 0, 0, 0],
-                     [2, 6, 4, 5, 4, 3, 0, 0, 0, 0]]))
-    self.assertAllEqual(
-        bert_inputs["input_mask"],
-        tf.constant([[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
-                     [1, 1, 1, 1, 1, 1, 0, 0, 0, 0]]))
-    self.assertAllEqual(
-        bert_inputs["input_type_ids"],
-        tf.constant([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-                     [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]))
-
-  @parameterized.named_parameters(("Bert", True), ("Albert", False))
-  def test_preprocessing_for_mlm(self, use_bert):
-    """Combines both SavedModel types and TF.text helpers for MLM."""
-    # Create the preprocessing SavedModel with a [MASK] token.
-    non_special_tokens = [
-        "hello", "world", "nice", "movie", "great", "actors", "quick", "fox",
-        "lazy", "dog"
-    ]
-
-    preprocess = tf.saved_model.load(
-        self._do_export(
-            non_special_tokens,
-            do_lower_case=True,
-            tokenize_with_offsets=use_bert,  # TODO(b/181866850): drop this.
-            experimental_disable_assert=True,  # TODO(b/175369555): drop this.
-            add_mask_token=True,
-            use_sp_model=not use_bert))
-    vocab_size = len(non_special_tokens) + (5 if use_bert else 7)
-
-    # Create the encoder SavedModel with an .mlm subobject.
-    hidden_size = 16
-    num_hidden_layers = 2
-    bert_config, encoder_config = _get_bert_config_or_encoder_config(
-        use_bert_config=use_bert,
-        hidden_size=hidden_size,
-        num_hidden_layers=num_hidden_layers,
-        vocab_size=vocab_size)
-    _, pretrainer = export_tfhub_lib._create_model(
-        bert_config=bert_config, encoder_config=encoder_config, with_mlm=True)
-    model_checkpoint_dir = os.path.join(self.get_temp_dir(), "checkpoint")
-    checkpoint = tf.train.Checkpoint(**pretrainer.checkpoint_items)
-    checkpoint.save(os.path.join(model_checkpoint_dir, "test"))
-    model_checkpoint_path = tf.train.latest_checkpoint(model_checkpoint_dir)
-    vocab_file, sp_model_file = _get_vocab_or_sp_model_dummy(  # Not used below.
-        self.get_temp_dir(), use_sp_model=not use_bert)
-    encoder_export_path = os.path.join(self.get_temp_dir(), "encoder_export")
-    export_tfhub_lib.export_model(
-        export_path=encoder_export_path,
-        bert_config=bert_config,
-        encoder_config=encoder_config,
-        model_checkpoint_path=model_checkpoint_path,
-        with_mlm=True,
-        vocab_file=vocab_file,
-        sp_model_file=sp_model_file,
-        do_lower_case=True)
-    encoder = tf.saved_model.load(encoder_export_path)
-
-    # Get special tokens from the vocab (and vocab size).
-    special_tokens_dict = preprocess.tokenize.get_special_tokens_dict()
-    self.assertEqual(int(special_tokens_dict["vocab_size"]), vocab_size)
-    padding_id = int(special_tokens_dict["padding_id"])
-    self.assertEqual(padding_id, 0)
-    start_of_sequence_id = int(special_tokens_dict["start_of_sequence_id"])
-    self.assertEqual(start_of_sequence_id, 2)
-    end_of_segment_id = int(special_tokens_dict["end_of_segment_id"])
-    self.assertEqual(end_of_segment_id, 3)
-    mask_id = int(special_tokens_dict["mask_id"])
-    self.assertEqual(mask_id, 4)
-
-    # A batch of 3 segment pairs.
-    raw_segments = [
-        tf.constant(["hello", "nice movie", "quick fox"]),
-        tf.constant(["world", "great actors", "lazy dog"])
-    ]
-    batch_size = 3
-
-    # Misc hyperparameters.
-    seq_length = 10
-    max_selections_per_seq = 2
-
-    # Tokenize inputs.
-    tokenized_segments = [preprocess.tokenize(s) for s in raw_segments]
-    # Trim inputs to eventually fit seq_lentgh.
-    num_special_tokens = len(raw_segments) + 1
-    trimmed_segments = text.WaterfallTrimmer(
-        seq_length - num_special_tokens).trim(tokenized_segments)
-    # Combine input segments into one input sequence.
-    input_ids, segment_ids = text.combine_segments(
-        trimmed_segments,
-        start_of_sequence_id=start_of_sequence_id,
-        end_of_segment_id=end_of_segment_id)
-    # Apply random masking controlled by policy objects.
-    (masked_input_ids, masked_lm_positions,
-     masked_ids) = text.mask_language_model(
-         input_ids=input_ids,
-         item_selector=text.RandomItemSelector(
-             max_selections_per_seq,
-             selection_rate=0.5,  # Adjusted for the short test examples.
-             unselectable_ids=[start_of_sequence_id, end_of_segment_id]),
-         mask_values_chooser=text.MaskValuesChooser(
-             vocab_size=vocab_size,
-             mask_token=mask_id,
-             # Always put [MASK] to have a predictable result.
-             mask_token_rate=1.0,
-             random_token_rate=0.0))
-    # Pad to fixed-length Transformer encoder inputs.
-    input_word_ids, _ = text.pad_model_inputs(
-        masked_input_ids, seq_length, pad_value=padding_id)
-    input_type_ids, input_mask = text.pad_model_inputs(
-        segment_ids, seq_length, pad_value=0)
-    masked_lm_positions, _ = text.pad_model_inputs(
-        masked_lm_positions, max_selections_per_seq, pad_value=0)
-    masked_lm_positions = tf.cast(masked_lm_positions, tf.int32)
-    num_predictions = int(tf.shape(masked_lm_positions)[1])
-
-    # Test transformer inputs.
-    self.assertEqual(num_predictions, max_selections_per_seq)
-    expected_word_ids = np.array([
-        # [CLS] hello [SEP] world [SEP]
-        [2, 5, 3, 6, 3, 0, 0, 0, 0, 0],
-        # [CLS] nice movie [SEP] great actors [SEP]
-        [2, 7, 8, 3, 9, 10, 3, 0, 0, 0],
-        # [CLS] brown fox [SEP] lazy dog [SEP]
-        [2, 11, 12, 3, 13, 14, 3, 0, 0, 0]
-    ])
-    for i in range(batch_size):
-      for j in range(num_predictions):
-        k = int(masked_lm_positions[i, j])
-        if k != 0:
-          expected_word_ids[i, k] = 4  # [MASK]
-    self.assertAllEqual(input_word_ids, expected_word_ids)
-
-    # Call the MLM head of the Transformer encoder.
-    mlm_inputs = dict(
-        input_word_ids=input_word_ids,
-        input_mask=input_mask,
-        input_type_ids=input_type_ids,
-        masked_lm_positions=masked_lm_positions,
-    )
-    mlm_outputs = encoder.mlm(mlm_inputs)
-    self.assertEqual(mlm_outputs["pooled_output"].shape,
-                     (batch_size, hidden_size))
-    self.assertEqual(mlm_outputs["sequence_output"].shape,
-                     (batch_size, seq_length, hidden_size))
-    self.assertEqual(mlm_outputs["mlm_logits"].shape,
-                     (batch_size, num_predictions, vocab_size))
-    self.assertLen(mlm_outputs["encoder_outputs"], num_hidden_layers)
-
-    # A real trainer would now compute the loss of mlm_logits
-    # trying to predict the masked_ids.
-    del masked_ids  # Unused.
-
-  @parameterized.named_parameters(("Bert", False), ("Sentencepiece", True))
-  def test_special_tokens_in_estimator(self, use_sp_model):
-    """Tests getting special tokens without an Eager init context."""
-    preprocess_export_path = self._do_export(["d", "ef", "abc", "xy"],
-                                             do_lower_case=True,
-                                             use_sp_model=use_sp_model,
-                                             tokenize_with_offsets=False)
-
-    def _get_special_tokens_dict(obj):
-      """Returns special tokens of restored tokenizer as Python values."""
-      if tf.executing_eagerly():
-        special_tokens_numpy = {
-            k: v.numpy() for k, v in obj.get_special_tokens_dict()
-        }
-      else:
-        with tf.Graph().as_default():
-          # This code expects `get_special_tokens_dict()` to be a tf.function
-          # with no dependencies (bound args) from the context it was loaded in,
-          # and boldly assumes that it can just be called in a dfferent context.
-          special_tokens_tensors = obj.get_special_tokens_dict()
-          with tf.compat.v1.Session() as sess:
-            special_tokens_numpy = sess.run(special_tokens_tensors)
-      return {
-          k: v.item()  # Numpy to Python.
-          for k, v in special_tokens_numpy.items()
-      }
-
-    def input_fn():
-      self.assertFalse(tf.executing_eagerly())
-      # Build a preprocessing Model.
-      sentences = tf_keras.layers.Input(shape=[], dtype=tf.string)
-      preprocess = tf.saved_model.load(preprocess_export_path)
-      tokenize = hub.KerasLayer(preprocess.tokenize)
-      special_tokens_dict = _get_special_tokens_dict(tokenize.resolved_object)
-      for k, v in special_tokens_dict.items():
-        self.assertIsInstance(v, int, "Unexpected type for {}".format(k))
-      tokens = tokenize(sentences)
-      packed_inputs = layers.BertPackInputs(
-          4, special_tokens_dict=special_tokens_dict)(
-              tokens)
-      preprocessing = tf_keras.Model(sentences, packed_inputs)
-      # Map the dataset.
-      ds = tf.data.Dataset.from_tensors(
-          (tf.constant(["abc", "D EF"]), tf.constant([0, 1])))
-      ds = ds.map(lambda features, labels: (preprocessing(features), labels))
-      return ds
-
-    def model_fn(features, labels, mode):
-      del labels  # Unused.
-      return tf_estimator.EstimatorSpec(
-          mode=mode, predictions=features["input_word_ids"])
-
-    estimator = tf_estimator.Estimator(model_fn=model_fn)
-    outputs = list(estimator.predict(input_fn))
-    self.assertAllEqual(outputs, np.array([[2, 6, 3, 0], [2, 4, 5, 3]]))
-
-  # TODO(b/175369555): Remove that code and its test.
-  @parameterized.named_parameters(("Bert", False), ("Sentencepiece", True))
-  def test_check_no_assert(self, use_sp_model):
-    """Tests the self-check during export without assertions."""
-    preprocess_export_path = self._do_export(["d", "ef", "abc", "xy"],
-                                             do_lower_case=True,
-                                             use_sp_model=use_sp_model,
-                                             tokenize_with_offsets=False,
-                                             experimental_disable_assert=False)
-    with self.assertRaisesRegex(AssertionError,
-                                r"failed to suppress \d+ Assert ops"):
-      export_tfhub_lib._check_no_assert(preprocess_export_path)
-
-
-def _result_shapes_in_tf_function(fn, *args, **kwargs):
-  """Returns shapes (as lists) observed on the result of `fn`.
-
-  Args:
-    fn: A callable.
-    *args: TensorSpecs for Tensor-valued arguments and actual values for
-      Python-valued arguments to fn.
-    **kwargs: Same for keyword arguments.
-
-  Returns:
-    The nest of partial tensor shapes (as lists) that is statically known inside
-    tf.function(fn)(*args, **kwargs) for the nest of its results.
-  """
-  # Use a captured mutable container for a side outout from the wrapper.
-  uninitialized = "uninitialized!"
-  result_shapes_container = [uninitialized]
-  assert result_shapes_container[0] is uninitialized
-
-  @tf.function
-  def shape_reporting_wrapper(*args, **kwargs):
-    result = fn(*args, **kwargs)
-    result_shapes_container[0] = tf.nest.map_structure(
-        lambda x: x.shape.as_list(), result)
-    return result
-
-  shape_reporting_wrapper.get_concrete_function(*args, **kwargs)
-  assert result_shapes_container[0] is not uninitialized
-  return result_shapes_container[0]
-
-
-if __name__ == "__main__":
-  tf.test.main()