export_tfhub_lib.py

# 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.

"""Library of components of export_tfhub.py. See docstring there for more."""

import contextlib
import hashlib
import os
import tempfile

from typing import Optional, Text, Tuple

# Import libraries
from absl import logging
import tensorflow as tf, tf_keras
# pylint: disable=g-direct-tensorflow-import  TODO(b/175369555): Remove these.
from tensorflow.core.protobuf import saved_model_pb2
from tensorflow.python.ops import control_flow_assert
# pylint: enable=g-direct-tensorflow-import
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.modeling import networks


def get_bert_encoder(bert_config):
  """Returns a BertEncoder with dict outputs."""
  bert_encoder = networks.BertEncoder(
      vocab_size=bert_config.vocab_size,
      hidden_size=bert_config.hidden_size,
      num_layers=bert_config.num_hidden_layers,
      num_attention_heads=bert_config.num_attention_heads,
      intermediate_size=bert_config.intermediate_size,
      activation=tf_utils.get_activation(bert_config.hidden_act),
      dropout_rate=bert_config.hidden_dropout_prob,
      attention_dropout_rate=bert_config.attention_probs_dropout_prob,
      max_sequence_length=bert_config.max_position_embeddings,
      type_vocab_size=bert_config.type_vocab_size,
      initializer=tf_keras.initializers.TruncatedNormal(
          stddev=bert_config.initializer_range),
      embedding_width=bert_config.embedding_size,
      dict_outputs=True)

  return bert_encoder


def get_do_lower_case(do_lower_case, vocab_file=None, sp_model_file=None):
  """Returns do_lower_case, replacing None by a guess from vocab file name."""
  if do_lower_case is not None:
    return do_lower_case
  elif vocab_file:
    do_lower_case = "uncased" in vocab_file
    logging.info("Using do_lower_case=%s based on name of vocab_file=%s",
                 do_lower_case, vocab_file)
    return do_lower_case
  elif sp_model_file:
    do_lower_case = True  # All public ALBERTs (as of Oct 2020) do it.
    logging.info("Defaulting to do_lower_case=%s for Sentencepiece tokenizer",
                 do_lower_case)
    return do_lower_case
  else:
    raise ValueError("Must set vocab_file or sp_model_file.")


def _create_model(
    *,
    bert_config: Optional[configs.BertConfig] = None,
    encoder_config: Optional[encoders.EncoderConfig] = None,
    with_mlm: bool,
) -> Tuple[tf_keras.Model, tf_keras.Model]:
  """Creates the model to export and the model to restore the checkpoint.

  Args:
    bert_config: A legacy `BertConfig` to create a `BertEncoder` object. Exactly
      one of encoder_config and bert_config must be set.
    encoder_config: An `EncoderConfig` to create an encoder of the configured
      type (`BertEncoder` or other).
    with_mlm: A bool to control the second component of the result. If True,
      will create a `BertPretrainerV2` object; otherwise, will create a
      `BertEncoder` object.

  Returns:
    A Tuple of (1) a Keras model that will be exported, (2) a `BertPretrainerV2`
    object or `BertEncoder` object depending on the value of `with_mlm`
    argument, which contains the first model and will be used for restoring
    weights from the checkpoint.
  """
  if (bert_config is not None) == (encoder_config is not None):
    raise ValueError("Exactly one of `bert_config` and `encoder_config` "
                     "can be specified, but got %s and %s" %
                     (bert_config, encoder_config))

  if bert_config is not None:
    encoder = get_bert_encoder(bert_config)
  else:
    encoder = encoders.build_encoder(encoder_config)

  # Convert from list of named inputs to dict of inputs keyed by name.
  # Only the latter accepts a dict of inputs after restoring from SavedModel.
  if isinstance(encoder.inputs, list) or isinstance(encoder.inputs, tuple):
    encoder_inputs_dict = {x.name: x for x in encoder.inputs}
  else:
    # encoder.inputs by default is dict for BertEncoderV2.
    encoder_inputs_dict = encoder.inputs
  encoder_output_dict = encoder(encoder_inputs_dict)
  # For interchangeability with other text representations,
  # add "default" as an alias for BERT's whole-input reptesentations.
  encoder_output_dict["default"] = encoder_output_dict["pooled_output"]
  core_model = tf_keras.Model(
      inputs=encoder_inputs_dict, outputs=encoder_output_dict)

  if with_mlm:
    if bert_config is not None:
      hidden_act = bert_config.hidden_act
    else:
      assert encoder_config is not None
      hidden_act = encoder_config.get().hidden_activation

    pretrainer = models.BertPretrainerV2(
        encoder_network=encoder,
        mlm_activation=tf_utils.get_activation(hidden_act))

    if isinstance(pretrainer.inputs, dict):
      pretrainer_inputs_dict = pretrainer.inputs
    else:
      pretrainer_inputs_dict = {x.name: x for x in pretrainer.inputs}
    pretrainer_output_dict = pretrainer(pretrainer_inputs_dict)
    mlm_model = tf_keras.Model(
        inputs=pretrainer_inputs_dict, outputs=pretrainer_output_dict)
    # Set `_auto_track_sub_layers` to False, so that the additional weights
    # from `mlm` sub-object will not be included in the core model.
    # TODO(b/169210253): Use a public API when available.
    core_model._auto_track_sub_layers = False  # pylint: disable=protected-access
    core_model.mlm = mlm_model
    return core_model, pretrainer
  else:
    return core_model, encoder


