def ner_model(bert_config, num_labels, use_crf=False):
encoder_network = bert_models.get_transformer_encoder(bert_config)
if use_crf:
model = BertNERCRFModel(encoder_network, num_labels)
else:
model = models.BertTokenClassifier(
network=encoder_network,
num_classes=num_labels,
initializer=tf.keras.initializers.TruncatedNormal(stddev=0.02),
dropout_rate=bert_config.hidden_dropout_prob,
output='logits')
return model, encoder_network
def build_model(self):
if self._hub_module:
encoder_network = utils.get_encoder_from_hub(self._hub_module)
else:
encoder_network = encoders.build_encoder(self.task_config.model.encoder)
return models.BertTokenClassifier(
network=encoder_network,
num_classes=len(self.task_config.class_names),
initializer=tf.keras.initializers.TruncatedNormal(
stddev=self.task_config.model.head_initializer_range),
dropout_rate=self.task_config.model.head_dropout,
output='logits')
def build_model(self):
if self._hub_module:
encoder_network = utils.get_encoder_from_hub(self._hub_module)
else:
encoder_network = encoders.instantiate_encoder_from_cfg(
self.task_config.model)
return models.BertTokenClassifier(
network=encoder_network,
num_classes=self.task_config.num_classes,
initializer=tf.keras.initializers.TruncatedNormal(
stddev=self.task_config.model.initializer_range),
dropout_rate=self.task_config.model.dropout_rate,
output='logits')
def build_model(self):
if self.task_config.hub_module_url and self.task_config.init_checkpoint:
raise ValueError('At most one of `hub_module_url` and '
'`init_checkpoint` can be specified.')
if self.task_config.hub_module_url:
encoder_network = utils.get_encoder_from_hub(
self.task_config.hub_module_url)
else:
encoder_network = encoders.build_encoder(
self.task_config.model.encoder)
return models.BertTokenClassifier(
network=encoder_network,
num_classes=len(self.task_config.class_names),
initializer=tf.keras.initializers.TruncatedNormal(
stddev=self.task_config.model.head_initializer_range),
dropout_rate=self.task_config.model.head_dropout,
output='logits')
def build_model(self, train_last_layer_only=False):
"""Modified version of official.nlp.tasks.tagging.build_model
Allows to freeze the underlying bert encoder, such that only the dense
layer is trained.
"""
if self.task_config.hub_module_url and self.task_config.init_checkpoint:
raise ValueError("At most one of `hub_module_url` and "
"`init_checkpoint` can be specified.")
if self.task_config.hub_module_url:
encoder_network = utils.get_encoder_from_hub(
self.task_config.hub_module_url)
else:
encoder_network = encoders.build_encoder(
self.task_config.model.encoder)
encoder_network.trainable = not train_last_layer_only
return models.BertTokenClassifier(
network=encoder_network,
num_classes=len(self.task_config.class_names),
initializer=tf.keras.initializers.TruncatedNormal(
stddev=self.task_config.model.head_initializer_range),
dropout_rate=self.task_config.model.head_dropout,
output="logits")
def classifier_model(bert_config,
num_labels,
max_seq_length=None,
final_layer_initializer=None,
hub_module_url=None,
hub_module_trainable=True):
"""BERT classifier model in functional API style.
Construct a Keras model for predicting `num_labels` outputs from an input with
maximum sequence length `max_seq_length`.
Args:
bert_config: BertConfig or AlbertConfig, the config defines the core BERT or
ALBERT model.
num_labels: integer, the number of classes.
max_seq_length: integer, the maximum input sequence length.
final_layer_initializer: Initializer for final dense layer. Defaulted
TruncatedNormal initializer.
hub_module_url: TF-Hub path/url to Bert module.
hub_module_trainable: True to finetune layers in the hub module.
Returns:
Combined prediction model (words, mask, type) -> (one-hot labels)
BERT sub-model (words, mask, type) -> (bert_outputs)
"""
if final_layer_initializer is not None:
initializer = final_layer_initializer
else:
initializer = tf.keras.initializers.TruncatedNormal(
stddev=bert_config.initializer_range)
if not hub_module_url:
bert_encoder = get_transformer_encoder(bert_config,
max_seq_length,
output_range=None)
# return models.BertClassifier(
# bert_encoder,
# num_classes=num_labels,
# dropout_rate=bert_config.hidden_dropout_prob,
# initializer=initializer), bert_encoder
return models.BertTokenClassifier(
bert_encoder,
num_classes=num_labels,
dropout_rate=bert_config.hidden_dropout_prob,
initializer=initializer), bert_encoder
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')
bert_model = hub.KerasLayer(hub_module_url, trainable=hub_module_trainable)
pooled_output, _ = bert_model([input_word_ids, input_mask, input_type_ids])
output = tf.keras.layers.Dropout(
rate=bert_config.hidden_dropout_prob)(pooled_output)
output = tf.keras.layers.Dense(num_labels,
kernel_initializer=initializer,
name='output')(output)
# print("SHAPE OUTPUT")
# print(output.shape)
return tf.keras.Model(inputs={
'input_word_ids': input_word_ids,
'input_mask': input_mask,
'input_type_ids': input_type_ids
},
outputs=output), bert_model
