in tensorflow_examples/lite/model_maker/core/task/model_spec/text_spec.py [0:0]
def create_qa_model(bert_config,
max_seq_length,
initializer=None,
hub_module_url=None,
hub_module_trainable=True,
is_tf2=True):
"""Returns BERT qa model along with core BERT model to import weights.
Args:
bert_config: BertConfig, the config defines the core Bert model.
max_seq_length: integer, the maximum input sequence length.
initializer: Initializer for the final dense layer in the span labeler.
Defaulted to TruncatedNormal initializer.
hub_module_url: TF-Hub path/url to Bert module.
hub_module_trainable: True to finetune layers in the hub module.
is_tf2: boolean, whether the hub module is in TensorFlow 2.x format.
Returns:
A tuple of (1) keras model that outputs start logits and end logits and
(2) the core BERT transformer encoder.
"""
if initializer is None:
initializer = tf.keras.initializers.TruncatedNormal(
stddev=bert_config.initializer_range)
input_word_ids = tf.keras.layers.Input(
shape=(max_seq_length,), dtype=tf.int32, name='input_word_ids')
input_mask = tf.keras.layers.Input(
shape=(max_seq_length,), dtype=tf.int32, name='input_mask')
input_type_ids = tf.keras.layers.Input(
shape=(max_seq_length,), dtype=tf.int32, name='input_type_ids')
if is_tf2:
core_model = hub.KerasLayer(hub_module_url, trainable=hub_module_trainable)
pooled_output, sequence_output = core_model(
[input_word_ids, input_mask, input_type_ids])
else:
bert_model = hub_loader.HubKerasLayerV1V2(
hub_module_url,
signature='tokens',
signature_outputs_as_dict=True,
trainable=hub_module_trainable)
outputs = bert_model({
'input_ids': input_word_ids,
'input_mask': input_mask,
'segment_ids': input_type_ids
})
pooled_output = outputs['pooled_output']
sequence_output = outputs['sequence_output']
bert_encoder = tf.keras.Model(
inputs=[input_word_ids, input_mask, input_type_ids],
outputs=[sequence_output, pooled_output],
name='core_model')
return models.BertSpanLabeler(
network=bert_encoder, initializer=initializer), bert_encoder