def export_model(export_path: Text,
                 *,
                 bert_config: Optional[configs.BertConfig] = None,
                 encoder_config: Optional[encoders.EncoderConfig] = None,
                 model_checkpoint_path: Text,
                 with_mlm: bool,
                 copy_pooler_dense_to_encoder: bool = False,
                 vocab_file: Optional[Text] = None,
                 sp_model_file: Optional[Text] = None,
                 do_lower_case: Optional[bool] = None) -> None:
  """Exports an Encoder as SavedModel after restoring pre-trained weights.

  The exported SavedModel implements a superset of the Encoder API for
  Text embeddings with Transformer Encoders described at
  https://www.tensorflow.org/hub/common_saved_model_apis/text.

  In particular, the exported SavedModel can be used in the following way:

  ```
  # Calls default interface (encoder only).

  encoder = hub.load(...)
  encoder_inputs = dict(
      input_word_ids=...,  # Shape [batch, seq_length], dtype=int32
      input_mask=...,      # Shape [batch, seq_length], dtype=int32
      input_type_ids=...,  # Shape [batch, seq_length], dtype=int32
  )
  encoder_outputs = encoder(encoder_inputs)
  assert encoder_outputs.keys() == {
    "pooled_output",   # Shape [batch_size, width], dtype=float32
    "default",         # Alias for "pooled_output" (aligns with other models).
    "sequence_output"  # Shape [batch_size, seq_length, width], dtype=float32
    "encoder_outputs", # List of Tensors with outputs of all transformer layers.
  }
  ```

  If `with_mlm` is True, the exported SavedModel can also be called in the
  following way:

  ```
  # Calls expanded interface that includes logits of the Masked Language Model.
  mlm_inputs = dict(
      input_word_ids=...,       # Shape [batch, seq_length], dtype=int32
      input_mask=...,           # Shape [batch, seq_length], dtype=int32
      input_type_ids=...,       # Shape [batch, seq_length], dtype=int32
      masked_lm_positions=...,  # Shape [batch, num_predictions], dtype=int32
  )
  mlm_outputs = encoder.mlm(mlm_inputs)
  assert mlm_outputs.keys() == {
    "pooled_output",   # Shape [batch, width], dtype=float32
    "sequence_output", # Shape [batch, seq_length, width], dtype=float32
    "encoder_outputs", # List of Tensors with outputs of all transformer layers.
    "mlm_logits"    # Shape [batch, num_predictions, vocab_size], dtype=float32
  }
  ```

  Args:
    export_path: The SavedModel output directory.
    bert_config: An optional `configs.BertConfig` object. Note: exactly one of
      `bert_config` and following `encoder_config` must be specified.
    encoder_config: An optional `encoders.EncoderConfig` object.
    model_checkpoint_path: The path to the checkpoint.
    with_mlm: Whether to export the additional mlm sub-object.
    copy_pooler_dense_to_encoder: Whether to copy the pooler's dense layer used
      in the next sentence prediction task to the encoder.
    vocab_file: The path to the wordpiece vocab file, or None.
    sp_model_file: The path to the sentencepiece model file, or None. Exactly
      one of vocab_file and sp_model_file must be set.
    do_lower_case: Whether to lower-case text before tokenization.
  """
  if with_mlm:
    core_model, pretrainer = _create_model(
        bert_config=bert_config,
        encoder_config=encoder_config,
        with_mlm=with_mlm)
    encoder = pretrainer.encoder_network
    # It supports both the new pretrainer checkpoint produced by TF-NLP and
    # the checkpoint converted from TF1 (original BERT, SmallBERTs).
    checkpoint_items = pretrainer.checkpoint_items
    checkpoint = tf.train.Checkpoint(**checkpoint_items)
  else:
    core_model, encoder = _create_model(
        bert_config=bert_config,
        encoder_config=encoder_config,
        with_mlm=with_mlm)
    checkpoint = tf.train.Checkpoint(
        model=encoder,  # Legacy checkpoints.
        encoder=encoder)
  checkpoint.restore(model_checkpoint_path).assert_existing_objects_matched()

  if copy_pooler_dense_to_encoder:
    logging.info("Copy pooler's dense layer to the encoder.")
    pooler_checkpoint = tf.train.Checkpoint(
        **{"next_sentence.pooler_dense": encoder.pooler_layer})
    pooler_checkpoint.restore(
        model_checkpoint_path).assert_existing_objects_matched()

  # Before SavedModels for preprocessing appeared in Oct 2020, the encoders
  # provided this information to let users do preprocessing themselves.
  # We keep doing that for now. It helps users to upgrade incrementally.
  # Moreover, it offers an escape hatch for advanced users who want the
  # full vocab, not the high-level operations from the preprocessing model.
  if vocab_file:
    core_model.vocab_file = tf.saved_model.Asset(vocab_file)
    if do_lower_case is None:
      raise ValueError("Must pass do_lower_case if passing vocab_file.")
    core_model.do_lower_case = tf.Variable(do_lower_case, trainable=False)
  elif sp_model_file:
    # This was used by ALBERT, with implied values of do_lower_case=True
    # and strip_diacritics=True.
    core_model.sp_model_file = tf.saved_model.Asset(sp_model_file)
  else:
    raise ValueError("Must set vocab_file or sp_model_file")
  core_model.save(export_path, include_optimizer=False, save_format="tf")


