official/nlp/tools/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])))