Add open-source docs for nlp/tools/export_tfhub.

PiperOrigin-RevId: 359256855

official/nlp/docs/tfhub.md

+# Exporting a pre-trained Encoder to TF Hub
+
+[TOC]
+
+## Overview
+
+This doc explains how to use TF-NLP's
+[export_tfhub](https://github.com/tensorflow/models/blob/master/official/nlp/tools/export_tfhub.py)
+tool to export pre-trained Transformer encoders to SavedModels suitable for
+publication on TF Hub. (For the steps after that, see TF Hub's
+[publisher guide](https://www.tensorflow.org/hub/publish).)
+For testing purposes, those SavedModels can also be used from their export
+locations on the filesystem.
+
+On TF Hub, Transformer encoders for text come as a pair of SavedModels:
+
+*   The preprocessing model applies a tokenizer with a fixed vocab plus some
+    additional logic to turn text into Transformer inputs.
+*   The encoder model (or "model" for short) applies the pre-trained Transformer
+    encoder.
+
+TF Hub defines
+[Common APIs](https://www.tensorflow.org/hub/common_saved_model_apis/text#transformer-encoders)
+for all SavedModels of those two respective types, encapsulating the particular
+choice of preprocessing logic and Encoder architecture.
+
+## Exporting the Encoder
+
+There is a choice between exporting just the encoder, or the encoder plus the
+prediction head for the masked language model (MLM) task from pre-training.
+
+Exporting just the encoder suffices for many straightforward applications.
+
+### Exporting the Encoder alone
+
+To export an encoder-only model, you can set `--export_type=model` and run the
+tool like this:
+
+```shell
+python official/nlp/tools/export_tfhub.py \
+  --encoder_config_file=${BERT_DIR:?}/bert_encoder.yaml \
+  --model_checkpoint_path=${BERT_DIR:?}/bert_model.ckpt \
+  --vocab_file=${BERT_DIR:?}/vocab.txt \
+  --export_type=model \
+  --export_path=/tmp/bert_model
+```
+
+The flag `--encoder_config_file` refers to a YAML file representing the
+[encoders.EncoderConfig](https://github.com/tensorflow/models/search?q=EncoderConfig+path%3Aofficial%2Fnlp%2Fconfigs+filename%3Aencoders.py)
+dataclass, which supports multiple encoders (e.g., BERT, ALBERT). Instead of
+`--encoder_config_file`, you can set `--bert_config_file` to a legacy
+`bert_config.json` file to export a BERT model. If the model definition involves
+[GIN](https://github.com/google/gin-config), the flags `--gin_file` and
+`--gin_params` must be set accordingly, consistent with pre-training.
+
+The `--model_checkpoint_path` refers to an object-based (TF2) checkpoint written
+by
+[BertPretrainerV2](https://github.com/tensorflow/models/search?q=BertPretrainerV2+filename%3Abert_pretrainer.py),
+or any other checkpoint that can be restored to
+`tf.train.Checkpoint(encoder=encoder)` for the encoder defined by the config
+flags. Legacy checkpoints with `model=` instead of `encoder=` are also supported
+for now.
+
+The exported SavedModel expects dict inputs and outputs as follows, implementing
+a specialization of the respective
+[Common SavedModel API](https://www.tensorflow.org/hub/common_saved_model_apis/text#transformer-encoders):
+
+```python
+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
+}
+```
+
+The encoder's pooler layer is restored from the `--model_checkpoint_path`.
+However, unlike classic BERT, `BertPretrainerV2` does not train the pooler layer
+of the encoder. You have three options to handle that:
+
+*   Set flag `--copy_pooler_dense_to_encoder` to copy the pooling layer from the
+    `ClassificationHead` passed to `BertPretrainerV2` for the next sentence
+    prediction task. This mimicks classic BERT, but is not recommended for new
+    models (see next item).
+*   Leave flag `--copy_pooler_dense_to_encoder` unset and export the untrained,
+    randomly initialized pooling layer of the encoder. Folklore (as of 2020) has
+    it that an untrained pooler gets fine-tuned better than a pre-trained
+    pooler, so this is the default.
+*   Leave flag `--copy_pooler_dense_to_encoder` unset and perform your own
+    initialization of the pooling layer before export. For example, Google's
+    [BERT Experts](https://tfhub.dev/google/collections/experts/bert/1)
+    published in October 2020 initialize it to the identity map, reporting equal
+    gains if fine-tuning, and more predictable behavior if not.
+
+In any case, at this time, the export tool requires the encoder model to *have*
+a `pooled_output`, whether trained or not. (This can be revised in the future.)
+
+The encoder model does not include any preprocessing logic, but for the benefit
+of users who take preprocessing into their own hands, the relevant information
+is attached from flags `--vocab_file` or `--sp_model_file`, resp., and
+`--do_lower_case`, which need to be set in exactly the same way as for the
+preprocessing model (see below).
+
+The root object of the exported SavedModel stores the resulting values as
+attributes on the root object:
+
+```python
+encoder = hub.load(...)
+# Gets the filename of the respective tf.saved_model.Asset object.
+if hasattr(encoder, "vocab_file"):
+  print("Wordpiece vocab at", encoder.vocab_file.asset_path.numpy())
+elif hasattr(encoder, "sp_model_file"):
+  print("SentencePiece model at", encoder.sp_model_file.asset_path.numpy())
+# Gets the value of a scalar bool tf.Variable.
+print("...using do_lower_case =", encoder.do_lower_case.numpy())
+```
+
+New users are encouraged to ignore these attributes and use the preprocessing
+model instead. However, there are legacy users, and advanced users that require
+access to the full vocab.
+
+### Exporting the Encoder with a Masked Language Model head
+
+To export an encoder and the masked language model it was trained with, first
+read the preceding section about exporting just the encoder. All the
+explanations there on setting the right flags apply here as well, up to the
+following differences.
+
+The masked language model is added to the export by changing flag
+`--export_type` from `model` to `model_with_mlm`, so the export command looks
+like this:
+
+```shell
+python official/nlp/tools/export_tfhub.py \
+  --encoder_config_file=${BERT_DIR:?}/bert_encoder.yaml \
+  --model_checkpoint_path=${BERT_DIR:?}/bert_model.ckpt \
+  --vocab_file=${BERT_DIR:?}/vocab.txt \
+  --export_type=model_with_mlm \
+  --export_path=/tmp/bert_model
+```
+
+The `--model_checkpoint_path` refers to an object-based (TF2) checkpoint written
+by
+[BertPretrainerV2](https://github.com/tensorflow/models/search?q=BertPretrainerV2+filename%3Abert_pretrainer.py),
+or any other checkpoint that can be restored to
+`tf.train.Checkpoint(**BertPretrainerV2(...).checkpoint_items)` with the encoder
+defined by the config flags.
+
+This is a more comprehensive requirement on the checkpoint than for
+`--export_type=model`; not all Transformer encoders and not all pre-training
+techniques can satisfy it. For example,
+[ELECTRA](https://arxiv.org/abs/2003.10555) uses the BERT architecture but is
+pre-trained without an MLM task.
+
+The root object of the exported SavedModel is called in the same way as above.
+In addition, the SavedModel has an `mlm` subobject that can be called as follows
+to output an `mlm_logits` tensor as well:
+
+```python
+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
+}
+```
+
+The extra subobject imposes a moderate size overhead.
+
+### Exporting from a TF1 BERT checkpoint
+
+A BERT model trained with the
+[original BERT implementation for TF1](https://github.com/google-research/bert)
+can be exported after converting its checkpoint with the
+[tf2_encoder_checkpoint_converter](https://github.com/tensorflow/models/blob/master/official/nlp/bert/tf2_encoder_checkpoint_converter.py)
+tool.
+
+After that, run
+[export_tfhub](https://github.com/tensorflow/models/blob/master/official/nlp/tools/export_tfhub.py)
+per the instructions above on the converted checkpoint. Do not set
+`--copy_pooler_dense_to_encoder`, because the pooler layer is part of the
+converted encoder. For `--vocab_file` and `--do_lower_case`, the values from TF1
+BERT can be used verbatim.
+
+## Exporting the preprocessing model
+
+You can skip this step if TF Hub already has a preprocessing model that does
+exactly what your encoder needs (same tokenizer, same vocab, same normalization
+settings (`do_lower_case`)). You can inspect its collection of
+[Transformer Encoders for Text](https://tfhub.dev/google/collections/transformer_encoders_text/1)
+and click through to models with a similar input domain to find their
+preprocessing models.
+
+To export the preprocessing model, set `--export_type=preprocessing` and run the
+export tool like this:
+
+```shell
+python official/nlp/tools/export_tfhub.py \
+  --vocab_file=${BERT_DIR:?}/vocab.txt \
+  --do_lower_case=True \
+  --export_type=preprocessing \
+  --export_path=/tmp/bert_preprocessing
+```
+
+Note: Set flag `--experimental_disable_assert_in_preprocessing` when exporting
+to users of the public TensorFlow releases 2.4.x to avoid a fatal ops placement
+issue when preprocessing is used within Dataset.map() on TPU workers.
+This is not an issue with TF2.3 and TF2.5+.
+
+Flag `--vocab_file` specifies the vocab file used with
+[BertTokenizer](https://github.com/tensorflow/models/search?q=BertTokenizer+filename%3Atext_layers.py).
+For models that use the
+[SentencepieceTokenizer](https://github.com/tensorflow/models/search?q=SentencepieceTokenizer+filename%3Atext_layers.py),
+set flag `--sp_model_file` instead.
+
+The boolean flag `--do_lower_case` controls text normalization (as in the
+respective tokenizer classes, so it's a bit more than just smashing case). If
+unset, do_lower_case will be enabled if 'uncased' appears in --vocab_file, or
+unconditionally if --sp_model_file is set, mimicking the conventions of BERT and
+ALBERT, respectively. For programmatic use, or if in doubt, it's best to set
+`--do_lower_case` explicity.
+
+If the definition of preprocessing involved
+[GIN](https://github.com/google/gin-config),
+the flags `--gin_file` and `--gin_params` would have to be set accordingly,
+consistent with pre-training. (At the time of this writing, no such GIN
+configurables exist in the code.)
+
+The exported SavedModel can be called in the following way for a single segment
+input.
+
+```python
+preprocessor = hub.load(...)
+text_input = ... # Shape [batch_size], dtype=tf.string
+encoder_inputs = preprocessor(text_input, seq_length=seq_length)
+assert encoder_inputs.keys() == {
+  "input_word_ids", # Shape [batch_size, seq_length], dtype=int32
+  "input_mask",     # Shape [batch_size, seq_length], dtype=int32
+  "input_type_ids"  # Shape [batch_size, seq_length], dtype=int32
+}
+```
+
+Flag `--default_seq_length` controls the value of `seq_length` if that argument
+is omitted in the usage example above. The flag defaults to 128, because
+mutiples of 128 work best for Cloud TPUs, yet the cost of attention computation
+grows quadratically with `seq_length`.
+
+Beyond this example, the exported SavedModel implements the full set interface
+from the preprocessor API for text embeddings with preprocessed inputs and with
+Transformer encoders from TF Hub's
+[Common APIs for text](https://www.tensorflow.org/hub/common_saved_model_apis/text).
+
+Please see
+[tfhub.dev/ensorflow/bert_en_uncased_preprocess](https://tfhub.dev/tensorflow/bert_en_uncased_preprocess)
+for the full documentation of one preprocessing model exported with this tool,
+especially how custom trimming of inputs can happen between `.tokenize` and
+`.bert_pack_inputs`.
+
+Using the `encoder.mlm()` interface requires masking of tokenized inputs by user
+code. The necessary information on the vocabulary encapsulated in the
+preprocessing model can be obtained like this (uniformly across tokenizers):
+
+```python
+special_tokens_dict = preprocess.tokenize.get_special_tokens_dict()
+vocab_size = int(special_tokens_dict["vocab_size"])
+padding_id = int(special_tokens_dict["padding_id"])  # [PAD] or <pad>
+start_of_sequence_id = int(special_tokens_dict["start_of_sequence_id"])  # [CLS]
+end_of_segment_id = int(special_tokens_dict["end_of_segment_id"])  # [SEP]
+mask_id = int(special_tokens_dict["mask_id"])  # [MASK]
+```
+
+## Testing the exported models
+
+Please test your SavedModels before publication by fine-tuning them on a
+suitable task and comparing performance and accuracy to a baseline experiment
+built from equivalent Python code.
+The
+[trainer doc](https://github.com/tensorflow/models/blob/master/official/nlp/docs/train.md)
+has instructions how to run BERT on MNLI and other tasks from the GLUE
+benchmark.