class BertPackInputsSavedModelWrapper(tf.train.Checkpoint):
  """Wraps a BertPackInputs layer for export to SavedModel.

  The wrapper object is suitable for use with `tf.saved_model.save()` and
  `.load()`. The wrapper object is callable with inputs and outputs like the
  BertPackInputs layer, but differs from saving an unwrapped Keras object:

    - The inputs can be a list of 1 or 2 RaggedTensors of dtype int32 and
      ragged rank 1 or 2. (In Keras, saving to a tf.function in a SavedModel
      would fix the number of RaggedTensors and their ragged rank.)
    - The call accepts an optional keyword argument `seq_length=` to override
      the layer's .seq_length hyperparameter. (In Keras, a hyperparameter
      could not be changed after saving to a tf.function in a SavedModel.)
  """

  def __init__(self, bert_pack_inputs: layers.BertPackInputs):
    super().__init__()

    # Preserve the layer's configured seq_length as a default but make it
    # overridable. Having this dynamically determined default argument
    # requires self.__call__ to be defined in this indirect way.
    default_seq_length = bert_pack_inputs.seq_length

    @tf.function(autograph=False)
    def call(inputs, seq_length=default_seq_length):
      return layers.BertPackInputs.bert_pack_inputs(
          inputs,
          seq_length=seq_length,
          start_of_sequence_id=bert_pack_inputs.start_of_sequence_id,
          end_of_segment_id=bert_pack_inputs.end_of_segment_id,
          padding_id=bert_pack_inputs.padding_id)

    self.__call__ = call

    for ragged_rank in range(1, 3):
      for num_segments in range(1, 3):
        _ = self.__call__.get_concrete_function([
            tf.RaggedTensorSpec([None] * (ragged_rank + 1), dtype=tf.int32)
            for _ in range(num_segments)
        ],
                                                seq_length=tf.TensorSpec(
                                                    [], tf.int32))


