export_tfhub.py 5.53 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""A script to export the BERT core model as a TF-Hub SavedModel."""
from __future__ import absolute_import
from __future__ import division
# from __future__ import google_type_annotations
from __future__ import print_function

Hongkun Yu's avatar
Hongkun Yu committed
21
# Import libraries
22
23
from absl import app
from absl import flags
24
from absl import logging
25
import tensorflow as tf
26
from typing import Text
27
from official.nlp.bert import bert_models
28
from official.nlp.bert import configs
29
30
31
32
33
34
35

FLAGS = flags.FLAGS

flags.DEFINE_string("bert_config_file", None,
                    "Bert configuration file to define core bert layers.")
flags.DEFINE_string("model_checkpoint_path", None,
                    "File path to TF model checkpoint.")
Chen Chen's avatar
Chen Chen committed
36
flags.DEFINE_string("export_path", None, "TF-Hub SavedModel destination path.")
Hongkun Yu's avatar
Hongkun Yu committed
37
38
flags.DEFINE_string("vocab_file", None,
                    "The vocabulary file that the BERT model was trained on.")
39
40
41
42
43
44
flags.DEFINE_bool(
    "do_lower_case", None, "Whether to lowercase. If None, "
    "do_lower_case will be enabled if 'uncased' appears in the "
    "name of --vocab_file")
flags.DEFINE_enum("model_type", "encoder", ["encoder", "squad"],
                  "What kind of BERT model to export.")
45
46


47
def create_bert_model(bert_config: configs.BertConfig) -> tf.keras.Model:
48
49
50
  """Creates a BERT keras core model from BERT configuration.

  Args:
51
    bert_config: A `BertConfig` to create the core model.
52
53
54
55
56
57
58
59
60
61
62

  Returns:
    A keras model.
  """
  # Adds input layers just as placeholders.
  input_word_ids = tf.keras.layers.Input(
      shape=(None,), dtype=tf.int32, name="input_word_ids")
  input_mask = tf.keras.layers.Input(
      shape=(None,), dtype=tf.int32, name="input_mask")
  input_type_ids = tf.keras.layers.Input(
      shape=(None,), dtype=tf.int32, name="input_type_ids")
Chen Chen's avatar
Chen Chen committed
63
  transformer_encoder = bert_models.get_transformer_encoder(
Zongwei Zhou's avatar
Zongwei Zhou committed
64
      bert_config, sequence_length=None)
Chen Chen's avatar
Chen Chen committed
65
66
67
68
69
70
71
  sequence_output, pooled_output = transformer_encoder(
      [input_word_ids, input_mask, input_type_ids])
  # To keep consistent with legacy hub modules, the outputs are
  # "pooled_output" and "sequence_output".
  return tf.keras.Model(
      inputs=[input_word_ids, input_mask, input_type_ids],
      outputs=[pooled_output, sequence_output]), transformer_encoder
72
73


74
def export_bert_tfhub(bert_config: configs.BertConfig,
75
76
77
78
                      model_checkpoint_path: Text,
                      hub_destination: Text,
                      vocab_file: Text,
                      do_lower_case: bool = None):
79
  """Restores a tf.keras.Model and saves for TF-Hub."""
80
81
82
83
84
85
  # If do_lower_case is not explicit, default to checking whether "uncased" is
  # in the vocab file name
  if do_lower_case is None:
    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)
Chen Chen's avatar
Chen Chen committed
86
  core_model, encoder = create_bert_model(bert_config)
87
88
89
  checkpoint = tf.train.Checkpoint(
      model=encoder,  # Legacy checkpoints.
      encoder=encoder)
André Susano Pinto's avatar
André Susano Pinto committed
90
  checkpoint.restore(model_checkpoint_path).assert_existing_objects_matched()
91
  core_model.vocab_file = tf.saved_model.Asset(vocab_file)
92
  core_model.do_lower_case = tf.Variable(do_lower_case, trainable=False)
93
94
95
  core_model.save(hub_destination, include_optimizer=False, save_format="tf")


96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def export_bert_squad_tfhub(bert_config: configs.BertConfig,
                            model_checkpoint_path: Text,
                            hub_destination: Text,
                            vocab_file: Text,
                            do_lower_case: bool = None):
  """Restores a tf.keras.Model for BERT with SQuAD and saves for TF-Hub."""
  # If do_lower_case is not explicit, default to checking whether "uncased" is
  # in the vocab file name
  if do_lower_case is None:
    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)
  span_labeling, _ = bert_models.squad_model(bert_config, max_seq_length=None)
  checkpoint = tf.train.Checkpoint(model=span_labeling)
  checkpoint.restore(model_checkpoint_path).assert_existing_objects_matched()
  span_labeling.vocab_file = tf.saved_model.Asset(vocab_file)
  span_labeling.do_lower_case = tf.Variable(do_lower_case, trainable=False)
  span_labeling.save(hub_destination, include_optimizer=False, save_format="tf")


116
def main(_):
117
  bert_config = configs.BertConfig.from_json_file(FLAGS.bert_config_file)
118
119
120
121
122
123
124
125
126
  if FLAGS.model_type == "encoder":
    export_bert_tfhub(bert_config, FLAGS.model_checkpoint_path,
                      FLAGS.export_path, FLAGS.vocab_file, FLAGS.do_lower_case)
  elif FLAGS.model_type == "squad":
    export_bert_squad_tfhub(bert_config, FLAGS.model_checkpoint_path,
                            FLAGS.export_path, FLAGS.vocab_file,
                            FLAGS.do_lower_case)
  else:
    raise ValueError("Unsupported model_type %s." % FLAGS.model_type)
127
128
129
130


if __name__ == "__main__":
  app.run(main)