def create_preprocessing(*,
                         vocab_file: Optional[str] = None,
                         sp_model_file: Optional[str] = None,
                         do_lower_case: bool,
                         tokenize_with_offsets: bool,
                         default_seq_length: int) -> tf_keras.Model:
  """Returns a preprocessing Model for given tokenization parameters.

  This function builds a Keras Model with attached subobjects suitable for
  saving to a SavedModel. The resulting SavedModel implements the Preprocessor
  API for Text embeddings with Transformer Encoders described at
  https://www.tensorflow.org/hub/common_saved_model_apis/text.

  Args:
    vocab_file: The path to the wordpiece vocab file, or None.
    sp_model_file: The path to the sentencepiece model file, or None. Exactly
      one of vocab_file and sp_model_file must be set. This determines the type
      of tokenzer that is used.
    do_lower_case: Whether to do lower case.
    tokenize_with_offsets: Whether to include the .tokenize_with_offsets
      subobject.
    default_seq_length: The sequence length of preprocessing results from root
      callable. This is also the default sequence length for the
      bert_pack_inputs subobject.

  Returns:
    A tf_keras.Model object with several attached subobjects, suitable for
    saving as a preprocessing SavedModel.
  """
  # Select tokenizer.
  if bool(vocab_file) == bool(sp_model_file):
    raise ValueError("Must set exactly one of vocab_file, sp_model_file")
  if vocab_file:
    tokenize = layers.BertTokenizer(
        vocab_file=vocab_file,
        lower_case=do_lower_case,
        tokenize_with_offsets=tokenize_with_offsets)
  else:
    tokenize = layers.SentencepieceTokenizer(
        model_file_path=sp_model_file,
        lower_case=do_lower_case,
        strip_diacritics=True,  #  Strip diacritics to follow ALBERT model.
        tokenize_with_offsets=tokenize_with_offsets)

  # The root object of the preprocessing model can be called to do
  # one-shot preprocessing for users with single-sentence inputs.
  sentences = tf_keras.layers.Input(shape=(), dtype=tf.string, name="sentences")
  if tokenize_with_offsets:
    tokens, start_offsets, limit_offsets = tokenize(sentences)
  else:
    tokens = tokenize(sentences)
  pack = layers.BertPackInputs(
      seq_length=default_seq_length,
      special_tokens_dict=tokenize.get_special_tokens_dict())
  model_inputs = pack(tokens)
  preprocessing = tf_keras.Model(sentences, model_inputs)

  # Individual steps of preprocessing are made available as named subobjects
  # to enable more general preprocessing. For saving, they need to be Models
  # in their own right.
  preprocessing.tokenize = tf_keras.Model(sentences, tokens)
  # Provide an equivalent to tokenize.get_special_tokens_dict().
  preprocessing.tokenize.get_special_tokens_dict = tf.train.Checkpoint()
  preprocessing.tokenize.get_special_tokens_dict.__call__ = tf.function(
      lambda: tokenize.get_special_tokens_dict(),  # pylint: disable=[unnecessary-lambda]
      input_signature=[])
  if tokenize_with_offsets:
    preprocessing.tokenize_with_offsets = tf_keras.Model(
        sentences, [tokens, start_offsets, limit_offsets])
    preprocessing.tokenize_with_offsets.get_special_tokens_dict = (
        preprocessing.tokenize.get_special_tokens_dict)
  # Conceptually, this should be
  # preprocessing.bert_pack_inputs = tf_keras.Model(tokens, model_inputs)
  # but technicalities require us to use a wrapper (see comments there).
  # In particular, seq_length can be overridden when calling this.
  preprocessing.bert_pack_inputs = BertPackInputsSavedModelWrapper(pack)

  return preprocessing


def _move_to_tmpdir(file_path: Optional[Text], tmpdir: Text) -> Optional[Text]:
  """Returns new path with same basename and hash of original path."""
  if file_path is None:
    return None
  olddir, filename = os.path.split(file_path)
  hasher = hashlib.sha1()
  hasher.update(olddir.encode("utf-8"))
  target_dir = os.path.join(tmpdir, hasher.hexdigest())
  target_file = os.path.join(target_dir, filename)
  tf.io.gfile.mkdir(target_dir)
  tf.io.gfile.copy(file_path, target_file)
  return target_file


def export_preprocessing(export_path: Text,
                         *,
                         vocab_file: Optional[Text] = None,
                         sp_model_file: Optional[Text] = None,
                         do_lower_case: bool,
                         tokenize_with_offsets: bool,
                         default_seq_length: int,
                         experimental_disable_assert: bool = False) -> None:
  """Exports preprocessing to a SavedModel for TF Hub."""
  with tempfile.TemporaryDirectory() as tmpdir:
    # TODO(b/175369555): Remove experimental_disable_assert and its use.
    with _maybe_disable_assert(experimental_disable_assert):
      preprocessing = create_preprocessing(
          vocab_file=_move_to_tmpdir(vocab_file, tmpdir),
          sp_model_file=_move_to_tmpdir(sp_model_file, tmpdir),
          do_lower_case=do_lower_case,
          tokenize_with_offsets=tokenize_with_offsets,
          default_seq_length=default_seq_length)
      preprocessing.save(export_path, include_optimizer=False, save_format="tf")
    if experimental_disable_assert:
      _check_no_assert(export_path)
  # It helps the unit test to prevent stray copies of the vocab file.
  if tf.io.gfile.exists(tmpdir):
    raise IOError("Failed to clean up TemporaryDirectory")


# TODO(b/175369555): Remove all workarounds for this bug of TensorFlow 2.4
# when this bug is no longer a concern for publishing new models.
# TensorFlow 2.4 has a placement issue with Assert ops in tf.functions called
# from Dataset.map() on a TPU worker. They end up on the TPU coordinator,
# and invoking them from the TPU worker is either inefficient (when possible)
# or impossible (notably when using "headless" TPU workers on Cloud that do not
# have a channel to the coordinator). The bug has been fixed in time for TF 2.5.
# To work around this, the following code avoids Assert ops in the exported
# SavedModels. It monkey-patches calls to tf.Assert from inside TensorFlow and
# replaces them by a no-op while building the exported model. This is fragile,
# so _check_no_assert() validates the result. The resulting model should be fine
# to read on future versions of TF, even if this workaround at export time
# may break eventually. (Failing unit tests will tell.)


def _dont_assert(condition, data, summarize=None, name="Assert"):
  """The no-op version of tf.Assert installed by _maybe_disable_assert."""
  del condition, data, summarize  # Unused.
  if tf.executing_eagerly():
    return
  with tf.name_scope(name):
    return tf.no_op(name="dont_assert")


@contextlib.contextmanager
def _maybe_disable_assert(disable_assert):
  """Scoped monkey patch of control_flow_assert.Assert to a no-op."""
  if not disable_assert:
    yield
    return

  original_assert = control_flow_assert.Assert
  control_flow_assert.Assert = _dont_assert
  yield
  control_flow_assert.Assert = original_assert


def _check_no_assert(saved_model_path):
  """Raises AssertionError if SavedModel contains Assert ops."""
  saved_model_filename = os.path.join(saved_model_path, "saved_model.pb")
  with tf.io.gfile.GFile(saved_model_filename, "rb") as f:
    saved_model = saved_model_pb2.SavedModel.FromString(f.read())

  assert_nodes = []
  graph_def = saved_model.meta_graphs[0].graph_def
  assert_nodes += [
      "node '{}' in global graph".format(n.name)
      for n in graph_def.node
      if n.op == "Assert"
  ]
  for fdef in graph_def.library.function:
    assert_nodes += [
        "node '{}' in function '{}'".format(n.name, fdef.signature.name)
        for n in fdef.node_def
        if n.op == "Assert"
    ]
  if assert_nodes:
    raise AssertionError(
        "Internal tool error: "
        "failed to suppress {} Assert ops in SavedModel:\n{}".format(
            len(assert_nodes), "\n".join(assert_nodes[:10])))